autosim.utils

autosim.utils#

Utility functions for plotting and visualizing simulator outputs.

generate_output_data(sim, n_train=200, n_valid=20, n_test=20)[source]#

Run simulations and save outputs in a dictionary.

Parameters:
plot_spatiotemporal_video(true, pred=None, pred_uq=None, batch_idx=0, fps=5, vmin=None, vmax=None, cmap='viridis', save_path=None, title='Ground Truth vs Prediction', pred_uq_label='Prediction UQ', colorbar_mode='none', colorbar_mode_uq='none', channel_names=None, preserve_aspect=False)[source]#

Create a video comparing ground truth and predicted spatiotemporal time series.

Parameters:

  • true (Tensor) – Ground-truth tensor.

  • pred (Tensor | None) – Optional predicted tensor of shape (B, T, W, H, C).

  • pred_uq (Tensor | None) – Optional prediction uncertainty tensor of shape (B, T, W, H, C).

  • batch_idx (int) – Which batch index to visualize (default: 0).

  • fps (int) – Frames per second for the video (default: 5).

  • vmin (float | None) – Minimum value for color scale (default: auto from data).

  • vmax (float | None) – Maximum value for color scale (default: auto from data).

  • cmap (str) – Colormap to use (default: “viridis”).

  • save_path (str | None) – Optional path to save the video (e.g., “output.mp4”).

  • title (str) – Title for the video (default: “Ground Truth vs Prediction”).

  • pred_uq_label (str) – Row label used when plotting prediction uncertainty.

  • colorbar_mode (Literal['none', 'row', 'column', 'all']) – Select how colorbars (and underlying color scales) are shared for the first two rows (true vs prediction): - “none”: every subplot gets its own colorbar (default). - “row”: a single colorbar per row (first two rows only). - “column”: a single colorbar per column (true/pred share per channel). - “all”: one colorbar shared across the first two rows.

  • colorbar_mode_uq (Literal['none', 'row']) – Select how colorbars are shared for the prediction uncertainty row.

  • channel_names (list[str] | None) – Optional list of channel names for titles.

  • preserve_aspect (bool) – If True, resize each subplot panel to match the spatial WxH ratio of the data so the image fills the panel without distortion. If False (default), panels are square and the image is stretched to fill via aspect='auto'.

Returns:

Animation object that can be displayed in notebooks.