| 554 | |
| 555 | |
| 556 | def make_contour_plot(array_2d, mode='log', ax=None): |
| 557 | if ax is None: |
| 558 | fig, ax = plt.subplots(figsize=(2.75, 2.75), dpi=300) |
| 559 | else: |
| 560 | fig = plt.gcf() |
| 561 | |
| 562 | if(mode == 'log'): |
| 563 | num_levels = 6 |
| 564 | levels_pos = np.logspace(-2, 0, num=num_levels) # logspace |
| 565 | levels_neg = -1. * levels_pos[::-1] |
| 566 | levels = np.concatenate((levels_neg, np.zeros((0)), levels_pos), axis=0) |
| 567 | colors = plt.get_cmap("Spectral")(np.linspace(0., 1., num=num_levels*2+1)) |
| 568 | elif(mode == 'lin'): |
| 569 | num_levels = 10 |
| 570 | levels = np.linspace(-.5, .5, num=num_levels) |
| 571 | colors = plt.get_cmap("Spectral")(np.linspace(0., 1., num=num_levels)) |
| 572 | |
| 573 | sample = np.flipud(array_2d) |
| 574 | CS = ax.contourf(sample, levels=levels, colors=colors) |
| 575 | fig.colorbar(CS, ax=ax) |
| 576 | |
| 577 | ax.contour(sample, levels=levels, colors='k', linewidths=0.1) |
| 578 | ax.contour(sample, levels=[0], colors='k', linewidths=0.3) |
| 579 | ax.axis('off') |
| 580 | return fig |
| 581 | |
| 582 | |
| 583 | def write_sdf_slice(model, writer, total_steps, prefix='train_', is_multi=False): |