dymad.utils

class dymad.utils.ControlInterpolator(t, u, axis=-2, order='linear')

Bases: Module

Interpolates the sampled control signal u(t_k) when the ODE solver requests u(t_query).

Parameters:

  • t (torch.Tensor) – 1-D tensor of shape (N,). Sampling times (must be ascending).

  • u (torch.Tensor) – Tensor of shape (…, N, m). Control samples, m inputs per step.

  • axis (int) – Axis of u that corresponds to time. Default is -2.

  • order (str) – Interpolation mode. One of {‘none’, ‘zoh’, ‘linear’, ‘cubic’, etc}.

Note

Not to be confused with dymad.utils.sampling._build_interpolant, which is for data generation, esp. with Numpy. ControlInterpolator is meant to be used in a Torch setting.

forward(t_query)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor | None

t: torch.Tensor
u: torch.Tensor | None
class dymad.utils.JaxWrapper(f_jax, jit=True)

Bases: Module

forward(*xs)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dymad.utils.TrajectorySampler(f, g=None, config=None, rng=None, config_mod=None)

Bases: object

Sampler for generating trajectories.

This class generates batches of trajectories based on a system defined by the functions f and g, which represent the system dynamics and observation model, respectively. The trajectories are sampled according to the configuration specified in the provided YAML file or dictionary.

The dynamics are

With input:

\[\begin{split}\begin{align*} \dot{x} &= f(t, x, u) \\ y &= g(t, x, u) \end{align*}\end{split}\]

Without input:

\[\begin{split}\begin{align*} \dot{x} &= f(t, x) \\ y &= g(t, x) \end{align*}\end{split}\]
Parameters:

  • f (Callable[[float, Array, Array], Array]) – Function defining the system dynamics. It should take time t, state x, and control input u as arguments and return the state derivative. Or just f(t, x) if the system is autonomous (no control input).

  • g (Callable[[float, Array, Array], Array], optional) – Function defining the observation model. It should take time t, state x, and control input u as arguments and return the observation. If not provided, it defaults to the identity function (g(t, x, u) = x). Or just g(t, x) if the system is autonomous (no control input).

  • config (Union[str, Dict], optional) – Path to a YAML configuration file or a dictionary containing the configuration for the sampler. The configuration should specify the dimensions of the states, inputs, and observations, as well as control and initial condition specifications.

  • rng (Union[int, np.random.Generator, None], optional) – Random number generator or seed for reproducibility. If an integer is provided, it is used to seed the default random number generator. If None, the default random generator is used.

  • config_mod (Dict, optional) – Additional configuration parameters to modify the loaded configuration. This should be a dictionary that updates or overrides the values in the loaded configuration.

dymad.utils.adj_to_edge(adj)

Convert dense adjacency matrix to edge index and edge weights.

The inputs can be an arbitrary (nested) list or array of adjacency matrices, The outputs will be lists of edge indices and edge weights of the same hierarchy.

dense_to_sparse is not directly used, because it aggregates all graphs into one big graph. But in our framework, we want to keep them separate.

Lastly, edge_index is transposed to shape (num_edges, 2) for downstream graph-series adapters.

Return type:

tuple[list[list[Tensor]] | list[Tensor] | Tensor | list[list[ndarray]] | list[ndarray] | ndarray, list[list[Tensor]] | list[Tensor] | Tensor | list[list[ndarray]] | list[ndarray] | ndarray]

dymad.utils.animate(fig_func, filename, fps=10, n_frames=None, writer_args=None, fig_args=None)

Create an animation by calling a figure-generating function for each frame.

Parameters:

  • fig_func (function) – Function that generates a figure for a given frame index. It should accept the frame index as its first argument, and return a matplotlib figure and axes.

  • filename (str) – Output filename for the output file.

  • fps (int) – Frames per second for the animation.

  • n_frames (int) – Total number of frames in the animation.

  • writer_args (dict) – Additional keyword arguments to pass to imageio.get_writer.

  • fig_args (dict) – Additional keyword arguments to pass to fig_func.

dymad.utils.compare_contour(x_true, x_pred, x=None, y=None, vmin=None, vmax=None, levels=20, figsize=(12, 4), colorbar=True, axes=None, label=None, mode='contourf', **kwargs)

Compare two contours with error contours.

dymad.utils.config_logger(logger, mode='info', prefix='.')

Configure a logger with specified mode and prefix. The function is intended for a logger instantiated in a class.

Parameters:

  • logger (logging.Logger) – Logger instance to configure.

  • mode (str) – Logging mode, either ‘debug’ or ‘info’. Default is ‘info’.

  • prefix (str) – Prefix for the log file. Default is ‘.’ (current directory).

Return type:

None

dymad.utils.load_config(config_path, config_mod=None)

Load a YAML configuration file and optionally merge with a dictionary.

Parameters:

  • config_path (str) – Path to the YAML configuration file.

  • config_mod (dict, optional) – Dictionary to merge into the config.

Returns:

Merged configuration dictionary.

Return type:

dict

dymad.utils.make_scheduler(scheduler=None, scheduler_type='', **kwargs)

Factory function to create a scheduler instance.

Parameters:

  • scheduler_type (str) – Type of the scheduler to create.

  • **kwargs – Additional arguments for the scheduler.

Returns:

An instance of the specified scheduler type.

Return type:

SchedulerBase

dymad.utils.plot_contour(arrays, x=None, y=None, vmin=None, vmax=None, levels=20, figsize=(12, 4), colorbar=True, axes=None, label=None, grid=None, mode='contourf', **kwargs)

Plot a grid of contour plots for a list of 2D arrays.

Parameters:

  • arrays (list of np.ndarray) – List of 2D arrays to plot.

  • titles (list of str) – Titles for each subplot.

  • x (np.ndarray) – Optional meshgrid for axes.

  • y (np.ndarray) – Optional meshgrid for axes.

  • vmin (float) – Color limits.

  • vmax (float) – Color limits.

  • levels (int or array) – Contour levels.

  • figsize (tuple) – Figure size.

  • colorbar (bool) – Whether to add a colorbar.

  • axes – Existing figure/axes tuple.

  • label (list) – Titles for each subplot.

  • grid (tuple) – Grid layout (n_rows, n_cols).

  • mode (str) – ‘contour’, ‘contourf’, or ‘tricontourf’.

  • **kwargs – Additional arguments for contourf.

Returns:

Matplotlib figure and axes.

Return type:

fig, axes

dymad.utils.plot_cv_results(cv_file, keys=None, ifclose=True, prefix='.', value_scale='log')

Plot cross-validation results.

Parameters:

  • cv_results (list) – List of CVResult objects.

  • metric_name (str) – Name of the metric to plot.

  • ifclose (bool) – Whether to close the plot after saving.

  • prefix (str) – Directory prefix for saving the plot.

dymad.utils.plot_hist(hist, crit, crit_name, model_name, ifclose=True, prefix='.')
dymad.utils.plot_multi_trajs(traj, ts, model_name, us=None, labels=None, ifclose=True, prefix='.', xidx=None, uidx=None, grid=None, xscl=None, uscl=None)

Multi-trajectory version of plot_trajectory - comparison of batches of trajectories.

dymad.utils.plot_summary(npz_files, labels=None, ifscl=True, ifclose=True, prefix='.', output_path=None)

Plot training losses and prediction criterion for multiple summary files on the same figure.

Parameters:

  • npz_files (list) – List of NPZ file paths containing summary data.

  • labels (list) – List of labels for each run (optional).

  • ifscl (bool) – If True, scale the loss by the first epoch loss.

  • ifclose (bool) – Whether to close the plot after saving.

  • prefix (str) – Directory prefix used when output_path is not an absolute path.

  • output_path (str, optional) – If provided, save the figure to this path. By default the plot is not written to disk.

dymad.utils.plot_trajectory(traj, ts, model_name=None, us=None, axes=None, labels=None, ifclose=True, prefix='.', xidx=None, uidx=None, grid=None, xscl=None, uscl=None, cmp_err=True)

Plot trajectories with optional control inputs and save the figure.

Parameters:

  • traj (np.ndarray) – Trajectory data, shape (n_steps, n_features) or (n_traj, n_steps, n_features)

  • ts (np.ndarray) – Time points corresponding to the trajectory data, shape (n_steps,)

  • model_name (str, optional) – Name of the model for the plot title and filename

  • us (np.ndarray, optional) – Control inputs, shape (n_steps, n_controls)

  • axes (list, optional) – List of axes to plot on. If None, creates new subplots

  • labels (list, optional) – Labels for each trajectory, length must match number of trajectories

  • ifclose (bool) – Whether to close the plot after saving

  • prefix (str) – Directory prefix for saving the plot

  • xidx (list, optional) – Indices of state features to plot

  • uidx (list, optional) – Indices of control inputs to plot

  • grid (tuple, optional) – Tuple (n_rows, n_cols) for subplot layout

  • xscl (str, optional) – Scaling mode for state features (‘01’, ‘-11’, ‘std’, or ‘none’)

  • uscl (str, optional) – Scaling mode for control inputs (‘01’, ‘-11’, ‘std’, or ‘none’)

  • cmp_err (bool) – Whether to compute and display RMSE between trajectories

Returns:

List of axes used for plotting

Return type:

axes (list)

dymad.utils.setup_logging(config_path='', mode='info', prefix='.')

Setup logging configuration based on the config file. Assuming the config file name is in the format ‘<case>.yaml’

Parameters:

  • config_path (str) – Path to the configuration file.

  • mode (str) – Logging mode, either ‘debug’ or ‘info’. Default is ‘info’.

  • prefix (str) – Directory prefix for the log file. Default is ‘.’ (current directory).

Return type:

None

Modules

control

graph

misc

plot

sampling

scheduler

wrapper