951 lines
34 KiB
Python
951 lines
34 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from scipy.integrate._ivp.common import num_jac
|
|
from scipy.optimize._numdiff import group_columns
|
|
from scipy.sparse import csc_matrix, csr_matrix
|
|
|
|
from app.simulation.solvers.jacobian import (
|
|
ExactColumnsUnavailable,
|
|
SparseSecantJacobian,
|
|
)
|
|
from app.simulation.solvers.solver import IntegrationCancelled
|
|
|
|
|
|
class _SwitchableLinearRhs:
|
|
def __init__(self, matrix: np.ndarray) -> None:
|
|
self.matrix = np.asarray(matrix, dtype=float)
|
|
self.evaluation_count = 0
|
|
|
|
def __call__(self, time: float, state: np.ndarray) -> np.ndarray:
|
|
del time
|
|
self.evaluation_count += 1
|
|
return self.matrix @ np.asarray(state, dtype=float)
|
|
|
|
|
|
class SparseSecantJacobianTests(unittest.TestCase):
|
|
def test_first_call_builds_full_sparse_finite_difference_jacobian(self) -> None:
|
|
matrix = np.asarray(
|
|
(
|
|
(2.0, 0.0, -1.0),
|
|
(0.0, 3.0, 0.0),
|
|
(4.0, 0.0, 5.0),
|
|
)
|
|
)
|
|
evaluator = _SwitchableLinearRhs(matrix)
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csr_matrix(matrix != 0.0),
|
|
atol=np.full(3, 1.0e-8),
|
|
)
|
|
|
|
jacobian = builder(0.0, np.asarray((1.0, -2.0, 0.5)))
|
|
|
|
np.testing.assert_allclose(jacobian.toarray(), matrix, rtol=1.0e-7)
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["fullBuildCount"], 1)
|
|
self.assertEqual(diagnostics["secantReuseCount"], 0)
|
|
self.assertEqual(diagnostics["baseRhsEvaluationCount"], 1)
|
|
self.assertGreater(
|
|
diagnostics["finiteDifferenceRhsEvaluationCount"],
|
|
0,
|
|
)
|
|
self.assertEqual(
|
|
evaluator.evaluation_count,
|
|
diagnostics["baseRhsEvaluationCount"]
|
|
+ diagnostics["finiteDifferenceRhsEvaluationCount"],
|
|
)
|
|
|
|
def test_same_time_secant_is_reused_once_after_directional_audit(self) -> None:
|
|
matrix = np.asarray(((2.0, -1.0), (0.5, 4.0)))
|
|
evaluator = _SwitchableLinearRhs(matrix)
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csr_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=1.0e-8,
|
|
)
|
|
builder(0.0, np.asarray((1.0, 1.0)))
|
|
first = np.asarray((1.5, -0.5))
|
|
second = np.asarray((1.75, -0.25))
|
|
builder.observe(0.1, first, matrix @ first)
|
|
builder.observe(0.1, second, matrix @ second)
|
|
|
|
reused = builder(0.1, second)
|
|
|
|
np.testing.assert_allclose(reused.toarray(), matrix, rtol=1.0e-7)
|
|
after_reuse = builder.diagnostics()
|
|
self.assertEqual(after_reuse["fullBuildCount"], 1)
|
|
self.assertEqual(after_reuse["secantReuseCount"], 1)
|
|
self.assertEqual(after_reuse["jvAuditEvaluationCount"], 1)
|
|
self.assertEqual(after_reuse["lastDecision"], "secantReuse")
|
|
self.assertEqual(after_reuse["jacobianEvaluationCount"], 2)
|
|
self.assertEqual(
|
|
after_reuse["segments"][0]["jvAuditRhsEvaluationCount"],
|
|
1,
|
|
)
|
|
self.assertGreater(after_reuse["segments"][0]["assemblySeconds"], 0.0)
|
|
|
|
rebuilt = builder(0.1, second)
|
|
|
|
np.testing.assert_allclose(rebuilt.toarray(), matrix, rtol=1.0e-7)
|
|
self.assertEqual(builder.diagnostics()["fullBuildCount"], 2)
|
|
|
|
def test_failed_directional_audit_falls_back_to_full_refresh(self) -> None:
|
|
initial_matrix = np.eye(2)
|
|
changed_matrix = np.asarray(((1.0, 0.0), (0.0, 10.0)))
|
|
evaluator = _SwitchableLinearRhs(initial_matrix)
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csr_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=1.0e-8,
|
|
audit_relative_tolerance=1.0e-3,
|
|
)
|
|
builder(0.0, np.asarray((1.0, 1.0)))
|
|
evaluator.matrix = changed_matrix
|
|
first = np.asarray((1.0, 1.0))
|
|
second = np.asarray((2.0, 1.0))
|
|
builder.observe(0.2, first, changed_matrix @ first)
|
|
builder.observe(0.2, second, changed_matrix @ second)
|
|
|
|
refreshed = builder(0.2, second)
|
|
|
|
np.testing.assert_allclose(
|
|
refreshed.toarray(),
|
|
changed_matrix,
|
|
rtol=1.0e-7,
|
|
)
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["fullBuildCount"], 2)
|
|
self.assertEqual(diagnostics["secantReuseCount"], 0)
|
|
self.assertEqual(diagnostics["auditFailureCount"], 1)
|
|
self.assertEqual(diagnostics["jvAuditEvaluationCount"], 1)
|
|
self.assertEqual(diagnostics["lastDecision"], "auditFallback")
|
|
|
|
def test_directional_audit_does_not_swallow_cancellation(self) -> None:
|
|
matrix = np.asarray(((2.0, -1.0), (0.5, 4.0)))
|
|
cancel_next = False
|
|
|
|
def evaluator(_time: float, state: np.ndarray) -> np.ndarray:
|
|
nonlocal cancel_next
|
|
if cancel_next:
|
|
cancel_next = False
|
|
raise IntegrationCancelled
|
|
return matrix @ np.asarray(state, dtype=float)
|
|
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csr_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=1.0e-8,
|
|
)
|
|
builder(0.0, np.asarray((1.0, 1.0)))
|
|
first = np.asarray((1.5, -0.5))
|
|
second = np.asarray((1.75, -0.25))
|
|
builder.observe(0.1, first, matrix @ first)
|
|
builder.observe(0.1, second, matrix @ second)
|
|
cancel_next = True
|
|
|
|
with self.assertRaises(IntegrationCancelled):
|
|
builder(0.1, second)
|
|
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["fullBuildCount"], 1)
|
|
self.assertEqual(diagnostics["auditFailureCount"], 0)
|
|
self.assertEqual(diagnostics["secantReuseCount"], 0)
|
|
|
|
def test_segment_reset_forces_a_new_full_build(self) -> None:
|
|
matrix = np.asarray(((3.0, 0.0), (0.0, -2.0)))
|
|
evaluator = _SwitchableLinearRhs(matrix)
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csr_matrix(matrix != 0.0),
|
|
atol=1.0e-8,
|
|
)
|
|
state = np.asarray((2.0, 4.0))
|
|
builder(0.0, state)
|
|
|
|
builder.start_segment()
|
|
rebuilt = builder(1.0, state)
|
|
|
|
np.testing.assert_allclose(rebuilt.toarray(), matrix, rtol=1.0e-7)
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["segmentStartCount"], 1)
|
|
self.assertEqual(diagnostics["fullBuildCount"], 2)
|
|
self.assertEqual(diagnostics["lastDecision"], "fullBuild")
|
|
self.assertEqual(len(diagnostics["segments"]), 2)
|
|
self.assertEqual(diagnostics["segments"][0]["fullBuildCount"], 1)
|
|
self.assertEqual(diagnostics["segments"][1]["fullBuildCount"], 1)
|
|
|
|
def test_exact_row_replaces_finite_difference_and_secant_row(self) -> None:
|
|
def nonlinear_rhs(time: float, state: np.ndarray) -> np.ndarray:
|
|
del time
|
|
return np.asarray(
|
|
(
|
|
state[0] * state[0] + 3.0 * state[1],
|
|
-2.0 * state[0] + state[1],
|
|
)
|
|
)
|
|
|
|
builder = SparseSecantJacobian(
|
|
nonlinear_rhs,
|
|
csr_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=1.0e-8,
|
|
exact_rows={0: {0: 7.0, 1: 8.0}},
|
|
)
|
|
|
|
jacobian = builder(0.0, np.asarray((2.0, 1.0))).toarray()
|
|
|
|
np.testing.assert_array_equal(jacobian[0], np.asarray((7.0, 8.0)))
|
|
np.testing.assert_allclose(jacobian[1], np.asarray((-2.0, 1.0)), rtol=1.0e-7)
|
|
|
|
def test_exact_columns_match_dense_num_jac_and_use_normalized_order(self) -> None:
|
|
def nonlinear_rhs(time: float, state: np.ndarray) -> np.ndarray:
|
|
del time
|
|
return np.asarray(
|
|
(
|
|
state[0] * state[0] + state[1] * state[2],
|
|
np.sin(state[0]) + 3.0 * state[1] - state[2],
|
|
np.exp(state[2]) + state[0] * state[1],
|
|
)
|
|
)
|
|
|
|
provider_calls: list[tuple[float, np.ndarray, tuple[int, ...]]] = []
|
|
|
|
def exact_column_provider(
|
|
time: float,
|
|
state: np.ndarray,
|
|
columns: tuple[int, ...],
|
|
) -> np.ndarray:
|
|
provider_calls.append((time, state.copy(), columns))
|
|
self.assertEqual(columns, (0, 2))
|
|
return np.asarray(
|
|
(
|
|
(2.0 * state[0], state[1]),
|
|
(np.cos(state[0]), -1.0),
|
|
(state[1], np.exp(state[2])),
|
|
)
|
|
)
|
|
|
|
state = np.asarray((1.25, -0.75, 0.2))
|
|
atol = np.full(3, 1.0e-8)
|
|
structure = csc_matrix(np.ones((3, 3), dtype=bool))
|
|
base_rhs = nonlinear_rhs(0.3, state)
|
|
|
|
def vectorized_rhs(time: float, states: np.ndarray) -> np.ndarray:
|
|
states_array = np.asarray(states, dtype=float)
|
|
if states_array.ndim == 1:
|
|
return nonlinear_rhs(time, states_array)
|
|
return np.column_stack(
|
|
[
|
|
nonlinear_rhs(time, states_array[:, column])
|
|
for column in range(states_array.shape[1])
|
|
]
|
|
)
|
|
|
|
expected, _factor = num_jac(
|
|
vectorized_rhs,
|
|
0.3,
|
|
state,
|
|
base_rhs,
|
|
atol,
|
|
None,
|
|
)
|
|
builder = SparseSecantJacobian(
|
|
nonlinear_rhs,
|
|
structure,
|
|
atol=atol,
|
|
exact_columns=((2, 0), exact_column_provider),
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
actual = builder(0.3, state).toarray()
|
|
|
|
np.testing.assert_allclose(actual, expected, rtol=2.0e-7, atol=2.0e-8)
|
|
self.assertEqual(len(provider_calls), 1)
|
|
self.assertEqual(provider_calls[0][0], 0.3)
|
|
np.testing.assert_array_equal(provider_calls[0][1], state)
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["exactColumnCount"], 2)
|
|
self.assertEqual(diagnostics["finiteDifferenceColumnCount"], 1)
|
|
self.assertEqual(diagnostics["colorGroupCount"], 1)
|
|
self.assertEqual(
|
|
diagnostics["exactColumnOutsidePatternNonzeroCount"],
|
|
0,
|
|
)
|
|
|
|
def test_exact_columns_reduce_seed_zero_coloring(self) -> None:
|
|
matrix = np.asarray(
|
|
(
|
|
(1.0, 2.0, 3.0, 4.0),
|
|
(5.0, 6.0, 7.0, 8.0),
|
|
(9.0, 10.0, 11.0, 12.0),
|
|
(13.0, 14.0, 15.0, 16.0),
|
|
)
|
|
)
|
|
structure = csc_matrix(np.ones((4, 4), dtype=bool))
|
|
baseline = SparseSecantJacobian(
|
|
_SwitchableLinearRhs(matrix),
|
|
structure,
|
|
atol=1.0e-8,
|
|
max_consecutive_reuses=0,
|
|
)
|
|
reduced = SparseSecantJacobian(
|
|
_SwitchableLinearRhs(matrix),
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_columns={1: matrix[:, 1], 3: matrix[:, 3]},
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
self.assertEqual(baseline.diagnostics()["colorGroupCount"], 4)
|
|
diagnostics = reduced.diagnostics()
|
|
self.assertEqual(diagnostics["coloringSeed"], 0)
|
|
self.assertEqual(diagnostics["colorGroupCount"], 2)
|
|
self.assertEqual(diagnostics["exactColumnCount"], 2)
|
|
self.assertEqual(diagnostics["finiteDifferenceColumnCount"], 2)
|
|
|
|
def test_exact_column_nonzeros_outside_pattern_are_preserved(self) -> None:
|
|
evaluator = _SwitchableLinearRhs(np.eye(3))
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csc_matrix(np.eye(3, dtype=bool)),
|
|
atol=1.0e-8,
|
|
exact_columns={1: np.asarray((4.0, 5.0, 6.0))},
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
jacobian = builder(0.0, np.asarray((1.0, 2.0, 3.0))).toarray()
|
|
|
|
np.testing.assert_array_equal(jacobian[:, 1], np.asarray((4.0, 5.0, 6.0)))
|
|
self.assertEqual(
|
|
builder.diagnostics()["exactColumnOutsidePatternNonzeroCount"],
|
|
2,
|
|
)
|
|
|
|
def test_exact_rows_override_exact_column_intersections(self) -> None:
|
|
builder = SparseSecantJacobian(
|
|
_SwitchableLinearRhs(np.eye(3)),
|
|
csc_matrix(np.eye(3, dtype=bool)),
|
|
atol=1.0e-8,
|
|
exact_rows={0: {0: 9.0, 1: 10.0}},
|
|
exact_columns={1: np.asarray((4.0, 5.0, 6.0))},
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
jacobian = builder(0.0, np.asarray((1.0, 2.0, 3.0))).toarray()
|
|
|
|
np.testing.assert_array_equal(jacobian[0], np.asarray((9.0, 10.0, 0.0)))
|
|
np.testing.assert_array_equal(jacobian[1:, 1], np.asarray((5.0, 6.0)))
|
|
|
|
def test_all_exact_columns_skip_every_finite_difference_rhs(self) -> None:
|
|
matrix = np.asarray(((2.0, -1.0), (3.0, 4.0)))
|
|
provider_calls = 0
|
|
evaluator = _SwitchableLinearRhs(matrix)
|
|
|
|
def exact_column_provider(
|
|
time: float,
|
|
state: np.ndarray,
|
|
columns: tuple[int, ...],
|
|
) -> np.ndarray:
|
|
nonlocal provider_calls
|
|
del time, state
|
|
provider_calls += 1
|
|
self.assertEqual(columns, (0, 1))
|
|
return matrix
|
|
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csc_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=1.0e-8,
|
|
exact_columns=((1, 0), exact_column_provider),
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
jacobian = builder(0.0, np.asarray((1.0, 2.0))).toarray()
|
|
|
|
np.testing.assert_array_equal(jacobian, matrix)
|
|
self.assertEqual(provider_calls, 1)
|
|
self.assertEqual(evaluator.evaluation_count, 1)
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["colorGroupCount"], 0)
|
|
self.assertEqual(diagnostics["finiteDifferenceColumnCount"], 0)
|
|
self.assertEqual(diagnostics["baseRhsEvaluationCount"], 1)
|
|
self.assertEqual(diagnostics["finiteDifferenceRhsEvaluationCount"], 0)
|
|
|
|
def test_exact_column_provider_runs_after_base_and_before_fd(self) -> None:
|
|
matrix = np.asarray(
|
|
(
|
|
(2.0, 1.0, 0.0),
|
|
(0.0, 3.0, 4.0),
|
|
(5.0, 0.0, 6.0),
|
|
)
|
|
)
|
|
state = np.asarray((1.0, 2.0, 3.0))
|
|
events: list[tuple[str, int]] = []
|
|
evaluated_states: list[np.ndarray] = []
|
|
rhs_epoch = 0
|
|
|
|
def evaluator(time: float, evaluation_state: np.ndarray) -> np.ndarray:
|
|
nonlocal rhs_epoch
|
|
del time
|
|
rhs_epoch += 1
|
|
events.append(("rhs", rhs_epoch))
|
|
evaluated_states.append(evaluation_state.copy())
|
|
return matrix @ evaluation_state
|
|
|
|
def provider(
|
|
time: float,
|
|
provider_state: np.ndarray,
|
|
columns: tuple[int, ...],
|
|
) -> np.ndarray:
|
|
del time
|
|
events.append(("provider", rhs_epoch))
|
|
self.assertEqual(rhs_epoch, 1)
|
|
self.assertEqual(columns, (1,))
|
|
np.testing.assert_array_equal(provider_state, state)
|
|
return matrix[:, [1]]
|
|
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csc_matrix(matrix != 0.0),
|
|
atol=1.0e-8,
|
|
exact_columns=((1,), provider),
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
builder(0.0, state)
|
|
|
|
self.assertEqual(events[0], ("rhs", 1))
|
|
self.assertEqual(events[1], ("provider", 1))
|
|
self.assertTrue(any(event[0] == "rhs" for event in events[2:]))
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(
|
|
len(evaluated_states),
|
|
1 + diagnostics["colorGroupCount"],
|
|
)
|
|
self.assertEqual(
|
|
diagnostics["finiteDifferenceRhsEvaluationCount"],
|
|
diagnostics["colorGroupCount"],
|
|
)
|
|
for perturbed_state in evaluated_states[1:]:
|
|
self.assertEqual(perturbed_state[1], state[1])
|
|
self.assertTrue(np.any(perturbed_state[[0, 2]] != state[[0, 2]]))
|
|
|
|
def test_exact_column_capture_request_is_always_cancelled(self) -> None:
|
|
matrix = np.asarray(((2.0, 1.0), (0.0, 3.0)))
|
|
|
|
class Provider:
|
|
def __init__(self) -> None:
|
|
self.pending = False
|
|
self.request_count = 0
|
|
self.cancel_count = 0
|
|
self.call_count = 0
|
|
|
|
def request_primal_capture(self) -> None:
|
|
self.request_count += 1
|
|
self.pending = True
|
|
|
|
def cancel_primal_capture(self) -> None:
|
|
self.cancel_count += 1
|
|
self.pending = False
|
|
|
|
def __call__(
|
|
self,
|
|
time: float,
|
|
state: np.ndarray,
|
|
columns: tuple[int, ...],
|
|
) -> np.ndarray:
|
|
del time, state
|
|
self.call_count += 1
|
|
self.assert_capture_was_cancelled(columns)
|
|
return matrix[:, [1]]
|
|
|
|
def assert_capture_was_cancelled(
|
|
self,
|
|
columns: tuple[int, ...],
|
|
) -> None:
|
|
if self.pending or columns != (1,):
|
|
raise AssertionError("The one-shot capture request leaked.")
|
|
|
|
provider = Provider()
|
|
builder = SparseSecantJacobian(
|
|
_SwitchableLinearRhs(matrix),
|
|
csc_matrix(matrix != 0.0),
|
|
atol=1.0e-8,
|
|
exact_columns=((1,), provider),
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
jacobian = builder(0.0, np.asarray((1.0, 2.0))).toarray()
|
|
|
|
np.testing.assert_allclose(jacobian, matrix, rtol=1.0e-7)
|
|
self.assertEqual(provider.request_count, 1)
|
|
self.assertEqual(provider.cancel_count, 1)
|
|
self.assertEqual(provider.call_count, 1)
|
|
self.assertFalse(provider.pending)
|
|
|
|
def test_exact_column_provider_output_is_validated_and_errors_propagate(self) -> None:
|
|
state = np.asarray((1.0, 2.0, 3.0))
|
|
structure = csc_matrix(np.eye(3, dtype=bool))
|
|
|
|
with self.assertRaisesRegex(ValueError, "duplicated"):
|
|
SparseSecantJacobian(
|
|
_SwitchableLinearRhs(np.eye(3)),
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_columns=((1, 1), lambda *_args: np.zeros((3, 2))),
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "out of range"):
|
|
SparseSecantJacobian(
|
|
_SwitchableLinearRhs(np.eye(3)),
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_columns=((3,), lambda *_args: np.zeros((3, 1))),
|
|
)
|
|
|
|
invalid_builders = (
|
|
(
|
|
SparseSecantJacobian(
|
|
_SwitchableLinearRhs(np.eye(3)),
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_columns=((0, 2), lambda *_args: np.zeros((2, 3))),
|
|
max_consecutive_reuses=0,
|
|
),
|
|
"must return shape",
|
|
),
|
|
(
|
|
SparseSecantJacobian(
|
|
_SwitchableLinearRhs(np.eye(3)),
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_columns=(
|
|
(0, 2),
|
|
lambda *_args: np.asarray(
|
|
((1.0, 0.0), (0.0, np.nan), (0.0, 1.0))
|
|
),
|
|
),
|
|
max_consecutive_reuses=0,
|
|
),
|
|
"finite values",
|
|
),
|
|
(
|
|
SparseSecantJacobian(
|
|
_SwitchableLinearRhs(np.eye(3)),
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_columns=(
|
|
(0, 2),
|
|
lambda *_args: {
|
|
2: np.zeros(3),
|
|
0: np.zeros(3),
|
|
},
|
|
),
|
|
max_consecutive_reuses=0,
|
|
),
|
|
"normalized column order",
|
|
),
|
|
)
|
|
for builder, message in invalid_builders:
|
|
with self.subTest(message=message):
|
|
with self.assertRaisesRegex(ValueError, message):
|
|
builder(0.0, state)
|
|
|
|
class ProviderFailure(RuntimeError):
|
|
pass
|
|
|
|
def failing_provider(*_args):
|
|
raise ProviderFailure("provider failed")
|
|
|
|
failure_evaluator = _SwitchableLinearRhs(np.eye(3))
|
|
failing_builder = SparseSecantJacobian(
|
|
failure_evaluator,
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_columns=((0,), failing_provider),
|
|
max_consecutive_reuses=0,
|
|
)
|
|
with self.assertRaisesRegex(ProviderFailure, "provider failed"):
|
|
failing_builder(0.0, state)
|
|
self.assertEqual(failure_evaluator.evaluation_count, 1)
|
|
failure_diagnostics = failing_builder.diagnostics()
|
|
self.assertEqual(failure_diagnostics["baseRhsEvaluationCount"], 1)
|
|
self.assertEqual(
|
|
failure_diagnostics["finiteDifferenceRhsEvaluationCount"],
|
|
0,
|
|
)
|
|
|
|
def test_exact_column_provider_is_refreshed_after_segment_reset(self) -> None:
|
|
calls: list[tuple[float, np.ndarray, tuple[int, ...]]] = []
|
|
|
|
def nonlinear_rhs(time: float, state: np.ndarray) -> np.ndarray:
|
|
return np.asarray(
|
|
(
|
|
2.0 * state[0] + (1.0 + time) * state[1],
|
|
state[0] * state[1],
|
|
)
|
|
)
|
|
|
|
def provider(
|
|
time: float,
|
|
state: np.ndarray,
|
|
columns: tuple[int, ...],
|
|
) -> np.ndarray:
|
|
calls.append((time, state.copy(), columns))
|
|
return np.asarray(((1.0 + time,), (state[0],)))
|
|
|
|
builder = SparseSecantJacobian(
|
|
nonlinear_rhs,
|
|
csc_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=1.0e-8,
|
|
exact_columns=((1,), provider),
|
|
max_consecutive_reuses=0,
|
|
)
|
|
first_state = np.asarray((1.0, 2.0))
|
|
second_state = np.asarray((3.0, 4.0))
|
|
|
|
first = builder(0.0, first_state).toarray()
|
|
builder.start_segment()
|
|
second = builder(1.0, second_state).toarray()
|
|
|
|
np.testing.assert_allclose(first[:, 1], np.asarray((1.0, 1.0)))
|
|
np.testing.assert_allclose(second[:, 1], np.asarray((2.0, 3.0)))
|
|
self.assertEqual([call[0] for call in calls], [0.0, 1.0])
|
|
self.assertEqual([call[2] for call in calls], [(1,), (1,)])
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["fullBuildCount"], 2)
|
|
self.assertEqual(diagnostics["segmentStartCount"], 1)
|
|
self.assertEqual(len(diagnostics["segments"]), 2)
|
|
|
|
def test_exact_column_unavailable_falls_back_and_next_build_recovers(self) -> None:
|
|
matrix = np.asarray(
|
|
(
|
|
(2.0, 1.0, 3.0, -1.0),
|
|
(4.0, 5.0, -2.0, 6.0),
|
|
(7.0, -3.0, 8.0, 2.0),
|
|
(-4.0, 9.0, 1.0, 10.0),
|
|
)
|
|
)
|
|
state = np.asarray((1.0, -2.0, 0.5, 3.0))
|
|
evaluated_states: list[np.ndarray] = []
|
|
provider_call_count = 0
|
|
|
|
def evaluator(time: float, evaluation_state: np.ndarray) -> np.ndarray:
|
|
del time
|
|
evaluated_states.append(evaluation_state.copy())
|
|
return matrix @ evaluation_state
|
|
|
|
def provider(
|
|
time: float,
|
|
provider_state: np.ndarray,
|
|
columns: tuple[int, ...],
|
|
) -> np.ndarray:
|
|
nonlocal provider_call_count
|
|
del time, provider_state
|
|
provider_call_count += 1
|
|
if provider_call_count == 1:
|
|
raise ExactColumnsUnavailable("tangent domain boundary")
|
|
return matrix[:, columns]
|
|
|
|
structure = csc_matrix(np.ones((4, 4), dtype=bool))
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_columns=((3, 1), provider),
|
|
max_consecutive_reuses=0,
|
|
)
|
|
baseline = SparseSecantJacobian(
|
|
_SwitchableLinearRhs(matrix),
|
|
structure,
|
|
atol=1.0e-8,
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
fallback = builder(0.0, state).toarray()
|
|
expected_fallback = baseline(0.0, state).toarray()
|
|
|
|
np.testing.assert_array_equal(fallback, expected_fallback)
|
|
self.assertEqual(provider_call_count, 1)
|
|
self.assertTrue(
|
|
any(
|
|
evaluation_state[1] != state[1]
|
|
for evaluation_state in evaluated_states[1:]
|
|
)
|
|
)
|
|
self.assertTrue(
|
|
any(
|
|
evaluation_state[3] != state[3]
|
|
for evaluation_state in evaluated_states[1:]
|
|
)
|
|
)
|
|
fallback_diagnostics = builder.diagnostics()
|
|
self.assertEqual(fallback_diagnostics["baseRhsEvaluationCount"], 1)
|
|
self.assertEqual(fallback_diagnostics["originalColorGroupCount"], 4)
|
|
self.assertEqual(fallback_diagnostics["remainingColorGroupCount"], 2)
|
|
self.assertEqual(fallback_diagnostics["exactColumnBuildCount"], 0)
|
|
self.assertEqual(fallback_diagnostics["exactColumnFallbackCount"], 1)
|
|
self.assertEqual(
|
|
fallback_diagnostics["lastExactColumnFallbackReason"],
|
|
"tangent domain boundary",
|
|
)
|
|
self.assertEqual(
|
|
fallback_diagnostics[
|
|
"lastBuildFiniteDifferenceRhsEvaluationCount"
|
|
],
|
|
4,
|
|
)
|
|
self.assertIsNotNone(builder._original_factor)
|
|
self.assertIsNone(builder._remaining_factor)
|
|
|
|
recovery_start = len(evaluated_states)
|
|
recovered = builder(0.1, state).toarray()
|
|
recovery_states = evaluated_states[recovery_start:]
|
|
|
|
np.testing.assert_allclose(recovered, matrix, rtol=1.0e-7)
|
|
self.assertEqual(provider_call_count, 2)
|
|
self.assertEqual(len(recovery_states), 3)
|
|
for perturbed_state in recovery_states[1:]:
|
|
self.assertEqual(perturbed_state[1], state[1])
|
|
self.assertEqual(perturbed_state[3], state[3])
|
|
recovery_diagnostics = builder.diagnostics()
|
|
self.assertEqual(recovery_diagnostics["exactColumnBuildCount"], 1)
|
|
self.assertEqual(recovery_diagnostics["exactColumnFallbackCount"], 1)
|
|
self.assertEqual(recovery_diagnostics["baseRhsEvaluationCount"], 2)
|
|
self.assertEqual(
|
|
recovery_diagnostics[
|
|
"lastBuildFiniteDifferenceRhsEvaluationCount"
|
|
],
|
|
2,
|
|
)
|
|
self.assertIsNotNone(builder._remaining_factor)
|
|
self.assertIsNot(builder._original_factor, builder._remaining_factor)
|
|
|
|
def test_exact_column_subset_retry_matches_scipy_active_columns(self) -> None:
|
|
offset = np.asarray((1.0e8, -3.0e8, 2.0e8))
|
|
slopes = np.asarray((1.0, 2.0, -3.0))
|
|
state = np.asarray((1.0, -2.0, 0.5))
|
|
evaluated_states: list[np.ndarray] = []
|
|
|
|
def evaluator(time: float, evaluation_state: np.ndarray) -> np.ndarray:
|
|
del time
|
|
evaluated_states.append(evaluation_state.copy())
|
|
return offset + slopes * evaluation_state
|
|
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csc_matrix(np.eye(3, dtype=bool)),
|
|
atol=1.0e-8,
|
|
exact_columns={1: np.asarray((0.0, slopes[1], 0.0))},
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
actual = builder(0.0, state).toarray()
|
|
|
|
active = np.asarray((0, 2))
|
|
active_state = state[active]
|
|
active_offset = offset[active]
|
|
active_slopes = slopes[active]
|
|
|
|
def active_rhs(time: float, states: np.ndarray) -> np.ndarray:
|
|
del time
|
|
states_array = np.asarray(states, dtype=float)
|
|
if states_array.ndim == 1:
|
|
return active_offset + active_slopes * states_array
|
|
return active_offset[:, None] + active_slopes[:, None] * states_array
|
|
|
|
active_structure = csc_matrix(np.eye(2, dtype=bool))
|
|
active_groups = group_columns(active_structure, order=0)
|
|
expected, _factor = num_jac(
|
|
active_rhs,
|
|
0.0,
|
|
active_state,
|
|
active_rhs(0.0, active_state),
|
|
np.full(2, 1.0e-8),
|
|
None,
|
|
(active_structure, active_groups),
|
|
)
|
|
|
|
np.testing.assert_array_equal(
|
|
actual[np.ix_(active, active)],
|
|
expected.toarray(),
|
|
)
|
|
diagnostics = builder.diagnostics()
|
|
self.assertGreater(
|
|
diagnostics["lastBuildFiniteDifferenceRhsEvaluationCount"],
|
|
diagnostics["remainingColorGroupCount"],
|
|
)
|
|
for perturbed_state in evaluated_states[1:]:
|
|
self.assertEqual(perturbed_state[1], state[1])
|
|
|
|
def test_audit_step_respects_tiny_state_absolute_tolerance(self) -> None:
|
|
evaluator = _SwitchableLinearRhs(np.eye(2))
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csr_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=np.asarray((1.0e-12, 1.0e-8)),
|
|
)
|
|
state = np.zeros(2)
|
|
|
|
step = builder._audit_step(state)
|
|
|
|
self.assertGreater(abs(step[0]), 0.0)
|
|
self.assertLessEqual(abs(step[0]), 1.0e-12)
|
|
self.assertGreater(abs(step[1]), 0.0)
|
|
self.assertLessEqual(abs(step[1]), 1.0e-8)
|
|
|
|
def test_optimized_finite_difference_skips_secant_observation_work(self) -> None:
|
|
matrix = np.asarray(((2.0, -1.0), (0.5, 4.0)))
|
|
evaluator = _SwitchableLinearRhs(matrix)
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csr_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=1.0e-8,
|
|
max_consecutive_reuses=0,
|
|
)
|
|
state = np.asarray((1.0, -2.0))
|
|
derivative = evaluator(0.0, state)
|
|
builder.observe(0.0, state, derivative)
|
|
|
|
jacobian = builder(0.0, state)
|
|
|
|
np.testing.assert_allclose(jacobian.toarray(), matrix, rtol=1.0e-7)
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["mode"], "optimizedSparseFiniteDifference")
|
|
self.assertEqual(diagnostics["fullBuildCount"], 1)
|
|
self.assertEqual(diagnostics["secantReuseCount"], 0)
|
|
self.assertEqual(diagnostics["baseRhsEvaluationCount"], 1)
|
|
self.assertEqual(diagnostics["secantUpdateCount"], 0)
|
|
self.assertEqual(diagnostics["coloringSeed"], 0)
|
|
|
|
def test_seed_zero_callable_matches_scipy_num_jac(self) -> None:
|
|
def nonlinear_rhs(time: float, state: np.ndarray) -> np.ndarray:
|
|
del time
|
|
return np.asarray(
|
|
(
|
|
state[0] * state[0] + 3.0 * state[1],
|
|
np.sin(state[0]) - state[1],
|
|
)
|
|
)
|
|
|
|
state = np.asarray((2.0, 1.0))
|
|
atol = np.full(2, 1.0e-8)
|
|
structure = csc_matrix(np.ones((2, 2), dtype=bool))
|
|
groups = group_columns(structure, order=0)
|
|
base_rhs = nonlinear_rhs(0.0, state)
|
|
|
|
def vectorized_rhs(time: float, states: np.ndarray) -> np.ndarray:
|
|
states_array = np.asarray(states, dtype=float)
|
|
if states_array.ndim == 1:
|
|
return nonlinear_rhs(time, states_array)
|
|
return np.column_stack(
|
|
[
|
|
nonlinear_rhs(time, states_array[:, column])
|
|
for column in range(states_array.shape[1])
|
|
]
|
|
)
|
|
|
|
expected, _factor = num_jac(
|
|
vectorized_rhs,
|
|
0.0,
|
|
state,
|
|
base_rhs,
|
|
atol,
|
|
None,
|
|
(structure, groups),
|
|
)
|
|
builder = SparseSecantJacobian(
|
|
nonlinear_rhs,
|
|
structure,
|
|
atol=atol,
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
actual = builder(0.0, state)
|
|
|
|
np.testing.assert_array_equal(actual.toarray(), expected.toarray())
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["coloringSeed"], 0)
|
|
self.assertEqual(diagnostics["baseRhsEvaluationCount"], 1)
|
|
|
|
def test_default_coloring_stays_seed_zero_when_another_seed_is_smaller(self) -> None:
|
|
structure = csc_matrix(
|
|
np.asarray(
|
|
(
|
|
(1, 1, 0, 0, 0),
|
|
(0, 1, 0, 0, 0),
|
|
(0, 0, 1, 0, 0),
|
|
(0, 0, 0, 1, 1),
|
|
(0, 1, 0, 0, 1),
|
|
),
|
|
dtype=bool,
|
|
)
|
|
)
|
|
seed_zero_count = int(group_columns(structure, order=0).max()) + 1
|
|
seed_54_count = int(group_columns(structure, order=54).max()) + 1
|
|
self.assertEqual(seed_zero_count, 3)
|
|
self.assertEqual(seed_54_count, 2)
|
|
|
|
builder = SparseSecantJacobian(
|
|
_SwitchableLinearRhs(structure.toarray().astype(float)),
|
|
structure,
|
|
atol=1.0e-8,
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["coloringSeed"], 0)
|
|
self.assertEqual(diagnostics["colorGroupCount"], 3)
|
|
self.assertEqual(diagnostics["defaultColorGroupCount"], 3)
|
|
|
|
def test_exact_rows_do_not_reorder_seed_zero_perturbation_batches(self) -> None:
|
|
structure = csc_matrix(np.asarray(((1, 1), (1, 0)), dtype=bool))
|
|
builder = SparseSecantJacobian(
|
|
_SwitchableLinearRhs(np.asarray(((1.0, 1.0), (1.0, 0.0)))),
|
|
structure,
|
|
atol=1.0e-8,
|
|
exact_rows={0: {0: 1.0, 1: 1.0}},
|
|
max_consecutive_reuses=0,
|
|
)
|
|
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["coloringSeed"], 0)
|
|
self.assertEqual(diagnostics["colorGroupCount"], 2)
|
|
|
|
def test_rejects_more_than_one_consecutive_secant_reuse(self) -> None:
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
"max_consecutive_reuses must be either 0 or 1",
|
|
):
|
|
SparseSecantJacobian(
|
|
_SwitchableLinearRhs(np.eye(2)),
|
|
csr_matrix(np.eye(2, dtype=bool)),
|
|
atol=1.0e-8,
|
|
max_consecutive_reuses=2,
|
|
)
|
|
|
|
def test_uninformative_zero_jv_audit_forces_full_refresh(self) -> None:
|
|
evaluator = _SwitchableLinearRhs(np.zeros((2, 2)))
|
|
builder = SparseSecantJacobian(
|
|
evaluator,
|
|
csr_matrix(np.ones((2, 2), dtype=bool)),
|
|
atol=1.0e-8,
|
|
)
|
|
state = np.asarray((1.0, 1.0))
|
|
builder(0.0, state)
|
|
builder.observe(0.1, state, np.zeros(2))
|
|
builder.observe(0.1, state + 1.0, np.zeros(2))
|
|
|
|
refreshed = builder(0.1, state + 1.0)
|
|
|
|
np.testing.assert_array_equal(refreshed.toarray(), np.zeros((2, 2)))
|
|
diagnostics = builder.diagnostics()
|
|
self.assertEqual(diagnostics["secantReuseCount"], 0)
|
|
self.assertEqual(diagnostics["auditFailureCount"], 1)
|
|
self.assertEqual(diagnostics["fullBuildCount"], 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|