1479 lines
59 KiB
Python
1479 lines
59 KiB
Python
"""Audited sparse secant Jacobians for implicit ODE solvers.
|
|
|
|
This module deliberately has no dependency on ``GenericFluidSystem`` or the
|
|
SciPy solver wrappers. It is an experimental, opt-in numerical kernel that can
|
|
be wired into BDF/Radau through their callable ``jac`` argument after its model
|
|
level eligibility rules have been checked.
|
|
|
|
The first Jacobian in every solver segment is a complete, sparsity-coloured
|
|
finite-difference build. Ordinary RHS samples at the same time coordinate can
|
|
then update that matrix with row-constrained sparse secants. At most one such
|
|
matrix is reused before another complete build, and reuse is allowed only after
|
|
a deterministic directional finite-difference audit. An audit mismatch falls
|
|
back to a complete build in the same call.
|
|
|
|
SciPy's private ``num_jac`` and ``group_columns`` helpers are isolated here so
|
|
an incompatible SciPy version fails with one explicit compatibility error
|
|
rather than changing solver behaviour silently.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
import inspect
|
|
import math
|
|
from time import perf_counter
|
|
from typing import TypeAlias
|
|
|
|
import numpy as np
|
|
|
|
from app.simulation.core.errors import RecoverableTrialStateError
|
|
|
|
|
|
JacobianEvaluate: TypeAlias = Callable[[float, np.ndarray], Sequence[float]]
|
|
ExactRowValues: TypeAlias = Sequence[float] | Mapping[int, float]
|
|
ExactRow: TypeAlias = ExactRowValues | Callable[[float, np.ndarray], ExactRowValues]
|
|
ExactColumnValues: TypeAlias = Sequence[float] | np.ndarray
|
|
ExactColumnProviderResult: TypeAlias = (
|
|
Sequence[Sequence[float]] | np.ndarray | Mapping[int, ExactColumnValues]
|
|
)
|
|
ExactColumnProvider: TypeAlias = Callable[
|
|
[float, np.ndarray, tuple[int, ...]],
|
|
ExactColumnProviderResult,
|
|
]
|
|
ExactColumns: TypeAlias = (
|
|
tuple[Sequence[int], ExactColumnProvider]
|
|
| Mapping[int, ExactColumnValues]
|
|
)
|
|
|
|
|
|
class SparseJacobianCompatibilityError(RuntimeError):
|
|
"""Raised when the installed SciPy cannot provide the expected helpers."""
|
|
|
|
|
|
class ExactColumnsUnavailable(RuntimeError):
|
|
"""Request a numerical-column fallback for one Jacobian assembly.
|
|
|
|
An exact-column provider may raise this exception when its analytical or
|
|
tangent calculation is not applicable at the current primal point. The
|
|
provider must raise before mutating the shared model closure: the Jacobian
|
|
builder deliberately reuses the real base RHS evaluation that immediately
|
|
preceded the provider call.
|
|
|
|
Shape errors, non-finite values, and all other provider exceptions remain
|
|
programming errors and are intentionally not converted into a fallback.
|
|
"""
|
|
|
|
def __init__(self, reason: str) -> None:
|
|
self.reason = str(reason)
|
|
super().__init__(self.reason)
|
|
|
|
|
|
class SparseSecantJacobian:
|
|
"""Build and cautiously reuse a sparse numerical ODE Jacobian.
|
|
|
|
Parameters
|
|
----------
|
|
evaluate:
|
|
Scalar-state RHS evaluator. The callable receives ``(time, state)``
|
|
and must return one derivative per state. It need not be vectorized.
|
|
sparsity:
|
|
Conservative square Jacobian structure accepted by ``scipy.sparse``.
|
|
Missing structural entries cannot be repaired by the secant update.
|
|
atol:
|
|
Scalar or state-aligned absolute tolerance used by SciPy's numerical
|
|
Jacobian step selection.
|
|
exact_rows:
|
|
Optional complete row providers keyed by row index. A provider (or
|
|
constant value) may return either a dense row or a ``column: value``
|
|
mapping. Exact rows are excluded from finite differences and secant
|
|
updates, then refreshed at the requested ``(time, state)``.
|
|
exact_columns:
|
|
Optional ``(column_indexes, provider)`` pair. The provider receives
|
|
``(time, state, normalized_column_indexes)`` and returns a dense
|
|
``(state_count, column_count)`` matrix in that exact column order. An
|
|
ordered ``column: dense_column`` mapping is also accepted from the
|
|
provider so its returned order can be checked. A fixed mapping may be
|
|
supplied directly as a convenience. Exact columns are removed before
|
|
seed-0 finite-difference colouring, then complete columns and complete
|
|
rows are applied in that order, so ``exact_rows`` wins at intersections.
|
|
A callable provider may raise :class:`ExactColumnsUnavailable` before
|
|
changing the shared model closure to request a one-build fallback to
|
|
the original seed-0 numerical Jacobian.
|
|
audit_relative_tolerance:
|
|
Maximum component-wise relative mismatch for the deterministic Jv
|
|
audit. The audit also includes a floating-point roundoff allowance.
|
|
audit_absolute_tolerance:
|
|
Optional absolute allowance for the Jv delta comparison.
|
|
max_consecutive_reuses:
|
|
``0`` selects complete seed-0 finite-difference builds and disables
|
|
observation work. ``1`` enables the experimental audited secant path.
|
|
|
|
Notes
|
|
-----
|
|
``observe`` must receive only ordinary RHS evaluations. Calls made by
|
|
this object while building or auditing the Jacobian are suppressed
|
|
automatically if the supplied evaluator itself reports them back through
|
|
``observe``.
|
|
"""
|
|
|
|
_MAX_PENDING_SECANT_PAIRS = 4
|
|
# Keep SciPy's deterministic compatibility baseline. Alternative greedy
|
|
# orders can change finite-difference perturbations and therefore event
|
|
# sequences in non-smooth models, even if they use fewer colour groups.
|
|
_COLORING_SEED = 0
|
|
|
|
def __init__(
|
|
self,
|
|
evaluate: JacobianEvaluate,
|
|
sparsity: object,
|
|
atol: float | Sequence[float],
|
|
exact_rows: Mapping[int, ExactRow] | None = None,
|
|
exact_columns: ExactColumns | None = None,
|
|
*,
|
|
audit_relative_tolerance: float = 5.0e-2,
|
|
audit_absolute_tolerance: float = 0.0,
|
|
max_consecutive_reuses: int = 1,
|
|
) -> None:
|
|
try:
|
|
from scipy.sparse import csc_matrix
|
|
except ImportError as exc: # pragma: no cover - project requires SciPy.
|
|
raise SparseJacobianCompatibilityError(
|
|
"Sparse secant Jacobians require scipy.sparse."
|
|
) from exc
|
|
|
|
if not callable(evaluate):
|
|
raise TypeError("evaluate must be callable.")
|
|
self._evaluate_rhs = evaluate
|
|
|
|
structure = csc_matrix(sparsity, dtype=bool)
|
|
structure.eliminate_zeros()
|
|
if len(structure.shape) != 2 or structure.shape[0] != structure.shape[1]:
|
|
raise ValueError("Jacobian sparsity must be a square matrix.")
|
|
if structure.shape[0] == 0:
|
|
raise ValueError("Jacobian sparsity must contain at least one state.")
|
|
self._state_count = int(structure.shape[0])
|
|
tolerance = np.asarray(atol, dtype=float)
|
|
if tolerance.ndim == 0:
|
|
tolerance = np.full(self._state_count, float(tolerance), dtype=float)
|
|
if tolerance.shape != (self._state_count,):
|
|
raise ValueError(
|
|
"atol must be scalar or contain one value per Jacobian state."
|
|
)
|
|
if not np.all(np.isfinite(tolerance)) or np.any(tolerance <= 0.0):
|
|
raise ValueError("atol values must be finite and greater than zero.")
|
|
self._atol = tolerance
|
|
|
|
if (
|
|
not math.isfinite(audit_relative_tolerance)
|
|
or audit_relative_tolerance < 0.0
|
|
):
|
|
raise ValueError(
|
|
"audit_relative_tolerance must be finite and non-negative."
|
|
)
|
|
if (
|
|
not math.isfinite(audit_absolute_tolerance)
|
|
or audit_absolute_tolerance < 0.0
|
|
):
|
|
raise ValueError(
|
|
"audit_absolute_tolerance must be finite and non-negative."
|
|
)
|
|
self._audit_relative_tolerance = float(audit_relative_tolerance)
|
|
self._audit_absolute_tolerance = float(audit_absolute_tolerance)
|
|
if (
|
|
not isinstance(max_consecutive_reuses, int)
|
|
or isinstance(max_consecutive_reuses, bool)
|
|
or max_consecutive_reuses not in {0, 1}
|
|
):
|
|
raise ValueError("max_consecutive_reuses must be either 0 or 1.")
|
|
self._max_consecutive_reuses = max_consecutive_reuses
|
|
|
|
self._exact_rows = dict(exact_rows or {})
|
|
for row_index in self._exact_rows:
|
|
if (
|
|
not isinstance(row_index, int)
|
|
or isinstance(row_index, bool)
|
|
or not 0 <= row_index < self._state_count
|
|
):
|
|
raise ValueError(
|
|
f"Exact Jacobian row index {row_index!r} is out of range."
|
|
)
|
|
self._exact_row_indexes = frozenset(self._exact_rows)
|
|
|
|
self._exact_column_provider: ExactColumnProvider | None = None
|
|
self._fixed_exact_column_values: np.ndarray | None = None
|
|
if exact_columns is None:
|
|
exact_column_indexes: tuple[int, ...] = ()
|
|
elif isinstance(exact_columns, Mapping):
|
|
exact_column_indexes = self._normalize_exact_column_indexes(
|
|
exact_columns.keys()
|
|
)
|
|
fixed_columns: list[np.ndarray] = []
|
|
for column_index in exact_column_indexes:
|
|
column = np.asarray(exact_columns[column_index], dtype=float)
|
|
if column.shape != (self._state_count,):
|
|
raise ValueError(
|
|
f"Exact Jacobian column {column_index} must have shape "
|
|
f"{(self._state_count,)}, received {column.shape}."
|
|
)
|
|
if not np.all(np.isfinite(column)):
|
|
raise ValueError(
|
|
"Exact Jacobian columns must contain finite values."
|
|
)
|
|
fixed_columns.append(column.copy())
|
|
self._fixed_exact_column_values = (
|
|
np.column_stack(fixed_columns)
|
|
if fixed_columns
|
|
else np.empty((self._state_count, 0), dtype=float)
|
|
)
|
|
else:
|
|
if not isinstance(exact_columns, tuple) or len(exact_columns) != 2:
|
|
raise TypeError(
|
|
"exact_columns must be a fixed mapping or a "
|
|
"(column_indexes, provider) pair."
|
|
)
|
|
requested_columns, provider = exact_columns
|
|
if not callable(provider):
|
|
raise TypeError("The exact Jacobian column provider must be callable.")
|
|
exact_column_indexes = self._normalize_exact_column_indexes(
|
|
requested_columns
|
|
)
|
|
self._exact_column_provider = provider
|
|
self._exact_column_indexes = exact_column_indexes
|
|
self._exact_column_index_set = frozenset(exact_column_indexes)
|
|
self._exact_column_pattern = np.asarray(
|
|
structure[:, list(exact_column_indexes)].toarray(),
|
|
dtype=bool,
|
|
)
|
|
self._exact_column_outside_pattern_nonzero_count = 0
|
|
|
|
self._original_sparsity = structure.copy()
|
|
finite_difference_sparsity = structure.tolil(copy=True)
|
|
for row_index in self._exact_row_indexes:
|
|
finite_difference_sparsity.rows[row_index] = []
|
|
finite_difference_sparsity.data[row_index] = []
|
|
self._original_finite_difference_sparsity = (
|
|
finite_difference_sparsity.tocsc().astype(bool)
|
|
)
|
|
self._remaining_finite_difference_sparsity = (
|
|
self._original_finite_difference_sparsity.copy()
|
|
)
|
|
for column_index in self._exact_column_indexes:
|
|
start = int(
|
|
self._remaining_finite_difference_sparsity.indptr[column_index]
|
|
)
|
|
stop = int(
|
|
self._remaining_finite_difference_sparsity.indptr[
|
|
column_index + 1
|
|
]
|
|
)
|
|
self._remaining_finite_difference_sparsity.data[start:stop] = False
|
|
self._remaining_finite_difference_sparsity.eliminate_zeros()
|
|
remaining_columns = np.asarray(
|
|
self._remaining_finite_difference_sparsity.getnnz(axis=0)
|
|
).reshape(-1) > 0
|
|
self._remaining_finite_difference_columns = np.flatnonzero(
|
|
remaining_columns
|
|
)
|
|
self._finite_difference_column_count = int(
|
|
self._remaining_finite_difference_columns.size
|
|
)
|
|
|
|
self._num_jac, group_columns = self._load_scipy_helpers()
|
|
if self._exact_column_indexes:
|
|
self._validate_subset_sparse_compatibility()
|
|
self._original_groups = np.asarray(
|
|
group_columns(
|
|
self._original_sparsity,
|
|
order=self._COLORING_SEED,
|
|
),
|
|
dtype=int,
|
|
)
|
|
if self._original_groups.shape != (self._state_count,):
|
|
raise SparseJacobianCompatibilityError(
|
|
"SciPy Jacobian column grouping returned an unexpected shape."
|
|
)
|
|
self._original_color_group_count = (
|
|
int(self._original_groups.max(initial=-1)) + 1
|
|
)
|
|
|
|
if self._exact_column_indexes:
|
|
# True subset finite differences: exact and otherwise inactive
|
|
# columns are absent from both colouring and perturbation batches.
|
|
grouped_columns = self._remaining_finite_difference_columns
|
|
grouping_structure = self._remaining_finite_difference_sparsity[
|
|
:, grouped_columns
|
|
]
|
|
self._remaining_groups = np.full(
|
|
self._state_count,
|
|
-1,
|
|
dtype=int,
|
|
)
|
|
else:
|
|
# Preserve the established SciPy seed-0 baseline when no exact
|
|
# columns are requested, including the existing exact-row policy.
|
|
grouped_columns = np.empty(0, dtype=int)
|
|
grouping_structure = None
|
|
self._remaining_groups = self._original_groups.copy()
|
|
if grouped_columns.size:
|
|
remaining_groups = np.asarray(
|
|
group_columns(
|
|
grouping_structure,
|
|
order=self._COLORING_SEED,
|
|
),
|
|
dtype=int,
|
|
)
|
|
if remaining_groups.shape != (grouped_columns.size,):
|
|
raise SparseJacobianCompatibilityError(
|
|
"SciPy Jacobian column grouping returned an unexpected shape."
|
|
)
|
|
self._remaining_groups[grouped_columns] = remaining_groups
|
|
self._remaining_color_group_count = (
|
|
int(remaining_groups.max(initial=-1)) + 1
|
|
)
|
|
elif self._exact_column_indexes:
|
|
self._remaining_color_group_count = 0
|
|
else:
|
|
self._remaining_color_group_count = self._original_color_group_count
|
|
if self._remaining_groups.shape != (self._state_count,):
|
|
raise SparseJacobianCompatibilityError(
|
|
"SciPy Jacobian column grouping returned an unexpected shape."
|
|
)
|
|
self._coloring_seed = self._COLORING_SEED
|
|
self._color_group_count = self._remaining_color_group_count
|
|
self._default_color_group_count = self._remaining_color_group_count
|
|
|
|
self._jacobian = None
|
|
self._original_factor = None
|
|
self._remaining_factor = None
|
|
self._last_factor = None
|
|
self._constructing = False
|
|
self._has_secant_update = False
|
|
self._consecutive_reuses = 0
|
|
self._pending_secants: list[
|
|
tuple[np.ndarray, np.ndarray, np.ndarray]
|
|
] = []
|
|
self._last_observation_time: float | None = None
|
|
self._last_observation_state: np.ndarray | None = None
|
|
self._last_observation_rhs: np.ndarray | None = None
|
|
|
|
self._full_build_count = 0
|
|
self._secant_reuse_count = 0
|
|
self._audit_failure_count = 0
|
|
self._finite_difference_rhs_count = 0
|
|
self._base_rhs_count = 0
|
|
self._jv_audit_count = 0
|
|
self._secant_update_count = 0
|
|
self._rejected_observation_count = 0
|
|
self._segment_start_count = 0
|
|
self._jacobian_evaluation_count = 0
|
|
self._assembly_seconds = 0.0
|
|
self._exact_column_build_count = 0
|
|
self._exact_column_fallback_count = 0
|
|
self._last_exact_column_fallback_reason: str | None = None
|
|
self._last_build_finite_difference_rhs_count = 0
|
|
self._last_decision = "uninitialized"
|
|
self._last_audit_relative_error: float | None = None
|
|
self._last_audit_rms_relative_error: float | None = None
|
|
self._last_audit_p95_relative_error: float | None = None
|
|
self._segments: list[dict[str, int | float]] = []
|
|
|
|
def _normalize_exact_column_indexes(
|
|
self,
|
|
values: Sequence[int] | object,
|
|
) -> tuple[int, ...]:
|
|
try:
|
|
requested = tuple(values) # type: ignore[arg-type]
|
|
except TypeError as exc:
|
|
raise TypeError(
|
|
"Exact Jacobian column indexes must be an iterable of integers."
|
|
) from exc
|
|
seen: set[int] = set()
|
|
for column_index in requested:
|
|
if (
|
|
not isinstance(column_index, int)
|
|
or isinstance(column_index, bool)
|
|
or not 0 <= column_index < self._state_count
|
|
):
|
|
raise ValueError(
|
|
f"Exact Jacobian column index {column_index!r} is out of range."
|
|
)
|
|
if column_index in seen:
|
|
raise ValueError(
|
|
f"Exact Jacobian column index {column_index} is duplicated."
|
|
)
|
|
seen.add(column_index)
|
|
return tuple(sorted(seen))
|
|
|
|
@staticmethod
|
|
def _load_scipy_helpers():
|
|
try:
|
|
from scipy.integrate._ivp.common import num_jac
|
|
from scipy.optimize._numdiff import group_columns
|
|
except (ImportError, AttributeError) as exc:
|
|
raise SparseJacobianCompatibilityError(
|
|
"The installed SciPy does not expose compatible numerical "
|
|
"Jacobian helpers."
|
|
) from exc
|
|
|
|
expected_num_jac = {
|
|
"fun",
|
|
"t",
|
|
"y",
|
|
"f",
|
|
"threshold",
|
|
"factor",
|
|
"sparsity",
|
|
}
|
|
expected_group_columns = {"A", "order"}
|
|
try:
|
|
num_jac_parameters = set(inspect.signature(num_jac).parameters)
|
|
group_parameters = set(inspect.signature(group_columns).parameters)
|
|
except (TypeError, ValueError) as exc:
|
|
raise SparseJacobianCompatibilityError(
|
|
"Unable to inspect SciPy numerical Jacobian helper signatures."
|
|
) from exc
|
|
if not expected_num_jac.issubset(num_jac_parameters) or not (
|
|
expected_group_columns.issubset(group_parameters)
|
|
):
|
|
raise SparseJacobianCompatibilityError(
|
|
"The installed SciPy numerical Jacobian helper API is unsupported."
|
|
)
|
|
return num_jac, group_columns
|
|
|
|
@staticmethod
|
|
def _validate_subset_sparse_compatibility() -> None:
|
|
"""Fail before integration when subset-FD private helpers are absent."""
|
|
|
|
try:
|
|
from scipy.integrate._ivp.common import (
|
|
EPS,
|
|
NUM_JAC_DIFF_BIG,
|
|
NUM_JAC_DIFF_REJECT,
|
|
NUM_JAC_DIFF_SMALL,
|
|
NUM_JAC_FACTOR_DECREASE,
|
|
NUM_JAC_FACTOR_INCREASE,
|
|
NUM_JAC_MIN_FACTOR,
|
|
)
|
|
from scipy.sparse import coo_matrix, csc_matrix, find
|
|
except (ImportError, AttributeError) as exc:
|
|
raise SparseJacobianCompatibilityError(
|
|
"The installed SciPy does not expose compatible subset "
|
|
"Jacobian helpers."
|
|
) from exc
|
|
del (
|
|
EPS,
|
|
NUM_JAC_DIFF_BIG,
|
|
NUM_JAC_DIFF_REJECT,
|
|
NUM_JAC_DIFF_SMALL,
|
|
NUM_JAC_FACTOR_DECREASE,
|
|
NUM_JAC_FACTOR_INCREASE,
|
|
NUM_JAC_MIN_FACTOR,
|
|
coo_matrix,
|
|
csc_matrix,
|
|
find,
|
|
)
|
|
|
|
def _subset_sparse_num_jac(
|
|
self,
|
|
fun,
|
|
time: float,
|
|
state: np.ndarray,
|
|
base_rhs: np.ndarray,
|
|
):
|
|
"""Run SciPy-compatible adaptive differences on active columns only."""
|
|
|
|
try:
|
|
from scipy.integrate._ivp.common import (
|
|
EPS,
|
|
NUM_JAC_DIFF_BIG,
|
|
NUM_JAC_DIFF_REJECT,
|
|
NUM_JAC_DIFF_SMALL,
|
|
NUM_JAC_FACTOR_DECREASE,
|
|
NUM_JAC_FACTOR_INCREASE,
|
|
NUM_JAC_MIN_FACTOR,
|
|
)
|
|
from scipy.sparse import coo_matrix, csc_matrix, find
|
|
except (ImportError, AttributeError) as exc: # pragma: no cover
|
|
raise SparseJacobianCompatibilityError(
|
|
"The installed SciPy does not expose compatible numerical "
|
|
"Jacobian constants."
|
|
) from exc
|
|
|
|
active_columns = self._remaining_finite_difference_columns
|
|
active_count = int(active_columns.size)
|
|
if active_count == 0: # pragma: no cover - guarded by _full_build.
|
|
return csc_matrix(
|
|
(self._state_count, self._state_count), dtype=float
|
|
), self._remaining_factor
|
|
|
|
if self._remaining_factor is None:
|
|
factor = np.full(self._state_count, EPS**0.5, dtype=float)
|
|
else:
|
|
factor = np.asarray(self._remaining_factor, dtype=float).copy()
|
|
if factor.shape != (self._state_count,):
|
|
raise SparseJacobianCompatibilityError(
|
|
"SciPy numerical Jacobian factor has an unexpected shape."
|
|
)
|
|
|
|
rhs_sign = 2.0 * (np.real(base_rhs) >= 0.0).astype(float) - 1.0
|
|
state_scale = rhs_sign * np.maximum(self._atol, np.abs(state))
|
|
step = np.zeros(self._state_count, dtype=float)
|
|
step[active_columns] = (
|
|
state[active_columns]
|
|
+ factor[active_columns] * state_scale[active_columns]
|
|
) - state[active_columns]
|
|
for column_index in active_columns[step[active_columns] == 0.0]:
|
|
while step[column_index] == 0.0:
|
|
factor[column_index] *= 10.0
|
|
step[column_index] = (
|
|
state[column_index]
|
|
+ factor[column_index] * state_scale[column_index]
|
|
) - state[column_index]
|
|
|
|
groups = self._remaining_groups[active_columns]
|
|
group_count = self._remaining_color_group_count
|
|
perturbations = np.zeros((self._state_count, group_count), dtype=float)
|
|
perturbations[active_columns, groups] = step[active_columns]
|
|
perturbed_rhs = fun(time, state[:, None] + perturbations)
|
|
differences = perturbed_rhs - base_rhs[:, None]
|
|
|
|
reduced_structure = self._remaining_finite_difference_sparsity[
|
|
:, active_columns
|
|
]
|
|
row_indexes, reduced_columns, _values = find(reduced_structure)
|
|
reduced_jacobian = coo_matrix(
|
|
(
|
|
differences[row_indexes, groups[reduced_columns]],
|
|
(row_indexes, reduced_columns),
|
|
),
|
|
shape=(self._state_count, active_count),
|
|
).tocsc()
|
|
maximum_rows = np.asarray(
|
|
abs(reduced_jacobian).argmax(axis=0)
|
|
).ravel()
|
|
reduced_range = np.arange(active_count)
|
|
maximum_difference = np.asarray(
|
|
np.abs(reduced_jacobian[maximum_rows, reduced_range])
|
|
).ravel()
|
|
difference_scale = np.maximum(
|
|
np.abs(base_rhs[maximum_rows]),
|
|
np.abs(perturbed_rhs[maximum_rows, groups]),
|
|
)
|
|
|
|
too_small = maximum_difference < NUM_JAC_DIFF_REJECT * difference_scale
|
|
if np.any(too_small):
|
|
rejected_reduced_columns = np.flatnonzero(too_small)
|
|
rejected_columns = active_columns[rejected_reduced_columns]
|
|
increased_factor = (
|
|
NUM_JAC_FACTOR_INCREASE * factor[rejected_columns]
|
|
)
|
|
increased_step = (
|
|
state[rejected_columns]
|
|
+ increased_factor * state_scale[rejected_columns]
|
|
) - state[rejected_columns]
|
|
increased_step_by_column = np.zeros(self._state_count, dtype=float)
|
|
increased_step_by_column[rejected_columns] = increased_step
|
|
|
|
rejected_groups = np.unique(groups[rejected_reduced_columns])
|
|
group_map = np.full(group_count, -1, dtype=int)
|
|
retry_perturbations = np.zeros(
|
|
(self._state_count, rejected_groups.size),
|
|
dtype=float,
|
|
)
|
|
for retry_index, group in enumerate(rejected_groups):
|
|
group_map[group] = retry_index
|
|
group_columns = active_columns[groups == group]
|
|
retry_perturbations[group_columns, retry_index] = (
|
|
increased_step_by_column[group_columns]
|
|
)
|
|
|
|
retry_rhs = fun(time, state[:, None] + retry_perturbations)
|
|
retry_differences = retry_rhs - base_rhs[:, None]
|
|
retry_structure = reduced_structure[:, rejected_reduced_columns]
|
|
retry_rows, retry_columns, _retry_values = find(retry_structure)
|
|
retry_group_columns = rejected_reduced_columns[retry_columns]
|
|
retry_jacobian = coo_matrix(
|
|
(
|
|
retry_differences[
|
|
retry_rows,
|
|
group_map[groups[retry_group_columns]],
|
|
],
|
|
(retry_rows, retry_columns),
|
|
),
|
|
shape=(self._state_count, rejected_reduced_columns.size),
|
|
).tocsc()
|
|
retry_maximum_rows = np.asarray(
|
|
abs(retry_jacobian).argmax(axis=0)
|
|
).ravel()
|
|
retry_range = np.arange(rejected_reduced_columns.size)
|
|
retry_maximum_difference = np.asarray(
|
|
np.abs(retry_jacobian[retry_maximum_rows, retry_range])
|
|
).ravel()
|
|
retry_scale = np.maximum(
|
|
np.abs(base_rhs[retry_maximum_rows]),
|
|
np.abs(
|
|
retry_rhs[
|
|
retry_maximum_rows,
|
|
group_map[groups[rejected_reduced_columns]],
|
|
]
|
|
),
|
|
)
|
|
use_retry = (
|
|
maximum_difference[rejected_reduced_columns] * retry_scale
|
|
< retry_maximum_difference
|
|
* difference_scale[rejected_reduced_columns]
|
|
)
|
|
if np.any(use_retry):
|
|
retry_selection = np.flatnonzero(use_retry)
|
|
selected_reduced_columns = rejected_reduced_columns[
|
|
retry_selection
|
|
]
|
|
selected_columns = active_columns[selected_reduced_columns]
|
|
factor[selected_columns] = increased_factor[retry_selection]
|
|
step[selected_columns] = increased_step[retry_selection]
|
|
reduced_jacobian[:, selected_reduced_columns] = retry_jacobian[
|
|
:, retry_selection
|
|
]
|
|
difference_scale[selected_reduced_columns] = retry_scale[
|
|
retry_selection
|
|
]
|
|
maximum_difference[selected_reduced_columns] = (
|
|
retry_maximum_difference[retry_selection]
|
|
)
|
|
|
|
reduced_jacobian.data /= np.repeat(
|
|
step[active_columns],
|
|
np.diff(reduced_jacobian.indptr),
|
|
)
|
|
factor[
|
|
active_columns[
|
|
maximum_difference < NUM_JAC_DIFF_SMALL * difference_scale
|
|
]
|
|
] *= NUM_JAC_FACTOR_INCREASE
|
|
factor[
|
|
active_columns[
|
|
maximum_difference > NUM_JAC_DIFF_BIG * difference_scale
|
|
]
|
|
] *= NUM_JAC_FACTOR_DECREASE
|
|
factor = np.maximum(factor, NUM_JAC_MIN_FACTOR)
|
|
|
|
rows, reduced_columns, values = find(reduced_jacobian)
|
|
jacobian = csc_matrix(
|
|
(
|
|
values,
|
|
(rows, active_columns[reduced_columns]),
|
|
),
|
|
shape=(self._state_count, self._state_count),
|
|
)
|
|
return jacobian, factor
|
|
|
|
@staticmethod
|
|
def _new_segment_diagnostics(index: int) -> dict[str, int | float]:
|
|
return {
|
|
"index": index,
|
|
"fullBuildCount": 0,
|
|
"secantReuseCount": 0,
|
|
"auditFailureCount": 0,
|
|
"exactColumnBuildCount": 0,
|
|
"exactColumnFallbackCount": 0,
|
|
"finiteDifferenceRhsEvaluationCount": 0,
|
|
"baseRhsEvaluationCount": 0,
|
|
"jvAuditRhsEvaluationCount": 0,
|
|
"jacobianEvaluationCount": 0,
|
|
"secantPairCount": 0,
|
|
"assemblySeconds": 0.0,
|
|
}
|
|
|
|
def _record_segment(self, key: str, value: int | float = 1) -> None:
|
|
if not self._segments:
|
|
self._segments.append(self._new_segment_diagnostics(0))
|
|
segment = self._segments[-1]
|
|
segment[key] = segment[key] + value
|
|
|
|
def start_segment(self) -> None:
|
|
"""Discard numerical state before a breakpoint/event solver restart."""
|
|
|
|
self._jacobian = None
|
|
self._original_factor = None
|
|
self._remaining_factor = None
|
|
self._last_factor = None
|
|
self._has_secant_update = False
|
|
self._consecutive_reuses = 0
|
|
self._pending_secants.clear()
|
|
self._clear_observation()
|
|
self._segment_start_count += 1
|
|
if not self._segments:
|
|
self._segments.append(self._new_segment_diagnostics(0))
|
|
elif any(
|
|
value
|
|
for key, value in self._segments[-1].items()
|
|
if key != "index"
|
|
):
|
|
self._segments.append(
|
|
self._new_segment_diagnostics(len(self._segments))
|
|
)
|
|
self._last_decision = "segmentReset"
|
|
self._last_audit_relative_error = None
|
|
self._last_audit_rms_relative_error = None
|
|
self._last_audit_p95_relative_error = None
|
|
|
|
def _clear_observation(self) -> None:
|
|
self._last_observation_time = None
|
|
self._last_observation_state = None
|
|
self._last_observation_rhs = None
|
|
|
|
def _state_array(self, values: Sequence[float] | np.ndarray) -> np.ndarray:
|
|
state = np.asarray(values, dtype=float)
|
|
if state.shape != (self._state_count,):
|
|
raise ValueError(
|
|
f"Expected a state vector of shape {(self._state_count,)}, "
|
|
f"received {state.shape}."
|
|
)
|
|
if not np.all(np.isfinite(state)):
|
|
raise ValueError("Jacobian state vectors must contain finite values.")
|
|
return state
|
|
|
|
def _rhs_array(self, values: Sequence[float] | np.ndarray) -> np.ndarray:
|
|
derivative = np.asarray(values, dtype=float)
|
|
if derivative.shape != (self._state_count,):
|
|
raise ValueError(
|
|
f"Expected an RHS vector of shape {(self._state_count,)}, "
|
|
f"received {derivative.shape}."
|
|
)
|
|
if not np.all(np.isfinite(derivative)):
|
|
raise ValueError("Jacobian RHS evaluations must contain finite values.")
|
|
return derivative
|
|
|
|
def _evaluate(
|
|
self,
|
|
time: float,
|
|
state: np.ndarray,
|
|
*,
|
|
purpose: str,
|
|
) -> np.ndarray:
|
|
if purpose == "base":
|
|
self._base_rhs_count += 1
|
|
self._record_segment("baseRhsEvaluationCount")
|
|
elif purpose == "finiteDifference":
|
|
self._finite_difference_rhs_count += 1
|
|
self._record_segment("finiteDifferenceRhsEvaluationCount")
|
|
elif purpose == "jvAudit":
|
|
self._jv_audit_count += 1
|
|
self._record_segment("jvAuditRhsEvaluationCount")
|
|
else: # pragma: no cover - internal programming error.
|
|
raise RuntimeError(f"Unknown Jacobian evaluation purpose: {purpose}.")
|
|
return self._rhs_array(self._evaluate_rhs(float(time), state.copy()))
|
|
|
|
def _remember_observation(
|
|
self,
|
|
time: float,
|
|
state: np.ndarray,
|
|
derivative: np.ndarray,
|
|
) -> None:
|
|
self._last_observation_time = float(time)
|
|
self._last_observation_state = state.copy()
|
|
self._last_observation_rhs = derivative.copy()
|
|
|
|
def observe(
|
|
self,
|
|
time: float,
|
|
state: Sequence[float] | np.ndarray,
|
|
derivative: Sequence[float] | np.ndarray,
|
|
) -> None:
|
|
"""Cache one ordinary RHS sample for a later same-time sparse secant."""
|
|
|
|
if self._constructing or self._max_consecutive_reuses == 0:
|
|
return
|
|
state_array = self._state_array(state)
|
|
derivative_array = self._rhs_array(derivative)
|
|
time_value = float(time)
|
|
if not math.isfinite(time_value):
|
|
raise ValueError("Jacobian observation time must be finite.")
|
|
|
|
previous_state = self._last_observation_state
|
|
previous_rhs = self._last_observation_rhs
|
|
same_time = (
|
|
self._last_observation_time is not None
|
|
and time_value == self._last_observation_time
|
|
)
|
|
if not same_time:
|
|
self._pending_secants.clear()
|
|
if (
|
|
self._jacobian is None
|
|
or not same_time
|
|
or previous_state is None
|
|
or previous_rhs is None
|
|
):
|
|
self._remember_observation(time_value, state_array, derivative_array)
|
|
return
|
|
|
|
state_delta = state_array - previous_state
|
|
rhs_delta = derivative_array - previous_rhs
|
|
state_scale = np.maximum.reduce(
|
|
(
|
|
np.abs(state_array),
|
|
np.abs(previous_state),
|
|
self._atol,
|
|
)
|
|
)
|
|
normalized_step = np.max(np.abs(state_delta) / state_scale)
|
|
if (
|
|
not math.isfinite(float(normalized_step))
|
|
or normalized_step <= 32.0 * np.finfo(float).eps
|
|
):
|
|
self._rejected_observation_count += 1
|
|
self._remember_observation(time_value, state_array, derivative_array)
|
|
return
|
|
|
|
self._pending_secants.append(
|
|
(state_delta.copy(), rhs_delta.copy(), state_scale.copy())
|
|
)
|
|
if len(self._pending_secants) > self._MAX_PENDING_SECANT_PAIRS:
|
|
del self._pending_secants[0]
|
|
self._record_segment("secantPairCount")
|
|
self._remember_observation(time_value, state_array, derivative_array)
|
|
|
|
def _apply_pending_secants(self) -> None:
|
|
"""Apply cached same-time pairs only when a new Jacobian is requested."""
|
|
|
|
if self._jacobian is None or not self._pending_secants:
|
|
return
|
|
jacobian = self._jacobian.tocsr(copy=True)
|
|
updated_row_count = 0
|
|
for state_delta, rhs_delta, state_scale in self._pending_secants:
|
|
normalized_step = state_delta / state_scale
|
|
for row_index in range(self._state_count):
|
|
if row_index in self._exact_row_indexes:
|
|
continue
|
|
start = int(jacobian.indptr[row_index])
|
|
stop = int(jacobian.indptr[row_index + 1])
|
|
if start == stop:
|
|
continue
|
|
columns = jacobian.indices[start:stop]
|
|
mutable = np.asarray(
|
|
[
|
|
column_index not in self._exact_column_index_set
|
|
for column_index in columns
|
|
],
|
|
dtype=bool,
|
|
)
|
|
if not np.any(mutable):
|
|
continue
|
|
mutable_columns = columns[mutable]
|
|
local_normalized_step = normalized_step[mutable_columns]
|
|
denominator = float(
|
|
np.dot(local_normalized_step, local_normalized_step)
|
|
)
|
|
if not math.isfinite(denominator) or denominator <= 0.0:
|
|
continue
|
|
row_values = jacobian.data[start:stop]
|
|
predicted_delta = float(
|
|
np.dot(row_values, state_delta[columns])
|
|
)
|
|
correction_scale = (
|
|
float(rhs_delta[row_index]) - predicted_delta
|
|
) / denominator
|
|
# Apply Broyden to A = J * diag(state_scale), then transform
|
|
# the corrected row back to the physical Jacobian coordinates.
|
|
correction = (
|
|
correction_scale
|
|
* local_normalized_step
|
|
/ state_scale[mutable_columns]
|
|
)
|
|
if not np.all(np.isfinite(correction)):
|
|
continue
|
|
row_values[mutable] = row_values[mutable] + correction
|
|
updated_row_count += 1
|
|
|
|
self._pending_secants.clear()
|
|
|
|
if updated_row_count:
|
|
self._jacobian = jacobian
|
|
self._has_secant_update = True
|
|
self._secant_update_count += 1
|
|
else:
|
|
self._rejected_observation_count += 1
|
|
|
|
def _resolve_exact_row(
|
|
self,
|
|
row_index: int,
|
|
time: float,
|
|
state: np.ndarray,
|
|
) -> np.ndarray:
|
|
specification = self._exact_rows[row_index]
|
|
values = (
|
|
specification(time, state.copy())
|
|
if callable(specification)
|
|
else specification
|
|
)
|
|
row = np.zeros(self._state_count, dtype=float)
|
|
if isinstance(values, Mapping):
|
|
for column_index, value in values.items():
|
|
if (
|
|
not isinstance(column_index, int)
|
|
or isinstance(column_index, bool)
|
|
or not 0 <= column_index < self._state_count
|
|
):
|
|
raise ValueError(
|
|
f"Exact Jacobian column index {column_index!r} is out of range."
|
|
)
|
|
row[column_index] = float(value)
|
|
else:
|
|
dense_values = np.asarray(values, dtype=float)
|
|
if dense_values.shape != (self._state_count,):
|
|
raise ValueError(
|
|
f"Exact Jacobian row {row_index} must have shape "
|
|
f"{(self._state_count,)}, received {dense_values.shape}."
|
|
)
|
|
row[:] = dense_values
|
|
if not np.all(np.isfinite(row)):
|
|
raise ValueError("Exact Jacobian rows must contain finite values.")
|
|
return row
|
|
|
|
def _with_exact_rows(self, jacobian, time: float, state: np.ndarray):
|
|
if not self._exact_rows:
|
|
return jacobian.tocsr(copy=True)
|
|
result = jacobian.tolil(copy=True)
|
|
for row_index in sorted(self._exact_rows):
|
|
row = self._resolve_exact_row(row_index, time, state)
|
|
columns = np.flatnonzero(row)
|
|
result.rows[row_index] = [int(column) for column in columns]
|
|
result.data[row_index] = [float(row[column]) for column in columns]
|
|
return result.tocsr()
|
|
|
|
def _resolve_exact_columns(
|
|
self,
|
|
time: float,
|
|
state: np.ndarray,
|
|
) -> np.ndarray:
|
|
column_count = len(self._exact_column_indexes)
|
|
if column_count == 0:
|
|
self._exact_column_outside_pattern_nonzero_count = 0
|
|
return np.empty((self._state_count, 0), dtype=float)
|
|
|
|
if self._fixed_exact_column_values is not None:
|
|
columns = self._fixed_exact_column_values.copy()
|
|
else:
|
|
assert self._exact_column_provider is not None
|
|
values = self._exact_column_provider(
|
|
float(time),
|
|
state.copy(),
|
|
self._exact_column_indexes,
|
|
)
|
|
if isinstance(values, Mapping):
|
|
returned_indexes = tuple(values)
|
|
if returned_indexes != self._exact_column_indexes:
|
|
raise ValueError(
|
|
"Exact Jacobian column provider mapping keys must match "
|
|
"the requested normalized column order."
|
|
)
|
|
provider_columns: list[np.ndarray] = []
|
|
for column_index in returned_indexes:
|
|
column = np.asarray(values[column_index], dtype=float)
|
|
if column.shape != (self._state_count,):
|
|
raise ValueError(
|
|
f"Exact Jacobian column {column_index} must have shape "
|
|
f"{(self._state_count,)}, received {column.shape}."
|
|
)
|
|
provider_columns.append(column)
|
|
columns = np.column_stack(provider_columns)
|
|
else:
|
|
columns = np.asarray(values, dtype=float)
|
|
|
|
expected_shape = (self._state_count, column_count)
|
|
if columns.shape != expected_shape:
|
|
raise ValueError(
|
|
"Exact Jacobian column provider must return shape "
|
|
f"{expected_shape}, received {columns.shape}."
|
|
)
|
|
if not np.all(np.isfinite(columns)):
|
|
raise ValueError("Exact Jacobian columns must contain finite values.")
|
|
self._exact_column_outside_pattern_nonzero_count = int(
|
|
np.count_nonzero((columns != 0.0) & ~self._exact_column_pattern)
|
|
)
|
|
return columns
|
|
|
|
def _evaluate_exact_column_base_rhs(
|
|
self,
|
|
time: float,
|
|
state: np.ndarray,
|
|
) -> np.ndarray:
|
|
"""Evaluate a base RHS while allowing a provider one-shot capture."""
|
|
|
|
request_capture = getattr(
|
|
self._exact_column_provider,
|
|
"request_primal_capture",
|
|
None,
|
|
)
|
|
cancel_capture = getattr(
|
|
self._exact_column_provider,
|
|
"cancel_primal_capture",
|
|
None,
|
|
)
|
|
try:
|
|
if request_capture is not None:
|
|
request_capture()
|
|
return self._evaluate(time, state, purpose="base")
|
|
finally:
|
|
if cancel_capture is not None:
|
|
cancel_capture()
|
|
|
|
def _with_exact_columns(
|
|
self,
|
|
jacobian,
|
|
time: float,
|
|
state: np.ndarray,
|
|
*,
|
|
resolved_columns: np.ndarray | None = None,
|
|
):
|
|
if not self._exact_column_indexes:
|
|
self._exact_column_outside_pattern_nonzero_count = 0
|
|
return jacobian.tocsr(copy=True)
|
|
|
|
from scipy.sparse import csc_matrix, diags
|
|
|
|
columns = (
|
|
self._resolve_exact_columns(time, state)
|
|
if resolved_columns is None
|
|
else resolved_columns
|
|
)
|
|
keep_column = np.ones(self._state_count, dtype=float)
|
|
keep_column[list(self._exact_column_indexes)] = 0.0
|
|
remaining = jacobian.tocsc(copy=True) @ diags(
|
|
keep_column,
|
|
format="csc",
|
|
)
|
|
row_indexes = np.tile(
|
|
np.arange(self._state_count, dtype=int),
|
|
len(self._exact_column_indexes),
|
|
)
|
|
column_indexes = np.repeat(
|
|
np.asarray(self._exact_column_indexes, dtype=int),
|
|
self._state_count,
|
|
)
|
|
exact = csc_matrix(
|
|
(
|
|
columns.ravel(order="F"),
|
|
(row_indexes, column_indexes),
|
|
),
|
|
shape=(self._state_count, self._state_count),
|
|
)
|
|
exact.eliminate_zeros()
|
|
return (remaining + exact).tocsr()
|
|
|
|
def _with_exact_overrides(
|
|
self,
|
|
jacobian,
|
|
time: float,
|
|
state: np.ndarray,
|
|
*,
|
|
resolved_columns: np.ndarray | None = None,
|
|
):
|
|
# Full columns are assembled first. Full rows intentionally win at
|
|
# row/column intersections, matching the documented precedence.
|
|
with_columns = self._with_exact_columns(
|
|
jacobian,
|
|
time,
|
|
state,
|
|
resolved_columns=resolved_columns,
|
|
)
|
|
return self._with_exact_rows(with_columns, time, state)
|
|
|
|
def _record_exact_column_fallback(
|
|
self,
|
|
unavailable: ExactColumnsUnavailable,
|
|
) -> None:
|
|
self._exact_column_fallback_count += 1
|
|
self._last_exact_column_fallback_reason = unavailable.reason
|
|
self._record_segment("exactColumnFallbackCount")
|
|
|
|
def _full_build(
|
|
self,
|
|
time: float,
|
|
state: np.ndarray,
|
|
*,
|
|
decision: str = "fullBuild",
|
|
base_rhs: np.ndarray | None = None,
|
|
exact_column_unavailable: ExactColumnsUnavailable | None = None,
|
|
):
|
|
from scipy.sparse import csr_matrix
|
|
|
|
finite_difference_rhs_before = self._finite_difference_rhs_count
|
|
self._constructing = True
|
|
try:
|
|
resolved_exact_columns = None
|
|
unavailable = exact_column_unavailable
|
|
if self._exact_column_indexes and unavailable is None:
|
|
# Tangent providers operate on the closure left by a real
|
|
# primal RHS at exactly (time, state). Resolve them before any
|
|
# grouped perturbation dirties the shared model.
|
|
if base_rhs is None:
|
|
base_rhs = self._evaluate_exact_column_base_rhs(
|
|
time,
|
|
state,
|
|
)
|
|
try:
|
|
resolved_exact_columns = self._resolve_exact_columns(
|
|
time,
|
|
state,
|
|
)
|
|
except ExactColumnsUnavailable as exc:
|
|
unavailable = exc
|
|
|
|
use_exact_columns = bool(
|
|
self._exact_column_indexes and unavailable is None
|
|
)
|
|
if unavailable is not None:
|
|
self._record_exact_column_fallback(unavailable)
|
|
self._exact_column_outside_pattern_nonzero_count = 0
|
|
if decision == "fullBuild":
|
|
decision = "exactColumnFallback"
|
|
|
|
finite_difference_sparsity = (
|
|
self._remaining_finite_difference_sparsity
|
|
if use_exact_columns
|
|
else self._original_finite_difference_sparsity
|
|
)
|
|
if finite_difference_sparsity.nnz == 0:
|
|
selected_factor = (
|
|
self._remaining_factor
|
|
if use_exact_columns
|
|
else self._original_factor
|
|
)
|
|
self._last_factor = (
|
|
None if selected_factor is None else selected_factor.copy()
|
|
)
|
|
jacobian = csr_matrix(
|
|
(self._state_count, self._state_count), dtype=float
|
|
)
|
|
else:
|
|
if base_rhs is None:
|
|
base_rhs = self._evaluate(time, state, purpose="base")
|
|
|
|
def vectorized_evaluate(
|
|
evaluation_time: float,
|
|
states: np.ndarray,
|
|
) -> np.ndarray:
|
|
states_array = np.asarray(states, dtype=float)
|
|
if states_array.ndim == 1:
|
|
return self._evaluate(
|
|
evaluation_time,
|
|
states_array.copy(),
|
|
purpose="finiteDifference",
|
|
)
|
|
if (
|
|
states_array.ndim != 2
|
|
or states_array.shape[0] != self._state_count
|
|
):
|
|
raise ValueError(
|
|
"SciPy requested an unexpected vectorized Jacobian shape."
|
|
)
|
|
return np.column_stack(
|
|
[
|
|
self._evaluate(
|
|
evaluation_time,
|
|
states_array[:, column_index].copy(),
|
|
purpose="finiteDifference",
|
|
)
|
|
for column_index in range(states_array.shape[1])
|
|
]
|
|
)
|
|
|
|
if use_exact_columns:
|
|
jacobian, self._remaining_factor = self._subset_sparse_num_jac(
|
|
vectorized_evaluate,
|
|
float(time),
|
|
state,
|
|
base_rhs,
|
|
)
|
|
self._last_factor = self._remaining_factor.copy()
|
|
else:
|
|
jacobian, self._original_factor = self._num_jac(
|
|
fun=vectorized_evaluate,
|
|
t=float(time),
|
|
y=state,
|
|
f=base_rhs,
|
|
threshold=self._atol,
|
|
factor=self._original_factor,
|
|
sparsity=(
|
|
self._original_finite_difference_sparsity,
|
|
self._original_groups,
|
|
),
|
|
)
|
|
self._last_factor = self._original_factor.copy()
|
|
jacobian = csr_matrix(jacobian, dtype=float)
|
|
if not np.all(np.isfinite(jacobian.data)):
|
|
raise ValueError(
|
|
"Numerical Jacobian construction returned non-finite values."
|
|
)
|
|
if use_exact_columns:
|
|
jacobian = self._with_exact_overrides(
|
|
jacobian,
|
|
time,
|
|
state,
|
|
resolved_columns=resolved_exact_columns,
|
|
)
|
|
else:
|
|
# This is exactly the no-provider seed-0 numerical path. Full
|
|
# rows may still overwrite their numerical values, but no
|
|
# reduced column pattern leaks into the fallback assembly.
|
|
jacobian = self._with_exact_rows(jacobian, time, state)
|
|
finally:
|
|
self._constructing = False
|
|
|
|
self._jacobian = jacobian
|
|
if use_exact_columns:
|
|
self._exact_column_build_count += 1
|
|
self._record_segment("exactColumnBuildCount")
|
|
self._last_build_finite_difference_rhs_count = (
|
|
self._finite_difference_rhs_count - finite_difference_rhs_before
|
|
)
|
|
self._full_build_count += 1
|
|
self._record_segment("fullBuildCount")
|
|
self._has_secant_update = False
|
|
self._consecutive_reuses = 0
|
|
self._pending_secants.clear()
|
|
self._last_decision = decision
|
|
if decision != "auditFallback":
|
|
self._last_audit_relative_error = None
|
|
self._last_audit_rms_relative_error = None
|
|
self._last_audit_p95_relative_error = None
|
|
return jacobian.copy()
|
|
|
|
def _audit_step(self, state: np.ndarray) -> np.ndarray:
|
|
epsilon_root = math.sqrt(np.finfo(float).eps)
|
|
factor = (
|
|
np.asarray(self._last_factor, dtype=float)
|
|
if self._last_factor is not None
|
|
else None
|
|
)
|
|
if factor is None or factor.shape != (self._state_count,):
|
|
factor = np.full(self._state_count, epsilon_root, dtype=float)
|
|
factor = np.where(
|
|
np.isfinite(factor) & (factor > 0.0),
|
|
factor,
|
|
epsilon_root,
|
|
)
|
|
magnitude = factor * np.maximum(np.abs(state), self._atol)
|
|
indexes = np.arange(self._state_count)
|
|
signs = np.where(indexes % 2 == 0, 1.0, -1.0)
|
|
# Do not cross zero for tiny positive masses/energies or other bounded
|
|
# coordinates. Away from zero, alternating signs reduce cancellation.
|
|
close_to_zero = np.abs(state) <= 4.0 * magnitude
|
|
signs = np.where(
|
|
close_to_zero,
|
|
np.where(state < 0.0, -1.0, 1.0),
|
|
signs,
|
|
)
|
|
requested_step = signs * magnitude
|
|
return (state + requested_step) - state
|
|
|
|
def _audit_candidate(
|
|
self,
|
|
time: float,
|
|
state: np.ndarray,
|
|
jacobian,
|
|
*,
|
|
base_rhs: np.ndarray | None = None,
|
|
) -> tuple[bool, np.ndarray | None]:
|
|
try:
|
|
if base_rhs is None:
|
|
base_rhs = self._evaluate(time, state, purpose="base")
|
|
step = self._audit_step(state)
|
|
perturbed_rhs = self._evaluate(
|
|
time,
|
|
state + step,
|
|
purpose="jvAudit",
|
|
)
|
|
except RecoverableTrialStateError:
|
|
self._last_audit_relative_error = math.inf
|
|
self._last_audit_rms_relative_error = math.inf
|
|
self._last_audit_p95_relative_error = math.inf
|
|
return False, base_rhs
|
|
|
|
observed_delta = perturbed_rhs - base_rhs
|
|
predicted_delta = np.asarray(jacobian @ step, dtype=float).reshape(-1)
|
|
audited_rows = np.ones(self._state_count, dtype=bool)
|
|
if self._exact_row_indexes:
|
|
audited_rows[list(self._exact_row_indexes)] = False
|
|
if not np.any(audited_rows):
|
|
self._last_audit_relative_error = 0.0
|
|
self._last_audit_rms_relative_error = 0.0
|
|
self._last_audit_p95_relative_error = 0.0
|
|
return True, base_rhs
|
|
|
|
observed = observed_delta[audited_rows]
|
|
predicted = predicted_delta[audited_rows]
|
|
error = np.abs(observed - predicted)
|
|
delta_scale = np.maximum(np.abs(observed), np.abs(predicted))
|
|
rhs_scale = np.maximum(
|
|
np.abs(base_rhs[audited_rows]),
|
|
np.abs(perturbed_rhs[audited_rows]),
|
|
)
|
|
roundoff_allowance = 64.0 * np.finfo(float).eps * np.maximum(
|
|
rhs_scale,
|
|
1.0,
|
|
)
|
|
tolerance = (
|
|
self._audit_absolute_tolerance
|
|
+ roundoff_allowance
|
|
+ self._audit_relative_tolerance * delta_scale
|
|
)
|
|
informative = delta_scale > (
|
|
roundoff_allowance + self._audit_absolute_tolerance
|
|
)
|
|
# A direction whose observed and predicted deltas are both buried in
|
|
# closure/floating-point noise did not validate the matrix. Treat it
|
|
# as inconclusive and refresh instead of accepting a vacuous 0 == 0.
|
|
passed = bool(np.any(informative) and np.all(error <= tolerance))
|
|
relative_denominator = np.maximum(
|
|
delta_scale,
|
|
roundoff_allowance + self._audit_absolute_tolerance,
|
|
)
|
|
relative_error = error / relative_denominator
|
|
self._last_audit_relative_error = float(
|
|
np.max(relative_error, initial=0.0)
|
|
)
|
|
self._last_audit_rms_relative_error = float(
|
|
np.sqrt(np.mean(relative_error * relative_error))
|
|
)
|
|
self._last_audit_p95_relative_error = float(
|
|
np.percentile(relative_error, 95.0)
|
|
)
|
|
return passed, base_rhs
|
|
|
|
def _evaluate_jacobian(
|
|
self,
|
|
time_value: float,
|
|
state_array: np.ndarray,
|
|
):
|
|
if self._jacobian is None:
|
|
return self._full_build(time_value, state_array)
|
|
|
|
if self._consecutive_reuses >= self._max_consecutive_reuses:
|
|
return self._full_build(time_value, state_array)
|
|
|
|
self._apply_pending_secants()
|
|
if not self._has_secant_update:
|
|
return self._full_build(time_value, state_array)
|
|
|
|
exact_base_rhs = None
|
|
resolved_exact_columns = None
|
|
if self._exact_column_indexes:
|
|
exact_base_rhs = self._evaluate_exact_column_base_rhs(
|
|
time_value,
|
|
state_array,
|
|
)
|
|
try:
|
|
resolved_exact_columns = self._resolve_exact_columns(
|
|
time_value,
|
|
state_array,
|
|
)
|
|
except ExactColumnsUnavailable as exc:
|
|
return self._full_build(
|
|
time_value,
|
|
state_array,
|
|
decision="exactColumnFallback",
|
|
base_rhs=exact_base_rhs,
|
|
exact_column_unavailable=exc,
|
|
)
|
|
candidate = self._with_exact_overrides(
|
|
self._jacobian,
|
|
time_value,
|
|
state_array,
|
|
resolved_columns=resolved_exact_columns,
|
|
)
|
|
audit_passed, _base_rhs = self._audit_candidate(
|
|
time_value,
|
|
state_array,
|
|
candidate,
|
|
base_rhs=exact_base_rhs,
|
|
)
|
|
if not audit_passed:
|
|
self._audit_failure_count += 1
|
|
self._record_segment("auditFailureCount")
|
|
return self._full_build(
|
|
time_value,
|
|
state_array,
|
|
decision="auditFallback",
|
|
)
|
|
|
|
self._jacobian = candidate
|
|
self._secant_reuse_count += 1
|
|
self._record_segment("secantReuseCount")
|
|
self._consecutive_reuses += 1
|
|
self._has_secant_update = False
|
|
self._last_decision = "secantReuse"
|
|
return candidate.copy()
|
|
|
|
def __call__(
|
|
self,
|
|
time: float,
|
|
state: Sequence[float] | np.ndarray,
|
|
):
|
|
"""Return a CSR Jacobian suitable for SciPy BDF/Radau."""
|
|
|
|
time_value = float(time)
|
|
if not math.isfinite(time_value):
|
|
raise ValueError("Jacobian evaluation time must be finite.")
|
|
state_array = self._state_array(state)
|
|
started_at = perf_counter()
|
|
self._jacobian_evaluation_count += 1
|
|
self._record_segment("jacobianEvaluationCount")
|
|
try:
|
|
return self._evaluate_jacobian(time_value, state_array)
|
|
finally:
|
|
elapsed = perf_counter() - started_at
|
|
self._assembly_seconds += elapsed
|
|
self._record_segment("assemblySeconds", elapsed)
|
|
|
|
def diagnostics(self) -> dict[str, object]:
|
|
"""Return cumulative work counters and the latest reuse decision."""
|
|
|
|
return {
|
|
"mode": (
|
|
"hybridSparseSecant"
|
|
if self._max_consecutive_reuses
|
|
else "optimizedSparseFiniteDifference"
|
|
),
|
|
"fullBuildCount": self._full_build_count,
|
|
"secantReuseCount": self._secant_reuse_count,
|
|
"auditFailureCount": self._audit_failure_count,
|
|
"finiteDifferenceRhsEvaluationCount": (
|
|
self._finite_difference_rhs_count
|
|
),
|
|
"baseRhsEvaluationCount": self._base_rhs_count,
|
|
"jvAuditEvaluationCount": self._jv_audit_count,
|
|
"secantUpdateCount": self._secant_update_count,
|
|
"rejectedObservationCount": self._rejected_observation_count,
|
|
"segmentStartCount": self._segment_start_count,
|
|
"jacobianEvaluationCount": self._jacobian_evaluation_count,
|
|
"assemblySeconds": self._assembly_seconds,
|
|
"colorGroupCount": self._color_group_count,
|
|
"defaultColorGroupCount": self._default_color_group_count,
|
|
"originalColorGroupCount": self._original_color_group_count,
|
|
"remainingColorGroupCount": self._remaining_color_group_count,
|
|
"coloringSeed": self._coloring_seed,
|
|
"exactRowCount": len(self._exact_rows),
|
|
"exactColumnCount": len(self._exact_column_indexes),
|
|
"finiteDifferenceColumnCount": self._finite_difference_column_count,
|
|
"exactColumnOutsidePatternNonzeroCount": (
|
|
self._exact_column_outside_pattern_nonzero_count
|
|
),
|
|
"exactColumnBuildCount": self._exact_column_build_count,
|
|
"exactColumnFallbackCount": self._exact_column_fallback_count,
|
|
"lastExactColumnFallbackReason": (
|
|
self._last_exact_column_fallback_reason
|
|
),
|
|
"lastBuildFiniteDifferenceRhsEvaluationCount": (
|
|
self._last_build_finite_difference_rhs_count
|
|
),
|
|
"consecutiveReuseCount": self._consecutive_reuses,
|
|
"lastDecision": self._last_decision,
|
|
"lastAuditRelativeError": self._last_audit_relative_error,
|
|
"lastAuditRmsRelativeError": self._last_audit_rms_relative_error,
|
|
"lastAuditP95RelativeError": self._last_audit_p95_relative_error,
|
|
"segments": [dict(segment) for segment in self._segments],
|
|
}
|