"""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], }