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, Any] | None = None, backend: str = 'jax')

Inverse modeling using JAX and diffrax.

This small helper class wraps the differentiable solver implementation and provides utilities to fit material parameters via gradient based optimization. It is intentionally lightweight so that new users can easily follow the individual steps: solving the PDE, computing residuals and running a least squares optimiser.

__init__(vf: Any, problem_cls: Type, pos_params: list[str] | None = None, problem_kwargs: dict[str, Any] | None = None, 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 Cahn–Hilliard equation for a given parameter set.

Parameters:
  • parameters (dict) – Dictionary containing the material parameters to solve with.

  • y0 (array-like) – Initial concentration 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 concentration fields with shape (len(saveat.ts), Nx, Ny, Nz).

Return type:

jax.Array

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" (concentration 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', time=None, 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.1, 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 ForwardEuler 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

fd_stencils

function_approximators

inversion

precompiled_solvers

Exports for the precompiled solver package.

problem_definition

profiler

solvers

timesteppers

voxelfields

voxelgrid