File size: 2,435 Bytes
1cd8a52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """Numerical integrators for multiscale polynomial memory systems."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Iterable, Optional, Sequence, Tuple, Union
import numpy as np
from scipy.integrate import solve_ivp
from .utils import vec, unvec
State = np.ndarray
Params = Union[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray]]
@dataclass
class IntegrationResult:
"""Container holding the final state and, optionally, the trajectory."""
state: State
trajectory: Optional[np.ndarray] = None
def _validate_inputs(xs: np.ndarray, dt: float) -> None:
if xs.ndim != 1:
raise ValueError("`xs` must be a one-dimensional array")
if not np.isscalar(dt) or dt <= 0:
raise ValueError("`dt` must be a positive scalar")
def integrate_system(
xs: np.ndarray,
dt: float,
params: Params,
*,
state: Optional[np.ndarray] = None,
return_trajectory: bool = False,
) -> IntegrationResult:
"""Integrate a discrete control signal with either single or multiscale parameters."""
_validate_inputs(xs, dt)
multiscale = len(params) == 3
if multiscale:
A, B, M = params
n, m = B.shape
if state is None:
state = np.zeros(n * m, dtype=float)
elif state.shape != (n * m,):
raise ValueError(f"Expected state of shape {(n * m,)}, got {state.shape}")
def get_fun(x_val: float):
return lambda _, y: vec(A @ unvec(y, n, m) @ M + x_val * B)
else:
A, b = params
d = b.shape[0]
if state is None:
state = np.zeros(d, dtype=float)
elif state.shape != (d,):
raise ValueError(f"Expected state of shape {(d,)}, got {state.shape}")
def get_fun(x_val: float):
return lambda _, y: A @ y + b * x_val
trajectory: Optional[list[np.ndarray]] = [] if return_trajectory else None
for x_val in xs:
fun = get_fun(float(x_val))
result = solve_ivp(fun, (0.0, float(dt)), state, method="RK45")
state = result.y[:, -1]
if return_trajectory:
trajectory.append(state.copy())
if return_trajectory and trajectory is not None:
history = np.stack(trajectory, axis=0)
else:
history = None
return IntegrationResult(state=state, trajectory=history)
__all__ = ["IntegrationResult", "integrate_system"]
|