evoxels

Public API for the evoxels package.

class evoxels.InversionModel(vf: Any, problem_cls: Type, pos_params: list[str] | None = None, problem_kwargs: dict[str, ~typing.Any] | None=None, timestepper_cls: Type[TimeStepper] = <class 'evoxels.timesteppers.PseudoSpectralIMEX'>, backend: str = 'jax')

Inverse modeling using JAX and diffrax.

This lightweight helper wraps differentiable forward solves for evoxels problem classes and provides utilities to fit model parameters via gradient-based optimization. It is intentionally minimal so that the individual steps of solving a PDE, computing residuals, and running a least-squares optimizer remain easy to follow.

__init__(vf: Any, problem_cls: Type, pos_params: list[str] | None = None, problem_kwargs: dict[str, ~typing.Any] | None=None, timestepper_cls: Type[TimeStepper] = <class 'evoxels.timesteppers.PseudoSpectralIMEX'>, backend: str = 'jax') → None

backend: str = 'jax'

forward_solve(parameters, fieldname, saveat, dt0=0.1, verbose=True)

pos_params: list[str] | None = None

problem_cls: Type

problem_kwargs: dict[str, Any] | None = None

residuals(parameters, y0s__values__saveat, adjoint=diffrax.ForwardMode)

Calculate residuals between measured and simulated states.

Parameters:

parameters (dict) – Current estimate of the model parameters.
y0s__values__saveat (tuple) – Tuple (y0s, values, saveat) where y0s contains the initial states for each sequence, values contains the observed states, and saveat specifies the time points of these observations.
adjoint – Differentiation mode for solve().

Returns:

Array of residuals with shape matching values.

Return type:

jax.Array

solve(parameters, y0, saveat, adjoint=diffrax.ForwardMode, dt0=0.1)

Integrate the configured problem for a given parameter set.

Parameters:

parameters (dict) – Dictionary containing the model parameters to solve with.
y0 (array-like) – Initial state field.
saveat (diffrax.SaveAt) – Time points at which the solution should be stored.
adjoint – Differentiation mode used by diffrax.diffeqsolve().
dt0 (float) – Initial step size for the time integrator.

Returns:

Array of saved state fields with shape (len(saveat.ts), Nx, Ny, Nz).

Return type:

jax.Array

timestepper_cls: alias of PseudoSpectralIMEX

train(initial_parameters, data, inds, adjoint=diffrax.ForwardMode, rtol=1e-06, atol=1e-06, verbose=True, max_steps=1000)

Fit parameters so that the model matches observed data.

This method assembles the observed sequences into a format suitable for optimistix.least_squares() and then runs a Levenberg–Marquardt optimisation to minimise the residuals returned by residuals().

Parameters:

initial_parameters (dict) – Initial guess for the parameters to be optimised.
data (dict) – Dictionary containing "ts" (time stamps) and "ys" (state fields) as produced by solve().
inds (list[list[int]]) – For each sequence, the indices in data that should be used for training. All sequences must have the same spacing.
adjoint – Differentiation mode used when evaluating the residuals.
rtol (float) – Tolerances for the optimiser.
atol (float) – Tolerances for the optimiser.
verbose (bool) – If True, prints optimisation progress.
max_steps (int) – Maximum number of optimisation steps.

Returns:

The optimiser state after termination.

Return type:

optimistix.State

vf: Any

class evoxels.VoxelFields(shape: Tuple[int, int, int], domain_size=(1, 1, 1), convention='cell_center')

Manage 3D voxel grids for simulation, visualization, and I/O.

This class provides a uniform, cell‐centered or staggered‐x voxel grid, handles spacing and origin, and stores any number of named 3D fields.

Parameters:

shape (tuple[int, int, int]) – Number of voxels (Nx, Ny, Nz).
domain_size (tuple[float, float, float], optional) – Physical dimensions (Lx, Ly, Lz). Defaults to (1, 1, 1).
convention (str, optional) – Grid convention, either ‘cell_center’ or ‘staggered_x’. Defaults to ‘cell_center’.

Raises:

ValueError – If domain_size is not length 3 or contains non-numeric values.
ValueError – If convention is not one of ‘cell_center’ or ‘staggered_x’.
Warning – If the spacing ratio max/min > 10, a warning is issued.

shape

Number of voxels (Nx, Ny, Nz).

Type:: tuple[int, int, int]

domain_size

Physical domain lengths.

Type:: tuple[float, float, float]

spacing

Grid spacing (dx, dy, dz).

Type:: tuple[float, float, float]

origin

Coordinates of the (0, 0, 0) corner for cell-centered or staggered grids.

Type:: tuple[float, float, float]

convention

Either ‘cell_center’ or ‘staggered_x’.

Type:: str

precision

NumPy floating-point type for grid coordinates.

Type:: type

grid

Meshgrid arrays (x, y, z) once created by add_grid(), else None.

Type:: tuple[np.ndarray, np.ndarray, np.ndarray] or None

fields

Mapping field names to 3D arrays.

Type:: dict[str, np.ndarray]

Example

>>> vf = VoxelFields((100, 100, 100), domain_size=(1, 1, 1))
>>> vf.add_field('temperature', np_array)
>>> x, y, z = vf.plot_slice('temperature', 10)

property Nx: int

property Ny: int

property Nz: int

__init__(shape: Tuple[int, int, int], domain_size=(1, 1, 1), convention='cell_center'): Create a voxel grid with shape cells.

add_field(name: str, array=None)

Adds a field to the voxel grid.

Parameters:

name (str) – Name of the field.
array (numpy.ndarray, optional) – 3D array to initialize the field. If None, initializes with zeros.

average(name: str): Return the average value of a stored field.

axes() → Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: Returns the 1D coordinate arrays along each axis.

export_to_vtk(filename='output.vtk', field_names=None)

Exports fields to a VTK file for visualization (e.g. VisIt or ParaView).

Parameters:

filename (str) – Name of the output VTK file.
field_names (list, optional) – List of field names to export. Exports all fields if None.

grid_info(): Return a Grid dataclass describing this domain.

meshgrid() → Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: Returns full 3D mesh grids for each axis.

plot_field_interactive(fieldname, direction='x', colormap='viridis', value_bounds=None)

Creates an interactive plot for exploring slices of a 3D field.

Parameters:

fieldname (str) – Name of the field to plot.
direction (str) – Direction of slicing (‘x’, ‘y’, or ‘z’).
dpi (int) – Resolution of the plot.
colormap (str) – Colormap to use for the plot.

Raises:

ValueError – If an invalid direction is provided.

plot_slice(fieldname, slice_index, direction='z', colormap='viridis', value_bounds=None)

Plots a 2D slice of a field along a specified direction.

Parameters:

fieldname (str) – Name of the field to plot.
slice_index (int) – Index of the slice to plot.
direction (str) – Normal direction of the slice (‘x’, ‘y’, or ‘z’).
dpi (int) – Resolution of the plot.
colormap (str) – Colormap to use for the plot.

Raises:

ValueError – If an invalid direction is provided.

set_field(name: str, array: numpy.ndarray)

Set field values for an existing field in the voxel grid.

Parameters:

name (str) – Name of the field.
array (numpy.ndarray, optional) – 3D array.

Raises:

ValueError – If the provided array does not match the voxel grid dimensions.
TypeError – If the provided array is not a numpy array.

set_voxel_sphere(name: str, center, radius, label: int | float = 1)

Create a voxelized representation of a sphere in 3D

Fill voxels within given radius around the given center with value provided by label.

evoxels.run_allen_cahn_solver(voxelfields, fieldnames: str | list[str], backend: str, jit: bool = True, device: str = 'cuda', time_increment: float = 0.5, frames: int = 10, max_iters: int = 100, eps: float = 2.0, gab: float = 1.0, M: float = 1.0, force: float = 0.0, curvature: float = 0.01, potential: Callable | None = None, vtk_out: bool = False, verbose: bool = True, plot_bounds=None): Solves time-dependent Allen-Cahn problem with RungeKutta4 timestepper.

evoxels.run_cahn_hilliard_solver(voxelfields, fieldnames: str | list[str], backend: str, jit: bool = True, device: str = 'cuda', time_increment: float = 0.1, frames: int = 10, max_iters: int = 100, eps: float = 3.0, diffusivity: float = 1.0, mu_hom: Callable | None = None, vtk_out: bool = False, verbose: bool = True, plot_bounds=None): Solves time-dependent Cahn-Hilliard problem with PseudoSpectralIMEX timestepper.

Modules

`boundary_conditions`
`diffrax_adapter`
`fd_stencils`
`function_approximators`
`inversion`
`precompiled_solvers`	Exports for the precompiled solver package.
`problem_definition`
`profiler`
`solvers`
`timesteppers`
`voxelfields`
`voxelgrid`