Model Core

model_core

Core class for managing the numerical state of a model.

This module provides a lightweight state container and time-grid manager for time-evolving models. It is designed to support:

• Non-uniform time grids via per-step dt accessors.
• Multi-axis state tensors (e.g., state x age x space x traits).
• Optional full history storage for post-analysis.
• A minimal backend hook to ease future NumPy->(JAX/CuPy) integration.

The core intentionally does not build RHS functions or construct operators; it only manages state, time, and shape/axis metadata in a solver-friendly manner.

ArrayBackend

Bases: Protocol

Minimal array backend interface for ModelCore numerical storage.

asarray(x, dtype=None)

Convert input to an array of the backend type.

Source code in src/op_engine/model_core.py

def asarray(
    self,
    x: object,
    dtype: object | None = None,
) -> np.ndarray:
    """Convert input to an array of the backend type."""
    ...

zeros(shape, dtype=None)

Return a new array of given shape filled with zeros.

Source code in src/op_engine/model_core.py

def zeros(
    self,
    shape: tuple[int, ...],
    dtype: object | None = None,
) -> np.ndarray:
    """Return a new array of given shape filled with zeros."""
    ...

ModelCore(n_states, n_subgroups, time_grid, *, options=None)

Core state and time manager for time-evolving models.

Initialize ModelCore.

Parameters:

Name	Type	Description	Default
n_states	int	Number of state variables.	required
n_subgroups	int	Number of subgroups (e.g., age groups).	required
time_grid	ndarray	1D array of times, shape (n_timesteps,).	required
options	ModelCoreOptions | None	Optional ModelCoreOptions for additional configuration.	None

Raises:

Type	Description
ValueError	if time_grid is invalid or any shapes mismatch.

Source code in src/op_engine/model_core.py

def __init__(
    self,
    n_states: int,
    n_subgroups: int,
    time_grid: np.ndarray,
    *,
    options: ModelCoreOptions | None = None,
) -> None:
    """
    Initialize ModelCore.

    Args:
        n_states: Number of state variables.
        n_subgroups: Number of subgroups (e.g., age groups).
        time_grid: 1D array of times, shape (n_timesteps,).
        options: Optional ModelCoreOptions for additional configuration.

    Raises:
        ValueError: if time_grid is invalid or any shapes mismatch.
    """
    opts = options or ModelCoreOptions()

    # Forward-compat hook; stored for later wiring. For now ModelCore remains
    # NumPy-first; higher layers can decide how to allocate/convert.
    self.xp = opts.xp

    self.dtype = np.dtype(opts.dtype)

    self.time_grid = np.asarray(time_grid, dtype=self.dtype)
    if self.time_grid.ndim != 1:
        raise ValueError(_TIMEGRID_1D_ERROR)

    self.n_timesteps = int(self.time_grid.size)
    if self.n_timesteps < 1:
        raise ValueError(_TIMEGRID_MIN_POINTS_ERROR)

    if self.n_timesteps > 1:
        dt_arr = np.diff(self.time_grid)
        if np.any(dt_arr <= 0):
            raise ValueError(_TIMEGRID_MONOTONE_ERROR)
        self.dt_grid = np.asarray(dt_arr, dtype=self.dtype)
        self.dt = float(self.dt_grid.mean())
    else:
        self.dt_grid = np.asarray([], dtype=self.dtype)
        self.dt = 0.0

    self.n_states = int(n_states)
    self.n_subgroups = int(n_subgroups)
    self._other_axes = tuple(int(x) for x in opts.other_axes)

    self.state_shape = (self.n_states, self.n_subgroups, *self._other_axes)
    self.store_history = bool(opts.store_history)

    if opts.axis_names is None:
        base = ["state", "subgroup"]
        extra = [f"axis{i}" for i in range(2, len(self.state_shape))]
        self.axis_names = tuple(base + extra)
    else:
        if len(opts.axis_names) != len(self.state_shape):
            raise ValueError(
                _AXIS_NAMES_LEN_ERROR.format(
                    actual=len(opts.axis_names),
                    expected=len(self.state_shape),
                )
            )
        self.axis_names = tuple(opts.axis_names)

    self.axis_coords: dict[str, np.ndarray] = {}
    if opts.axis_coords is not None:
        self.axis_coords = {
            str(k): np.asarray(v, dtype=self.dtype)
            for k, v in opts.axis_coords.items()
        }

    self.current_step = 0

    # Per-timestep working state (contiguous).
    self.current_state = np.zeros(self.state_shape, dtype=self.dtype)

    # Optional full history: (n_timesteps, *state_shape)
    self.state_array: FloatArray | None
    if self.store_history:
        self.state_array = cast(
            "FloatArray",
            np.zeros((self.n_timesteps, *self.state_shape), dtype=self.dtype),
        )
    else:
        self.state_array = None

current_time property

Current simulation time t = time_grid[current_step].

Returns:

Type	Description
float	Current time as a float.

state_ndim property

Number of state tensor dimensions (rank).

Returns:

Type	Description
int	State tensor rank as an integer.

advance_timestep(next_state)

Alias for apply_next_state, for solver-friendly naming.

Source code in src/op_engine/model_core.py

def advance_timestep(self, next_state: np.ndarray) -> None:
    """Alias for apply_next_state, for solver-friendly naming."""
    self.apply_next_state(next_state)

apply_deltas(deltas)

Apply state deltas, advancing the timestep.

Supports alternate solver implementations (e.g., splitting updates, additive increments) without forcing allocation of y_next.

Parameters:

Name	Type	Description	Default
deltas	ndarray	State deltas to apply, shape state_shape.	required

Raises:

Type	Description
ValueError	if deltas has incorrect shape.

Source code in src/op_engine/model_core.py

def apply_deltas(self, deltas: np.ndarray) -> None:
    """
    Apply state deltas, advancing the timestep.

    Supports alternate solver implementations (e.g., splitting updates,
    additive increments) without forcing allocation of y_next.

    Args:
        deltas: State deltas to apply, shape state_shape.

    Raises:
        ValueError: if deltas has incorrect shape.
    """
    deltas_arr = np.asarray(deltas, dtype=self.dtype)
    if deltas_arr.shape != self.state_shape:
        raise ValueError(
            _DELTAS_SHAPE_ERROR.format(
                actual=deltas_arr.shape, expected=self.state_shape
            )
        )

    self._check_can_advance()

    self.current_state += deltas_arr

    self.current_step += 1
    if self.store_history and self.state_array is not None:
        self.state_array[self.current_step] = self.current_state

apply_next_state(next_state)

Set the next state directly, advancing the timestep.

Parameters:

Name	Type	Description	Default
next_state	ndarray	State at the next timestep, shape state_shape.	required

Raises:

Type	Description
ValueError	if next_state has incorrect shape.

Source code in src/op_engine/model_core.py

def apply_next_state(self, next_state: np.ndarray) -> None:
    """
    Set the next state directly, advancing the timestep.

    Args:
        next_state: State at the next timestep, shape state_shape.

    Raises:
        ValueError: if next_state has incorrect shape.
    """
    next_state_arr = np.asarray(next_state, dtype=self.dtype)
    if next_state_arr.shape != self.state_shape:
        raise ValueError(
            _NEXT_STATE_SHAPE_ERROR.format(
                actual=next_state_arr.shape, expected=self.state_shape
            )
        )

    self._check_can_advance()

    np.copyto(self.current_state, next_state_arr)

    self.current_step += 1
    if self.store_history and self.state_array is not None:
        self.state_array[self.current_step] = self.current_state

axis_index(axis)

Resolve an axis name or integer into an axis index.

Parameters:

Name	Type	Description	Default
axis	str | int	Axis name or index.	required

Raises:

Type	Description
IndexError	if axis index is out of bounds.
ValueError	if axis name is unknown.

Returns:

Type	Description
int	Axis index as an integer.

Source code in src/op_engine/model_core.py

def axis_index(self, axis: str | int) -> int:
    """
    Resolve an axis name or integer into an axis index.

    Args:
        axis: Axis name or index.

    Raises:
        IndexError: if axis index is out of bounds.
        ValueError: if axis name is unknown.

    Returns:
        Axis index as an integer.
    """
    if isinstance(axis, int):
        if not (0 <= axis < self.state_ndim):
            raise IndexError(_AXIS_INDEX_OOB_ERROR.format(axis=axis))
        return axis

    try:
        return self.axis_names.index(axis)
    except ValueError as exc:
        raise ValueError(_AXIS_UNKNOWN_ERROR.format(axis=axis)) from exc

get_axis_coords(axis)

Return coordinate array for a given axis, or None if not set.

Parameters:

Name	Type	Description	Default
axis	str | int	Axis name or index.	required

Returns:

Type	Description
ndarray | None	Coordinate array for the axis, or None if not set.

Source code in src/op_engine/model_core.py

def get_axis_coords(self, axis: str | int) -> np.ndarray | None:
    """
    Return coordinate array for a given axis, or None if not set.

    Args:
        axis: Axis name or index.

    Returns:
        Coordinate array for the axis, or None if not set.
    """
    idx = self.axis_index(axis)
    name = self.axis_names[idx]
    return self.axis_coords.get(name)

get_current_state()

Retrun the current state.

Returns:

Type	Description
ndarray	Current state, shape state_shape.

Source code in src/op_engine/model_core.py

def get_current_state(self) -> np.ndarray:
    """
    Retrun the current state.

    Returns:
        Current state, shape state_shape.
    """
    return self.current_state

get_dt(step_idx)

Return dt for the step [t_step_idx, t_step_idx+1].

Parameters:

Name	Type	Description	Default
step_idx	int	Timestep index in [0, n_timesteps - 1].	required

Raises:

Type	Description
IndexError	if step_idx is out of bounds.

Returns:

Type	Description
float	dt as a float.

Source code in src/op_engine/model_core.py

def get_dt(self, step_idx: int) -> float:
    """
    Return dt for the step [t_step_idx, t_step_idx+1].

    Args:
        step_idx: Timestep index in [0, n_timesteps - 1].

    Raises:
        IndexError: if step_idx is out of bounds.

    Returns:
        dt as a float.
    """
    if self.n_timesteps <= 1:
        return 0.0
    if not (0 <= step_idx < self.n_timesteps - 1):
        raise IndexError(_DT_INDEX_OOB_ERROR.format(idx=step_idx))
    return float(self.dt_grid[step_idx])

get_state_at(step)

Return the state at a given timestep from history.

Parameters:

Name	Type	Description	Default
step	int	Timestep index in [0, n_timesteps).	required

Returns:

Type	Description
FloatArray	State at the given timestep, shape state_shape.

Raises:

Type	Description
RuntimeError	if history is not stored.
IndexError	if step is out of bounds.

Source code in src/op_engine/model_core.py

def get_state_at(self, step: int) -> FloatArray:
    """
    Return the state at a given timestep from history.

    Args:
        step: Timestep index in [0, n_timesteps).

    Returns:
        State at the given timestep, shape state_shape.

    Raises:
        RuntimeError: if history is not stored.
        IndexError: if step is out of bounds.
    """
    if not self.store_history or self.state_array is None:
        raise RuntimeError(_HISTORY_NOT_STORED_ERROR)

    if not (0 <= step < self.n_timesteps):
        raise IndexError(_STEP_OOB_ERROR)

    # NumPy typing stubs often type ndarray.__getitem__ as Any under mypy,
    # which triggers --strict [no-any-return] without an explicit cast.
    return cast("FloatArray", self.state_array[step])

get_time_at(step_idx)

Return time at a given step index.

Parameters:

Name	Type	Description	Default
step_idx	int	Timestep index in [0, n_timesteps).	required

Raises:

Type	Description
IndexError	if step_idx is out of bounds.

Returns:

Type	Description
float	Time as a float.

Source code in src/op_engine/model_core.py

def get_time_at(self, step_idx: int) -> float:
    """
    Return time at a given step index.

    Args:
        step_idx: Timestep index in [0, n_timesteps).

    Raises:
        IndexError: if step_idx is out of bounds.

    Returns:
        Time as a float.
    """
    if not (0 <= step_idx < self.n_timesteps):
        raise IndexError(_TIME_INDEX_OOB_ERROR.format(idx=step_idx))
    return float(self.time_grid[step_idx])

reshape_for_axis_solve(x, axis)

Reshape a state-like tensor into 2D for an axis-local operator solve.

Parameters:

Name	Type	Description	Default
x	ndarray	State-like tensor of shape state_shape.	required
axis	str | int	Axis name or index along which to reshape.	required

Returns:

Type	Description
tuple[ndarray, tuple[int, ...], int]	(x2d, original_shape, axis_index)

Raises:

Type	Description
ValueError	if x does not have shape state_shape.

Contract

• x2d has shape (axis_len, batch)
• batch is the product of all non-axis dimensions
• unreshape_from_axis_solve(inverse) reconstructs exactly.

Source code in src/op_engine/model_core.py

def reshape_for_axis_solve(
    self,
    x: np.ndarray,
    axis: str | int,
) -> tuple[np.ndarray, tuple[int, ...], int]:
    """Reshape a state-like tensor into 2D for an axis-local operator solve.

    Args:
        x: State-like tensor of shape state_shape.
        axis: Axis name or index along which to reshape.

    Returns:
        (x2d, original_shape, axis_index)

    Raises:
        ValueError: if x does not have shape state_shape.

    Contract:
        - x2d has shape (axis_len, batch)
        - batch is the product of all non-axis dimensions
        - unreshape_from_axis_solve(inverse) reconstructs exactly.
    """
    x_arr = np.asarray(x, dtype=self.dtype)
    if x_arr.shape != self.state_shape:
        raise ValueError(
            _NEXT_STATE_SHAPE_ERROR.format(
                actual=x_arr.shape, expected=self.state_shape
            )
        )

    original_shape = x_arr.shape
    axis_idx = self.axis_index(axis)
    axis_len = int(original_shape[axis_idx])

    moved = np.moveaxis(x_arr, axis_idx, 0)
    x2d = moved.reshape(axis_len, -1)
    return x2d, original_shape, axis_idx

set_initial_state(initial_state)

Set the initial state at time_grid[0].

Parameters:

Name	Type	Description	Default
initial_state	ndarray	Initial state, shape state_shape.	required

Raises:

Type	Description
ValueError	if initial_state has incorrect shape.

Source code in src/op_engine/model_core.py

def set_initial_state(self, initial_state: np.ndarray) -> None:
    """
    Set the initial state at time_grid[0].

    Args:
        initial_state: Initial state, shape state_shape.

    Raises:
        ValueError: if initial_state has incorrect shape.
    """
    initial_state_arr = np.asarray(initial_state, dtype=self.dtype)
    if initial_state_arr.shape != self.state_shape:
        raise ValueError(
            _INITIAL_STATE_SHAPE_ERROR.format(
                actual=initial_state_arr.shape,
                expected=self.state_shape,
            )
        )

    np.copyto(self.current_state, initial_state_arr)

    if self.store_history and self.state_array is not None:
        self.state_array[0] = self.current_state

    self.current_step = 0

unreshape_from_axis_solve(x2d, original_shape, axis)

Inverse of reshape_for_axis_solve.

Parameters:

Name	Type	Description	Default
x2d	ndarray	2D array of shape (axis_len, batch).	required
original_shape	tuple[int, ...]	Original full shape before reshape.	required
axis	str | int	Axis name or index along which the reshape was done.	required

Returns:

Type	Description
ndarray	Reconstructed array of shape original_shape.

Source code in src/op_engine/model_core.py

def unreshape_from_axis_solve(
    self,
    x2d: np.ndarray,
    original_shape: tuple[int, ...],
    axis: str | int,
) -> np.ndarray:
    """
    Inverse of reshape_for_axis_solve.

    Args:
        x2d: 2D array of shape (axis_len, batch).
        original_shape: Original full shape before reshape.
        axis: Axis name or index along which the reshape was done.

    Returns:
        Reconstructed array of shape original_shape.
    """
    axis_idx = self.axis_index(axis)
    axis_len = int(original_shape[axis_idx])

    # Rebuild shape with axis leading, then move it back.
    trailing = tuple(d for i, d in enumerate(original_shape) if i != axis_idx)
    arr = np.asarray(x2d, dtype=self.dtype).reshape((axis_len, *trailing))
    return np.moveaxis(arr, 0, axis_idx)

validate_state_shape(arr, *, msg=None)

Validate that arr has state_shape.

This is intentionally small and solver-friendly; higher layers can reuse it to avoid duplicating shape checks for intermediate/stage states.

Parameters:

Name	Type	Description	Default
arr	ndarray	Array to validate.	required
msg	str | None	Optional custom error message prefix.	None

Raises:

Type	Description
ValueError	if arr does not have shape state_shape.

Source code in src/op_engine/model_core.py

def validate_state_shape(self, arr: np.ndarray, *, msg: str | None = None) -> None:
    """
    Validate that arr has state_shape.

    This is intentionally small and solver-friendly; higher layers can reuse it
    to avoid duplicating shape checks for intermediate/stage states.

    Args:
        arr: Array to validate.
        msg: Optional custom error message prefix.

    Raises:
        ValueError: if arr does not have shape state_shape.
    """
    arr_shape = np.asarray(arr).shape
    if arr_shape != self.state_shape:
        raise ValueError(
            (msg or _NEXT_STATE_SHAPE_ERROR).format(
                actual=arr_shape, expected=self.state_shape
            )
        )

ModelCoreOptions(other_axes=(), axis_names=None, axis_coords=None, store_history=True, dtype=np.float64, xp=np) dataclass

Optional configuration for ModelCore.

This object groups non-essential constructor parameters to keep the ModelCore initializer compact and stable while allowing future extensions without breaking the public API.

Attributes:

Name	Type	Description
other_axes	tuple[int, ...]	Additional axis sizes appended after the default (state, subgroup) axes. Example: (n_space,) or (n_space, n_trait).
axis_names	tuple[str, ...] | None	Optional names for each axis in the state tensor. Length must equal the total state rank.
axis_coords	Mapping[str, ndarray] | None	Optional coordinate arrays keyed by axis name. This is metadata only and enables non-uniform grids and FV/FDM operator builders.
store_history	bool	Whether to store the full time history.
dtype	DTypeLike	Floating-point dtype for internal arrays.
xp	object	Array backend module (default NumPy). This is a forward- compatibility hook for GPU backends.