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()