autosim.simulations.spatiotemporal.advection_diffusion_multichannel

Advection-diffusion simulator with multi-channel outputs.

Returns channels: [vorticity, u, v, streamfunction].

Implementation uses spectral Poisson solve (FFT) for streamfunction and central finite differences (via numpy.roll) for spatial derivatives.

class AdvectionDiffusionMultichannel(parameters_range=None, output_names=None, output_indices=None, return_timeseries=False, log_level='progress_bar', n=32, L=10.0, T=80.0, dt=0.25, integrator_kwargs=None)

Bases: SpatioTemporalSimulator

Differentiable advection-diffusion simulator exposing multi-channel outputs.

The simulator evolves a vorticity field according to:

\[\begin{aligned} \partial_t \omega &= \nu \nabla^2 \omega - \mu (u \partial_x \omega + v \partial_y \omega) \end{aligned}\]

It returns selected channels from \([\omega, u, v, \psi]\), where \(\psi\) is the streamfunction.

Parameters:

• parameters_range (dict[str, tuple[float, float]] | None) – Bounds on the sampled viscosity (nu) and advection strength (mu).

• output_names (list[str] | None) – Human-readable names for the returned channels.

• output_indices (list[int] | None) – Channel indices to keep from [vorticity, u, v, streamfunction]. Defaults to all channels in canonical order.

• return_timeseries (bool) – Whether forward returns the entire trajectory instead of a single snapshot.

• log_level (str) – Logging verbosity passed to the base Simulator.

• n (int) – Number of spatial points per dimension.

• L (float) – Domain length in each spatial direction.

• T (float) – Total integration time.

• dt (float) – Temporal resolution used for the ODE solver outputs.

• integrator_kwargs (dict | None) – Extra keyword arguments forwarded to scipy.integrate.solve_ivp.

Notes

Each grid point emits four channels [vorticity, u, v, streamfunction].

forward_samples_spatiotemporal(n, random_seed=None, ensure_exact_n=False)

Produce simulator rollouts along with the sampled parameters.

Parameters:

• n (int) – Number of trajectories to sample.

• random_seed (int | None) – Seed for reproducible parameter draws.

• ensure_exact_n (bool) – Whether to resample failed trajectories until exactly n succeed.

Returns:

Spatiotemporal payload with data, constant_scalars, and constant_fields keys. Channels follow the configured output_indices order.

Return type:

dict

advection_diffusion_rhs(_t, w_flat, n, dx, nu, mu, K3)

Evaluate the vorticity time derivative.

Parameters:

• _t (float) – Time (ignored, but required by solve_ivp).

• w_flat (ndarray) – Flattened vorticity field of length n * n.

• n (int) – Grid resolution per spatial axis.

• dx (float) – Grid spacing.

• nu (float) – Diffusion coefficient.

• mu (float) – Advection strength.

• K3 (ndarray) – Spectral multiplier used for the Poisson solve.

Returns:

Flattened dw/dt matching the shape of w_flat.

Return type:

ndarray

Notes

Implements dw/dt = nu * laplacian(w) - mu * (u * d/dx + v * d/dy) where (u, v) is recovered from the streamfunction via u = dpsi/dy and v = -dpsi/dx.

simulate_advection_diffusion(x, return_timeseries=False, n=50, L=10.0, T=80.0, dt=0.25, integrator_kwargs=None)

Integrate the advection-diffusion system and emit physical channels.

Parameters:

• x (ndarray) – Two-element vector [nu, mu] with viscosity and advection parameters.

• return_timeseries (bool) – If True, return the entire trajectory; otherwise only the terminal state.

• n (int) – Grid resolution per spatial dimension.

• L (float) – Domain size along each axis.

• T (float) – End time for integration.

• dt (float) – Step between recorded solver outputs.

• integrator_kwargs (dict | None) – Extra keyword arguments forwarded to solve_ivp.

Returns:

If return_timeseries is True, an array of shape (n_time, n, n, 4); otherwise shape (n, n, 4). Channels are ordered [vorticity, u, v, streamfunction].

Return type:

ndarray