openswmm_gymnasium.viz#

Plotly-only visualisation layer for optimisation trajectories. Plan §5.5 + §10 P6.5.

This subpackage is gated by an optional install extra. If Plotly is not installed, every public import here raises ImportError with the install instruction; the rest of openswmm_gymnasium is unaffected.

Public API:

Trajectory — single JSONL episode loader.

TrajectoryRun — directory-of-JSONL multi-episode loader.

figures.reward_curves — scalar/vector reward over time.

figures.pareto_front — 2-D / 3-D scatter, non-dominated highlighted.

figures.hypervolume_trace — normalised HV per episode.

figures.action_timeseries — per-step action heatmap.

figures.network_state_heatmap — node depths / link flows over time.

figures.flooding_attribution — stacked area of cost-component contributions over time.

figures.objective_radar — radar comparing per-objective totals across multiple trajectories.

figures.trajectory_replay — animated network schematic on user-supplied node coordinates.

The module is post-hoc only: it consumes the JSONL trajectories already written by openswmm_gymnasium.wrappers.RecordTrajectory and does not touch a live solver.

author:: Caleb Buahin
copyright:: Copyright (c) 2026 Caleb Buahin
license:: MIT

class openswmm_gymnasium.viz.Trajectory(records)[source]#

Bases: object

One episode loaded from a single JSONL file.

@ivar reset_record: The first "reset" record from the file. :type reset_record: dict @ivar step_records: All "step" records, in order. :type step_records: list[dict]

Parameters:: records (Sequence[dict[str, Any]])

property actions: list[Any]#

List of per-step actions, in their original (possibly nested) form.

Return type:: list

cumulative_cost_vector()[source]#

Per-objective episode cost vector.

For each term name in info["reward_components"], sums the per-step contributions. Useful for HV scoring and Pareto-front plots.

Return type:: numpy.ndarray of shape (n_terms,)

property cumulative_reward: ndarray#

Cumulative reward per step. Same shape as rewards.

Return type:: numpy.ndarray

classmethod from_jsonl(path)[source]#

Load one trajectory from a JSONL file.

Parameters:: path (str | PathLike)
Return type:: Trajectory

property n_steps: int#

Return type:: int

property observations: ndarray#

Stacked observations. Shape (n_steps + 1, obs_dim).

Return type:: numpy.ndarray

reward_components()[source]#

Per-step reward-component values keyed by term name.

Reads info["reward_components"] from each step record. Steps missing the key contribute 0.0 for that step.

Return type:: dict[str, numpy.ndarray]

property rewards: ndarray#

Per-step rewards. Shape (n_steps,) or (n_steps, n_obj).

Return type:: numpy.ndarray

class openswmm_gymnasium.viz.TrajectoryRun(trajectories)[source]#

Bases: object

A collection of Trajectory objects from one experiment directory.

@ivar trajectories: Loaded trajectories, in episode order. :type trajectories: list[Trajectory]

Parameters:: trajectories (Sequence[Trajectory])

cumulative_cost_matrix()[source]#

Stack each trajectory’s cumulative-cost vector into a matrix.

Returns:: Array of shape (n_episodes, n_terms). If trajectories have heterogeneous term sets, the matrix uses the union of all term names (sorted), with 0.0 for missing entries.
Return type:: numpy.ndarray

classmethod from_dir(path, pattern='episode_*.jsonl')[source]#

Load all JSONL files matching pattern from a directory.

Parameters:

path (str or os.PathLike) – Directory containing JSONL files.
pattern (str) – Glob pattern. Defaults to "episode_*.jsonl" (the openswmm_gymnasium.wrappers.RecordTrajectory default).

Return type:

TrajectoryRun

property n_episodes: int#

Figures#

Plotly figure factories.

Each public function returns a plotly.graph_objects.Figure. Plan §5.5.

author:: Caleb Buahin
copyright:: Copyright (c) 2026 Caleb Buahin
license:: MIT

openswmm_gymnasium.viz.figures.action_timeseries(trajectory, *, title='Action timeseries')[source]#

Heatmap of per-step action components.

Rows = action component (flattened from nested Dict, joined by ). Columns = env step. Values = scalar component value.

Parameters:

trajectory (Trajectory) – Loaded episode.
title (str) – Figure title.

Return type:

plotly.graph_objects.Figure

openswmm_gymnasium.viz.figures.flooding_attribution(trajectory, *, title='Reward-component attribution')[source]#

Stacked area of per-step reward-component contributions.

Despite the name (inherited from plan §5.5), this works for any reward-component breakdown — flooding, CSO, energy, etc.

Parameters:

trajectory (Trajectory) – Loaded episode.
title (str) – Figure title.

Return type:

plotly.graph_objects.Figure

openswmm_gymnasium.viz.figures.hypervolume_trace(run, *, ideal, reference, title='Normalized hypervolume')[source]#

Per-episode single-point normalised HV, in [0, 1].

Each episode contributes one point — its cumulative cost vector — against the supplied ideal and reference for normalisation.

Parameters:

run (TrajectoryRun) – Multi-episode run.
ideal (sequence of float) – Per-objective best-case cost (minimisation).
reference (sequence of float) – Per-objective nadir cost.
title (str) – Figure title.

Return type:

plotly.graph_objects.Figure

openswmm_gymnasium.viz.figures.network_state_heatmap(trajectory, *, feature_labels=None, title='Network state')[source]#

Heatmap of observation features over time.

Parameters:

trajectory (Trajectory) – Loaded episode.
feature_labels (sequence of str or None) – Optional names for each observation index. Defaults to "feat[0]", "feat[1]", ....
title (str) – Figure title.

Return type:

plotly.graph_objects.Figure

openswmm_gymnasium.viz.figures.objective_radar(trajectories, *, labels=None, title='Per-objective totals')[source]#

One radar trace per trajectory comparing cumulative-cost vectors.

The set of objective axes is the B{union} of reward-component names across all input trajectories. Trajectories missing a component contribute 0 on that axis.

Parameters:

trajectories (sequence of Trajectory) – Trajectories to compare. Two or more typical.
labels (sequence of str or None) – Per-trajectory labels for the legend. Defaults to "policy[0]", "policy[1]", ....
title (str) – Figure title.

Return type:

plotly.graph_objects.Figure

openswmm_gymnasium.viz.figures.pareto_front(run, *, dimensions=(0, 1), title='Pareto front')[source]#

Scatter of cumulative cost vectors across episodes.

Non-dominated episodes are highlighted in a distinct color. Supports 2-D (default) and 3-D plots; higher dimensions raise.

Parameters:

run (TrajectoryRun) – Multi-episode run.
dimensions (tuple of int) – Indices into the cumulative-cost vector to plot.
title (str) – Figure title.

Return type:

plotly.graph_objects.Figure

Raises:

ValueError – If len(dimensions) not in {2, 3}.

openswmm_gymnasium.viz.figures.reward_curves(trajectory, *, rolling_window=10, title='Reward curves')[source]#

Per-step reward + cumulative reward + rolling-mean overlay.

For scalar reward: three traces on one axis. For vector reward (B{MO env}): one trace per objective, plus per-objective cumulative traces on a secondary axis.

Parameters:

trajectory (Trajectory) – Loaded episode.
rolling_window (int) – Window size for the rolling-mean overlay. Set to 0 or 1 to disable.
title (str) – Figure title.

Return type:

plotly.graph_objects.Figure

openswmm_gymnasium.viz.figures.trajectory_replay(trajectory, *, coords, obs_index, title='Trajectory replay')[source]#

Animated scatter on user-supplied node coordinates.

Parameters:

trajectory (Trajectory) – Loaded episode.
coords (mapping) – Mapping from node ID to (x, y) coordinates.
obs_index (mapping) – Mapping from node ID to its column index in Trajectory.observations.
title (str) – Figure title.

Return type:

plotly.graph_objects.Figure

Raises:

ValueError – If coords and obs_index have mismatched keys.