trphysx.viz¶

trphysx.viz.viz_auto¶

class trphysx.viz.viz_auto.AutoViz¶

Bases: object

Helper class for initializing visualization classes.

Raises:	EnvironmentError – If direct initialization of this class is attempted.

classmethod load_viz(viz_name: str, *args, **kwargs) → trphysx.viz.viz_model.Viz¶

Loads built in visualization class. Currently supports: “lorenz”, “cylinder”, “grayscott”

Parameters:	viz_name (str) – Keyword/name of visualization class
Raises:	KeyError – If viz_name is not a supported visualization type
Returns:	Initialized viz class
Return type:	(Viz)

trphysx.viz.viz_cylinder¶

class trphysx.viz.viz_cylinder.CylinderViz(plot_dir: str = None)¶

Bases: trphysx.viz.viz_model.Viz

Visualization class for flow around a cylinder

Parameters:	plot_dir (str, optional) – Directory to save visualizations in. Defaults to None.

plotPrediction(y_pred: torch.Tensor, y_target: torch.Tensor, plot_dir: str = None, epoch: int = None, pid: int = 0, nsteps: int = 10, stride: int = 20) → None¶

Plots the predicted x-velocity, y-velocity and pressure field contours

Parameters:

y_pred (Tensor) – [T, 3, H, W] Prediction tensor.
y_target (Tensor) – [T, 3, H, W] Target tensor.
plot_dir (str, optional) – Directory to save figure, overrides plot_dir one if provided. Defaults to None.
epoch (int, optional) – Current epoch, used for file name. Defaults to None.
pid (int, optional) – Optional plotting id for indexing file name manually. Defaults to 0.
nsteps (int, optional) – Number of timesteps to plot. Defaults to 10.
stride (int, optional) – Number of timesteps in between plots. Defaults to 10.

plotPredictionVorticity(y_pred: torch.Tensor, y_target: torch.Tensor, plot_dir: str = None, epoch: int = None, pid: int = 0, nsteps: int = 10, stride: int = 10) → None¶

Plots vorticity contours of flow around a cylinder at several time-steps. Vorticity gradients are calculated using standard smoothed central finite difference.

Parameters:

y_pred (Tensor) – [T, 3, H, W] Prediction tensor.
y_target (Tensor) – [T, 3, H, W] Target tensor.
plot_dir (str, optional) – Directory to save figure, overrides class plot_dir if provided. Defaults to None.
epoch (int, optional) – Current epoch, used for file name. Defaults to None.
pid (int, optional) – Optional plotting id for indexing file name manually. Defaults to 0.
nsteps (int, optional) – Number of timesteps to plot. Defaults to 10.
stride (int, optional) – Number of timesteps in between plots. Defaults to 5.

plotEmbeddingPrediction(y_pred: torch.Tensor, y_target: torch.Tensor, plot_dir: str = None, epoch: int = None, bidx: int = None, tidx: int = None, pid: int = 0) → None¶

Plots the predicted x-velocity, y-velocity and pressure field contours

Parameters:

y_pred (Tensor) – [B, T, 3, H, W] Prediction tensor.
y_target (Tensor) – [B, T, 3, H, W] Target tensor.
plot_dir (str, optional) – Directory to save figure, overrides plot_dir one if provided. Defaults to None.
epoch (int, optional) – Current epoch, used for file name. Defaults to None.
bidx (int, optional) – Batch index to plot. Defaults to None (plot random example in batch).
tidx (int, optional) – Timestep index to plot. Defaults to None (plot random time-step).
pid (int, optional) – Optional plotting id for indexing file name manually. Defaults to 0.

trphysx.viz.viz_grayscott¶

class trphysx.viz.viz_grayscott.GrayScottViz(plot_dir: str = None)¶

Bases: trphysx.viz.viz_model.Viz

Visualization class for the 3D Gray-scott system.

Parameters:	plot_dir (str, optional) – Directory to save visualizations in. Defaults to None.

plotPrediction(y_pred: torch.Tensor, y_target: torch.Tensor, plot_dir: str = None, epoch: int = None, pid: int = 0, nsteps: int = 10, stride: int = 5) → None¶

Plots z-slice of Gray-Scott prediction along the z-axis and saves to file

Parameters:

y_pred (torch.Tensor) – [T, 2, H, W, D] prediction time-series of states
y_target (torch.Tensor) – [T, 2, H, W, D] target time-series of states
plot_dir (str, optional) – Directory to save figure, overrides class plot_dir if provided. Defaults to None.
epoch (int, optional) – Current epoch, used for file name. Defaults to None.
pid (int, optional) – Optional plotting id for indexing file name manually. Defaults to 0.
nsteps (int, optional) – Number of timesteps to plot. Defaults to 10.
stride (int, optional) – Number of timesteps in between plots. Defaults to 5.

trphysx.viz.viz_lorenz¶

class trphysx.viz.viz_lorenz.HandlerColormap(cmap: <module 'matplotlib.cm' from '/home/docs/checkouts/readthedocs.org/user_builds/transformer-physx/envs/latest/lib/python3.8/site-packages/matplotlib/cm.py'>, num_stripes: int = 8, **kw)¶

Bases: matplotlib.legend_handler.HandlerBase

Class for creating colormap legend rectangles

Parameters:	cmap (matplotlib.cm) – Matplotlib colormap num_stripes (int) – Number of countour levels (strips) in rectangle

create_artists(legend, orig_handle, xdescent, ydescent, width, height, fontsize, trans)¶

class trphysx.viz.viz_lorenz.LorenzViz(plot_dir: str = None)¶

Bases: trphysx.viz.viz_model.Viz

Visualization class for Lorenz ODE

Parameters:	plot_dir (str, optional) – Directory to save visualizations in. Defaults to None.

plotPrediction(y_pred: torch.Tensor, y_target: torch.Tensor, plot_dir: str = None, epoch: int = None, pid: int = 0) → None¶

Plots a 3D line of a single Lorenz prediction

Parameters:

y_pred (Tensor) – [T, 3] Prediction tensor.
y_target (Tensor) – [T, 3] Target tensor.
plot_dir (str, optional) – Directory to save figure, overrides plot_dir one if provided. Defaults to None.
epoch (int, optional) – Current epoch, used for file name. Defaults to None.
pid (int, optional) – Optional plotting id for indexing file name manually. Defaults to 0.

plotMultiPrediction(y_pred: torch.Tensor, y_target: torch.Tensor, plot_dir: str = None, epoch: int = None, pid: int = 0, nplots: int = 2) → None¶

Plots the 3D lines of multiple Lorenz predictions

Parameters:

y_pred (Tensor) – [B, T, 3] Prediction tensor.
y_target (Tensor) – [B, T, 3] Target tensor.
plot_dir (str, optional) – Directory to save figure, overrides plot_dir one if provided. Defaults to None.
epoch (int, optional) – Current epoch, used for file name. Defaults to None.
pid (int, optional) – Optional plotting id for indexing file name manually. Defaults to 0.
nplots (int, optional) – Number of cases to plot. Defaults to 2.

plotPredictionScatter(y_pred: torch.Tensor, plot_dir: str = None, epoch: int = None, pid: int = 0) → None¶

Plots scatter plots of a Lorenz prediction contoured based on distance from the basins. This will only contour correctly for the parameters s=10, r=28, b=2.667

Parameters:	y_pred (Tensor) – [T, 3] Prediction tensor. plot_dir (str, optional) – Directory to save figure, overrides plot_dir one if provided. Defaults to None. epoch (int, optional) – Current epoch, used for file name. Defaults to None. pid (int, optional) – Optional plotting id for indexing file name manually. Defaults to 0.

trphysx.viz.viz_model¶

class trphysx.viz.viz_model.Viz(plot_dir: str = None)¶

Bases: object

Parent class for visualization

Parameters:	plot_dir (str, optional) – Directory to save visualizations in. Defaults to None.

plotPrediction(y_pred: torch.Tensor, y_target: torch.Tensor, plot_dir: str = None, **kwargs) → None¶

Plots model prediction and target values

Parameters:	y_pred (Tensor) – prediction tensor y_target (Tensor) – target tensor plot_dir (str, optional) – Directory to save plot at. Defaults to None. **kwargs – Additional keyword arguments.
Raises:	NotImplementedError – If function has not been overridden by a child dataset class.

plotEmbeddingPrediction(y_pred: torch.Tensor, y_target: torch.Tensor, plot_dir: str = None, **kwargs) → None¶

Plots model prediction and target values during the embedding training

Parameters:	y_pred (Tensor) – mini-batch of prediction tensor y_target (Tensor) – mini-batch target tensor plot_dir (str, optional) – Directory to save plot at. Defaults to None. **kwargs – Additional keyword arguments.
Raises:	NotImplementedError – If function has not been overridden by a child dataset class.

saveFigure(plot_dir: str = None, file_name: str = 'plot', savepng: bool = True, savepdf: bool = False) → None¶

Saves active matplotlib figure to file

Parameters:	plot_dir (str, optional) – Directory to save plot at, will use class plot_dir if none provided. Defaults to None. file_name (str, optional) – File name of the saved figure. Defaults to ‘plot’. savepng (bool, optional) – Save figure in png format. Defaults to True. savepdf (bool, optional) – Save figure in pdf format. Defaults to False.