def at(directory: str or tuple or list or math.Tensor or 'Scene', id: int or math.Tensor or None = None) -> 'Scene': """ Creates a `Scene` for an existing directory. See Also: `Scene.create()`, `Scene.list()`. Args: directory: Either directory containing scene folder if `id` is given, or scene path if `id=None`. id: (Optional) Scene `id`, will be determined from `directory` if not specified. Returns: `Scene` object for existing scene. """ if isinstance(directory, Scene): assert id is None, f"Got id={id} but directory is already a Scene." return directory if isinstance(directory, (tuple, list)): directory = math.wrap(directory, batch('scenes')) directory = math.map(lambda d: expanduser(d), math.wrap(directory)) if id is None: paths = directory else: id = math.wrap(id) paths = math.map(lambda d, i: join(d, f"sim_{i:06d}"), directory, id) # test all exist for path in math.flatten(paths): if not isdir(path): raise IOError(f"There is no scene at '{path}'") return Scene(paths)
def single_write(f, name): name = _slugify_filename(name) files = math.map(lambda dir_: _filename(dir_, name, frame), directory) if isinstance(f, SampledField): write(f, files) elif isinstance(f, math.Tensor): raise NotImplementedError() elif isinstance(f, Field): raise ValueError("write_sim_frame: only SampledField instances are saved. Resample other Fields before saving them.") else: raise ValueError(f"write_sim_frame: only SampledField instances can be saved but got {f}")
def create(parent_directory: str, shape: math.Shape = math.EMPTY_SHAPE, copy_calling_script=True, **dimensions) -> 'Scene': """ Creates a new `Scene` or a batch of new scenes inside `parent_directory`. See Also: `Scene.at()`, `Scene.list()`. Args: parent_directory: Directory to hold the new `Scene`. If it doesn't exist, it will be created. shape: Determines number of scenes to create. Multiple scenes will be represented by a `Scene` with `is_batch=True`. copy_calling_script: Whether to copy the Python file that invoked this method into the `src` folder of all created scenes. See `Scene.copy_calling_script()`. dimensions: Additional batch dimensions Returns: Single `Scene` object representing the new scene(s). """ shape = (shape & math.shape(**dimensions)).to_batch() parent_directory = expanduser(parent_directory) abs_dir = abspath(parent_directory) if not isdir(abs_dir): os.makedirs(abs_dir) next_id = 0 else: indices = [ int(name[4:]) for name in os.listdir(abs_dir) if name.startswith("sim_") ] next_id = max([-1] + indices) + 1 ids = math.wrap(tuple(range(next_id, next_id + shape.volume))).vector.split(shape) paths = math.map(lambda id_: join(parent_directory, f"sim_{id_:06d}"), ids) scene = Scene(paths) scene.mkdir() if copy_calling_script: try: scene.copy_calling_script() except IOError as err: warnings.warn( f"Failed to copy calling script to scene during Scene.create(): {err}" ) return scene
def subpath(self, name: str, create: bool = False) -> str or tuple: """ Resolves the relative path `name` with this `Scene` as the root folder. Args: name: Relative path with this `Scene` as the root folder. create: Whether to create a directory of that name. Returns: Relative path including the path to this `Scene`. In batch mode, returns a `tuple`, else a `str`. """ def single_subpath(path): path = join(path, name) if create and not isdir(path): os.mkdir(path) return path result = math.map(single_subpath, self._paths) if result.rank == 0: return result.native() else: return result
def shape_type(self) -> Tensor: return math.map(lambda s: f"rot{s}", self._geometry.shape_type)
def single_read(name): name = _slugify_filename(name) files = math.map(lambda dir_: _filename(dir_, name, frame), directory) return read(files, convert_to_backend=convert_to_backend)
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