def new_scene(self, count=None):
if count is None:
count = 1 if self.world.batch_size is None else self.world.batch_size
if count > 1:
self.scene = Scene.create(os.path.join(self.base_dir, self.scene_summary()), batch=count)
else:
self.scene = Scene.create(os.path.join(self.base_dir, self.scene_summary()))
def test_properties(self):
scene = Scene.create(DIR)
self.assertEqual(0, len(scene.properties))
scene.put_property('a', 1)
scene.put_properties({'b': 2, 'c': 3}, d=4)
scene = Scene.at(scene.path)
self.assertEqual(4, len(scene.properties))
scene.remove()
def test_list_scenes(self):
scene = Scene.create(DIR, count=2)
scenes = Scene.list(DIR, include_other=True)
scenes_ = Scene.list(DIR, include_other=False)
self.assertEqual(scenes, scenes_)
self.assertGreaterEqual(len(scenes), 2)
scene_ = Scene.list(DIR, dim='batch')
self.assertGreaterEqual(scene_.shape.volume, 2)
scene.remove()
def test_create_remove_at_equality_batch(self):
scene = Scene.create(DIR, batch=2, config=3)
self.assertEqual(6, scene.shape.volume)
self.assertEqual(('batch', 'config'), scene.shape.names)
scene_ = Scene.at(scene.paths)
self.assertEqual(scene, scene_)
repr(scene)
scene.remove()
try:
Scene.at(scene.paths)
self.fail(
"Scene.at() should fail with IOError if the directory does not exist."
)
except IOError:
pass
def test_create_remove_at_equality_single(self):
scene = Scene.create(DIR)
self.assertEqual(basename(scene.path)[:4], "sim_")
self.assertEqual(1, scene.shape.volume)
scene_ = Scene.at(scene.path)
self.assertEqual(scene, scene_)
repr(scene)
scene.remove()
try:
Scene.at(scene.path)
self.fail(
"Scene.at() should fail with IOError if the directory does not exist."
)
except IOError:
pass
def test_write_read(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1)
vel = DOMAIN.staggered_grid(2)
# write
scene = Scene.create(DIR)
scene.write(smoke=smoke, vel=vel)
self.assertEqual(1, len(scene.frames))
self.assertEqual(1, len(scene.complete_frames))
self.assertEqual(2, len(scene.fieldnames))
# read single
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
self.assertEqual(smoke.extrapolation, smoke_.extrapolation)
self.assertEqual(vel.extrapolation, vel_.extrapolation)
# read multiple
smoke__, vel__ = scene.read(['smoke', 'vel']) # deprecated
field.assert_close(smoke, smoke__)
field.assert_close(vel, vel__)
smoke__, vel__ = scene.read('smoke', 'vel')
field.assert_close(smoke, smoke__)
field.assert_close(vel, vel__)
scene.remove()
def test_write_read_batch_batched_files(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1) * math.random_uniform(count=2, config=3)
vel = DOMAIN.staggered_grid(2) * math.random_uniform(count=2, vel=2)
# write
scene = Scene.create(DIR, count=2)
scene.write({'smoke': smoke, 'vel': vel})
# read batch
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
scene.remove()
def test_write_read_batch_matching(self):
smoke = CenteredGrid(1, extrapolation.BOUNDARY, x=32,
y=32) * math.random_uniform(batch(count=2))
vel = StaggeredGrid(2, 0, x=32, y=32) * math.random_uniform(
batch(count=2))
# write
scene = Scene.create(DIR, count=2)
scene.write({'smoke': smoke, 'vel': vel})
# read batch
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
scene.remove()
def test_write_read_batch_duplicate(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1) * math.random_uniform(count=2)
vel = DOMAIN.staggered_grid(2) * math.random_uniform(count=2)
# write
scene = Scene.create(DIR, more=2)
scene.write({'smoke': smoke, 'vel': vel})
# read batch
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
self.assertEqual(4, smoke_.shape.batch.volume)
self.assertEqual(4, vel_.shape.batch.volume)
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
scene.remove()
def __init__(self, name: str = None, description: str = "", scene: Scene = None):
self.start_time = time.time()
""" Time of creation (`App` constructor invocation) """
self.name = name if name is not None else self.__class__.__name__
""" Human-readable name. """
self.description = description
""" Description to be displayed. """
self.scene = scene
""" Directory to which data and logging information should be written as `Scene` instance. """
self.uses_existing_scene = scene.exist_properties() if scene is not None else False
self.steps = 0
""" Counts the number of times `step()` has been called. May be set by the user. """
self.progress_lock = Lock()
self.pre_step = [] # callback(vis)
self.post_step = [] # callback(vis)
self.progress_available = [] # callback(vis)
self.progress_unavailable = [] # callback(vis)
self.growing_dims = () # tuple or list, used by GUI to determine whether to scroll to last element
self.message = None
self.log_file = None
def test_write_read(self):
smoke = CenteredGrid(1, extrapolation.BOUNDARY, x=32, y=32)
vel = StaggeredGrid(2, 0, x=32, y=32)
# write
scene = Scene.create(DIR)
scene.write(smoke=smoke, vel=vel)
self.assertEqual(1, len(scene.frames))
self.assertEqual(1, len(scene.complete_frames))
self.assertEqual(2, len(scene.fieldnames))
# read single
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
self.assertEqual(smoke.extrapolation, smoke_.extrapolation)
self.assertEqual(vel.extrapolation, vel_.extrapolation)
# read multiple
smoke__, vel__ = scene.read(['smoke', 'vel']) # deprecated
field.assert_close(smoke, smoke__)
field.assert_close(vel, vel__)
smoke__, vel__ = scene.read('smoke', 'vel')
field.assert_close(smoke, smoke__)
field.assert_close(vel, vel__)
scene.remove()
def plot_scalars(scene: str or tuple or list or Scene or math.Tensor,
names: str or tuple or list or math.Tensor = None,
reduce: str or tuple or list or math.Shape = 'names',
down='',
smooth=1,
smooth_alpha=0.2,
smooth_linewidth=2.,
size=(8, 6),
transform: Callable = None,
tight_layout=True,
grid: str or dict = 'y',
log_scale='',
legend='upper right',
x='steps',
xlim=None,
ylim=None,
titles=True,
labels: math.Tensor = None,
xlabel: str = None,
ylabel: str = None,
colors: math.Tensor = 'default'):
"""
Args:
scene: `str` or `Tensor`. Scene paths containing the data to plot.
names: Data files to plot for each scene. The file must be located inside the scene directory and have the name `log_<name>.txt`.
reduce: Tensor dimensions along which all curves are plotted in the same diagram.
down: Tensor dimensions along which diagrams are ordered top-to-bottom instead of left-to-right.
smooth: `int` or `Tensor`. Number of data points to average, -1 for all.
smooth_alpha: Opacity of the non-smoothed curves under the smoothed curves.
smooth_linewidth: Line width of the smoothed curves.
size: Figure size in inches.
transform: Function `T(x,y) -> (x,y)` transforming the curves.
tight_layout:
grid:
log_scale:
legend:
x:
xlim:
ylim:
titles:
labels:
xlabel:
ylabel:
colors: Line colors as `str`, `int` or `Tensor`. Integers are interpreted as indices of the default color list.
Returns:
MatPlotLib [figure](https://matplotlib.org/stable/api/figure_api.html#matplotlib.figure.Figure)
"""
scene = Scene.at(scene)
additional_reduce = ()
if names is None:
first_path = next(iter(math.flatten(scene.paths)))
names = [_str(n) for n in os.listdir(first_path)]
names = [n[4:-4] for n in names if n.endswith('.txt') and n.startswith('log_')]
names = math.wrap(names, batch('names'))
additional_reduce = ['names']
elif isinstance(names, str):
names = math.wrap(names)
elif isinstance(names, (tuple, list)):
names = math.wrap(names, batch('names'))
else:
assert isinstance(names, math.Tensor), f"Invalid argument 'names': {type(names)}"
if not isinstance(colors, math.Tensor):
colors = math.wrap(colors)
if xlabel is None:
xlabel = 'Iterations' if x == 'steps' else 'Time (s)'
shape = (scene.shape & names.shape)
batches = shape.without(reduce).without(additional_reduce)
cycle = list(plt.rcParams['axes.prop_cycle'].by_key()['color'])
fig, axes = plt.subplots(batches.only(down).volume, batches.without(down).volume, figsize=size)
axes = axes if isinstance(axes, numpy.ndarray) else [axes]
for b, axis in zip(batches.meshgrid(), axes):
assert isinstance(axis, plt.Axes)
names_equal = names[b].rank == 0
paths_equal = scene.paths[b].rank == 0
if titles is not None and titles is not False:
if isinstance(titles, str):
axis.set_title(titles)
elif names_equal:
axis.set_title(display_name(str(names[b])))
elif paths_equal:
axis.set_title(os.path.basename(scene.paths[b].native()))
if labels is not None:
curve_labels = labels
elif names_equal:
curve_labels = math.map(os.path.basename, scene.paths[b])
elif paths_equal:
curve_labels = names[b]
else:
curve_labels = math.map(lambda p, n: f"{os.path.basename(p)} - {n}", scene.paths[b], names[b])
def single_plot(name, path, label, i, color, smooth):
logging.debug(f"Reading {os.path.join(path, f'log_{name}.txt')}")
curve = numpy.loadtxt(os.path.join(path, f"log_{name}.txt"))
if curve.ndim == 2:
x_values, values, *_ = curve.T
else:
values = curve
x_values = np.arange(len(values))
if x == 'steps':
pass
else:
assert x == 'time', f"x must be 'steps' or 'time' but got {x}"
logging.debug(f"Reading {os.path.join(path, 'log_step_time.txt')}")
_, x_values, *_ = numpy.loadtxt(os.path.join(path, "log_step_time.txt")).T
values = values[:len(x_values)]
x_values = np.cumsum(x_values[:len(values)])
if transform:
x_values, values = transform(np.stack([x_values, values]))
if color == 'default':
color = cycle[i]
try:
color = int(color)
except ValueError:
pass
if isinstance(color, Number):
color = cycle[int(color)]
logging.debug(f"Plotting curve {label}")
axis.plot(x_values, values, color=color, alpha=smooth_alpha, linewidth=1)
curve = np.stack([x_values, values], -1)
axis.plot(*smooth_uniform_curve(curve, smooth), color=color, linewidth=smooth_linewidth, label=label)
if grid:
if isinstance(grid, dict):
axis.grid(**grid)
else:
grid_axis = 'both' if 'x' in grid and 'y' in grid else grid
axis.grid(which='both', axis=grid_axis, linestyle='--', linewidth=size[1] * 0.3)
if 'x' in log_scale:
axis.set_xscale('log')
if 'y' in log_scale:
axis.set_yscale('log')
if xlim:
axis.set_xlim(xlim)
if ylim:
axis.set_ylim(ylim)
if xlabel:
axis.set_xlabel(xlabel)
if ylabel:
axis.set_ylabel(ylabel)
return name
math.map(single_plot, names[b], scene.paths[b], curve_labels, math.range_tensor(shape.after_gather(b)), colors, smooth)
if legend:
axis.legend(loc=legend)
# Final touches
if tight_layout:
plt.tight_layout()
return fig