solve-nivp 0.1.3.dev0__tar.gz → 0.2.0.dev0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (78) hide show
  1. {solve_nivp-0.1.3.dev0/solve_nivp.egg-info → solve_nivp-0.2.0.dev0}/PKG-INFO +1 -1
  2. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/pyproject.toml +1 -1
  3. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/setup.py +1 -1
  4. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/__init__.py +42 -4
  5. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/adaptive_integrator.py +4 -0
  6. solve_nivp-0.2.0.dev0/solve_nivp/block_system.py +408 -0
  7. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/contact.py +2 -1
  8. solve_nivp-0.2.0.dev0/solve_nivp/desaxce_contact.py +1874 -0
  9. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/integrations.py +997 -50
  10. solve_nivp-0.2.0.dev0/solve_nivp/macklin_contact.py +276 -0
  11. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/ncp_contact.py +413 -13
  12. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/nonlinear_solvers.py +20 -4
  13. solve_nivp-0.2.0.dev0/solve_nivp/pcr.py +127 -0
  14. solve_nivp-0.2.0.dev0/solve_nivp/projected_radau_contact.py +1697 -0
  15. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/projections.py +11 -5
  16. solve_nivp-0.2.0.dev0/solve_nivp/rattle_contact.py +1750 -0
  17. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0/solve_nivp.egg-info}/PKG-INFO +1 -1
  18. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp.egg-info/SOURCES.txt +17 -0
  19. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_3d_and_anisotropic_contact.py +0 -113
  20. solve_nivp-0.2.0.dev0/tests/test_block_system.py +90 -0
  21. solve_nivp-0.2.0.dev0/tests/test_bouncing_ball_schur_comparison.py +150 -0
  22. solve_nivp-0.2.0.dev0/tests/test_desaxce_contact.py +453 -0
  23. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_dilatancy.py +0 -70
  24. solve_nivp-0.2.0.dev0/tests/test_general_moreau_projection.py +74 -0
  25. solve_nivp-0.2.0.dev0/tests/test_macklin_contact.py +49 -0
  26. solve_nivp-0.2.0.dev0/tests/test_ncp_schur_integration.py +275 -0
  27. solve_nivp-0.2.0.dev0/tests/test_pcr.py +92 -0
  28. solve_nivp-0.2.0.dev0/tests/test_projected_radau_contact.py +732 -0
  29. solve_nivp-0.2.0.dev0/tests/test_radau_iia.py +502 -0
  30. solve_nivp-0.2.0.dev0/tests/test_rattle_integrator.py +551 -0
  31. solve_nivp-0.2.0.dev0/tests/test_rattle_local_slider.py +248 -0
  32. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/LICENSE +0 -0
  33. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/README.md +0 -0
  34. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/setup.cfg +0 -0
  35. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/ODESolver.py +0 -0
  36. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/ODESystem.py +0 -0
  37. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/_numba_accel.py +0 -0
  38. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/_selftest.py +0 -0
  39. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/alart_curnier_contact.py +0 -0
  40. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/rl/__init__.py +0 -0
  41. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/rl/callbacks.py +0 -0
  42. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/rl/dependency.py +0 -0
  43. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp/rl/env.py +0 -0
  44. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp.egg-info/dependency_links.txt +0 -0
  45. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp.egg-info/entry_points.txt +0 -0
  46. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp.egg-info/requires.txt +0 -0
  47. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/solve_nivp.egg-info/top_level.txt +0 -0
  48. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/__init__.py +0 -0
  49. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_active_set_filter.py +0 -0
  50. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_alart_curnier_contact.py +0 -0
  51. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_algebraic_constraint_projection.py +0 -0
  52. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_anisotropic_soc_projection.py +0 -0
  53. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_auto_h0_and_dae_weighting.py +0 -0
  54. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_build_impulse_contact.py +0 -0
  55. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_c_extract_contact.py +0 -0
  56. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_composite_contact_projection.py +0 -0
  57. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_coulomb_projection.py +0 -0
  58. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_coulomb_projection_jacobian.py +0 -0
  59. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_error_predictive_rejection.py +0 -0
  60. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_globalization.py +0 -0
  61. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_identity_newton_linear_solver.py +0 -0
  62. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_import_and_api.py +0 -0
  63. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_integrators_added.py +0 -0
  64. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_jacobian_scaling.py +0 -0
  65. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_large_scale_solver_fixes.py +0 -0
  66. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_mu_scaled_soc_projection.py +0 -0
  67. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_ncp_contact.py +0 -0
  68. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_nl_recovery_cap.py +0 -0
  69. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_nonlinear_solvers_added.py +0 -0
  70. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_per_dof_tolerances.py +0 -0
  71. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_prestress_soc.py +0 -0
  72. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_prestressed_fault_dynamic_helper.py +0 -0
  73. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_projection_batch_equivalence.py +0 -0
  74. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_projections_added.py +0 -0
  75. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_sdirk2.py +0 -0
  76. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_sdirk2_soc_contact.py +0 -0
  77. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_sparse_semismooth_newton.py +0 -0
  78. {solve_nivp-0.1.3.dev0 → solve_nivp-0.2.0.dev0}/tests/test_threading_time_and_fk.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: solve_nivp
3
- Version: 0.1.3.dev0
3
+ Version: 0.2.0.dev0
4
4
  Summary: A Python toolkit for integrating nonsmooth dynamical systems
5
5
  Author: David Riley, Ioannis Stefanou
6
6
  License: MIT
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "solve_nivp"
7
- version = "0.1.3.dev0"
7
+ version = "0.2.0.dev0"
8
8
  description = "A Python toolkit for integrating nonsmooth dynamical systems"
9
9
  readme = "README.md"
10
10
  license = {text = "MIT"}
@@ -50,7 +50,7 @@ class BuildSphinx(Command):
50
50
 
51
51
  setup(
52
52
  name="solve_nivp",
53
- version="0.1.3.dev0",
53
+ version="0.2.0.dev0",
54
54
  packages=find_packages(), # automatically discovers packages
55
55
  description="A solver package for implicit ODEs and projection-based solvers",
56
56
  cmdclass={'build_sphinx': BuildSphinx},
@@ -62,7 +62,8 @@ from .projections import (
62
62
  CompositeContactProjection,
63
63
  )
64
64
  from .nonlinear_solvers import ImplicitEquationSolver, UMFPACK_AVAILABLE, PETSC_AVAILABLE
65
- from .integrations import BackwardEuler, Trapezoidal, ThetaMethod, CompositeMethod, EmbeddedBETR, SDIRK2 # , BDFMethod
65
+ from .integrations import BackwardEuler, BackwardEulerSchur, RadauIIASchur, Trapezoidal, ThetaMethod, CompositeMethod, EmbeddedBETR, SDIRK2, RadauIIA # , BDFMethod
66
+ from .block_system import SchurComplementSolver, BlockStructuredSystem
66
67
  from .ODESystem import ODESystem
67
68
  from .ODESolver import ODESolver
68
69
  from .contact import build_impulse_contact, ContactSystem
@@ -73,6 +74,24 @@ from .alart_curnier_contact import (
73
74
  from .ncp_contact import (
74
75
  build_ncp_contact,
75
76
  build_dynamic_ncp_contact,
77
+ build_ncp_contact_blocked,
78
+ )
79
+ from .desaxce_contact import (
80
+ build_dynamic_desaxce_contact,
81
+ build_dynamic_desaxce_projected_contact,
82
+ build_dynamic_desaxce_residual_contact,
83
+ )
84
+ from .rattle_contact import (
85
+ build_dynamic_rattle_contact,
86
+ solve_dynamic_rattle_contact,
87
+ build_rattle_system,
88
+ RattleMechanicalSystem,
89
+ RattleContactSpec,
90
+ RattleBilateralSpec,
91
+ RattleAlgebraicSpec,
92
+ RattleContactSystem,
93
+ RattleSolveResult,
94
+ RattleSolver,
76
95
  )
77
96
 
78
97
  # Curated public API
@@ -83,7 +102,7 @@ __all__ = [
83
102
  # Nonlinear solver
84
103
  'ImplicitEquationSolver',
85
104
  # Integrators
86
- 'BackwardEuler', 'Trapezoidal', 'ThetaMethod', 'CompositeMethod', 'EmbeddedBETR', 'SDIRK2',
105
+ 'BackwardEuler', 'BackwardEulerSchur', 'RadauIIASchur', 'Trapezoidal', 'ThetaMethod', 'CompositeMethod', 'EmbeddedBETR', 'SDIRK2', 'RadauIIA',
87
106
  # Projections
88
107
  'Projection',
89
108
  'CoulombProjection', 'SignProjection', 'IdentityProjection',
@@ -93,7 +112,22 @@ __all__ = [
93
112
  # Contact helpers
94
113
  'build_impulse_contact', 'build_alart_curnier_contact',
95
114
  'build_dynamic_alart_curnier_contact', 'build_ncp_contact',
96
- 'build_dynamic_ncp_contact', 'ContactSystem',
115
+ 'build_dynamic_ncp_contact', 'build_ncp_contact_blocked',
116
+ 'SchurComplementSolver', 'BlockStructuredSystem',
117
+ 'build_dynamic_desaxce_contact',
118
+ 'build_dynamic_desaxce_projected_contact',
119
+ 'build_dynamic_desaxce_residual_contact',
120
+ 'build_dynamic_rattle_contact',
121
+ 'solve_dynamic_rattle_contact',
122
+ 'build_rattle_system',
123
+ 'ContactSystem',
124
+ 'RattleMechanicalSystem',
125
+ 'RattleContactSpec',
126
+ 'RattleBilateralSpec',
127
+ 'RattleAlgebraicSpec',
128
+ 'RattleContactSystem',
129
+ 'RattleSolveResult',
130
+ 'RattleSolver',
97
131
  ]
98
132
 
99
133
 
@@ -140,7 +174,9 @@ def solve_nivp(
140
174
  Initial state.
141
175
  method : str, default 'composite'
142
176
  Time stepping scheme: ``'backward_euler'``, ``'trapezoidal'``, ``'theta'``,
143
- ``'composite'`` (TR-BE like second order), ``'embedded_betr'``, ``'sdirk2'``.
177
+ ``'composite'`` (TR-BE like second order), ``'embedded_betr'``, ``'sdirk2'``,
178
+ ``'radau_iia'`` (L-stable, stiffly accurate, order 3 or 5; stages controlled
179
+ via ``integrator_opts={'stages': 2}`` or ``{'stages': 3}``).
144
180
  projection : str or Projection or None, default 'identity'
145
181
  Name of projection to build: ``'coulomb'``, ``'sign'``, ``'identity'`` or
146
182
  ``None``. At the high-level API, ``None`` is promoted to the identity
@@ -348,6 +384,8 @@ def solve_nivp(
348
384
  integrator = EmbeddedBETR(solver=solver_instance, A=A)
349
385
  elif m == 'sdirk2':
350
386
  integrator = SDIRK2(solver=solver_instance, A=A, **_integrator_opts)
387
+ elif m in ('radau_iia', 'radau'):
388
+ integrator = RadauIIA(solver=solver_instance, A=A, **_integrator_opts)
351
389
  # elif m == 'bdf':
352
390
  # integrator = BDFMethod(solver=solver_instance, atol=atol, rtol=rtol)
353
391
  else:
@@ -416,6 +416,10 @@ class AdaptiveStepping:
416
416
  # ------------------------------------------------------------------
417
417
 
418
418
  def _infer_method_order(self, integrator) -> int:
419
+ p_emb = getattr(integrator, 'embedded_order', None)
420
+ if isinstance(p_emb, (int, float)) and p_emb > 0:
421
+ return int(p_emb)
422
+
419
423
  p = getattr(integrator, 'order', None)
420
424
  if isinstance(p, (int, float)) and p > 0:
421
425
  return int(p)
@@ -0,0 +1,408 @@
1
+ """Block-structured Newton solver using Schur-complement reduction.
2
+
3
+ Generalises the compliance-formulation Newton method from Macklin et al.
4
+ (2019) §6 to non-symmetric off-diagonal blocks. The 2×2 block system
5
+
6
+ [H, B_top] [Δu] [g]
7
+ [B_bot, C ] [Δλ] = [h_c]
8
+
9
+ is reduced via Schur complement
10
+
11
+ S = C - B_bot H⁻¹ B_top
12
+ S Δλ = h_c - B_bot H⁻¹ g
13
+ Δu = H⁻¹ (g - B_top Δλ)
14
+
15
+ When B_top = -J^T and B_bot = J (Macklin velocity-level constraints)
16
+ this recovers S = J H⁻¹ J^T + C. For position-level NCP constraints
17
+ the off-diagonal blocks differ and must be kept separate.
18
+ """
19
+
20
+ from __future__ import annotations
21
+
22
+ from typing import Any, Dict, Optional, Protocol, Tuple, runtime_checkable
23
+
24
+ import numpy as np
25
+ import scipy.linalg as la
26
+ import scipy.sparse as sp
27
+ import scipy.sparse.linalg as spla
28
+
29
+ from .pcr import pcr_solve
30
+
31
+
32
+ @runtime_checkable
33
+ class BlockStructuredSystem(Protocol):
34
+ """Protocol for systems that expose 4-block Newton Jacobian."""
35
+
36
+ n_phys: int
37
+ n_react: int
38
+
39
+ def assemble_blocks(
40
+ self,
41
+ y: np.ndarray,
42
+ t: float,
43
+ h: float,
44
+ y_prev: np.ndarray,
45
+ ) -> Dict[str, Any]:
46
+ """Return the block components for the Newton system.
47
+
48
+ Returns
49
+ -------
50
+ dict with keys:
51
+ H : ndarray or sparse, (n_phys, n_phys)
52
+ B_top : ndarray or sparse, (n_phys, n_react)
53
+ B_bot : ndarray or sparse, (n_react, n_phys)
54
+ C : ndarray or sparse, (n_react, n_react)
55
+ g : ndarray, (n_phys,)
56
+ h_c : ndarray, (n_react,)
57
+ precond_diag : ndarray, (n_react,) or None
58
+ """
59
+ ...
60
+
61
+
62
+ class SchurComplementSolver:
63
+ """Newton solver using Schur-complement reduction.
64
+
65
+ Parameters
66
+ ----------
67
+ maxiter : int
68
+ Maximum Newton iterations.
69
+ tol : float
70
+ Convergence tolerance on the full residual norm.
71
+ pcr_maxiter : int
72
+ Maximum PCR iterations for the Schur system.
73
+ pcr_tol : float
74
+ PCR convergence tolerance.
75
+ damped_step_fraction : float
76
+ Macklin §8.1 damped Newton step alpha in (0, 1].
77
+ diagonal_regularization : float
78
+ Macklin §8.4 epsilon added to Schur diagonal.
79
+ use_preconditioner : bool
80
+ Apply the Macklin complementarity preconditioner.
81
+ linear_solver : str
82
+ ``"direct"`` assembles the full saddle-point matrix and uses
83
+ dense LU / sparse SPLU. ``"pcr"`` uses the Schur complement
84
+ with PCR.
85
+ """
86
+
87
+ def __init__(
88
+ self,
89
+ maxiter: int = 20,
90
+ tol: float = 1e-10,
91
+ pcr_maxiter: int = 100,
92
+ pcr_tol: float = 1e-10,
93
+ damped_step_fraction: float = 0.75,
94
+ diagonal_regularization: float = 0.0,
95
+ use_preconditioner: bool = True,
96
+ linear_solver: str = "direct",
97
+ ):
98
+ self.maxiter = int(maxiter)
99
+ self.tol = float(tol)
100
+ self.pcr_maxiter = int(pcr_maxiter)
101
+ self.pcr_tol = float(pcr_tol)
102
+ self.damped_step_fraction = float(damped_step_fraction)
103
+ self.diagonal_regularization = float(diagonal_regularization)
104
+ self.use_preconditioner = bool(use_preconditioner)
105
+ self.linear_solver = str(linear_solver).strip().lower()
106
+
107
+ def _to_dense(self, M):
108
+ if sp.issparse(M):
109
+ return M.toarray()
110
+ return np.asarray(M, dtype=float)
111
+
112
+ def _compute_H_inv_action(self, H):
113
+ """Return a callable v -> H^{-1} v."""
114
+ n = H.shape[0]
115
+ if sp.issparse(H) and n > 200:
116
+ H_lu = spla.splu(H.tocsc())
117
+ return lambda v: H_lu.solve(v)
118
+ else:
119
+ H_dense = self._to_dense(H)
120
+ H_lu_piv = la.lu_factor(H_dense)
121
+ return lambda v, _lu=H_lu_piv: la.lu_solve(_lu, v)
122
+
123
+ def _solve_linear_direct(self, H, B_top, B_bot, C, g, h_c):
124
+ """Solve the full saddle-point system with a direct factorisation."""
125
+ n_p = g.shape[0]
126
+ H_arr = self._to_dense(H)
127
+ Bt = self._to_dense(B_top)
128
+ Bb = self._to_dense(B_bot)
129
+ C_arr = self._to_dense(C)
130
+
131
+ A_full = np.block([[H_arr, Bt],
132
+ [Bb, C_arr]])
133
+ rhs_full = np.concatenate([g, h_c])
134
+
135
+ x = la.solve(A_full, rhs_full)
136
+ residual = float(np.linalg.norm(A_full @ x - rhs_full))
137
+ return x[:n_p], x[n_p:], {
138
+ "converged": True, "iterations": 1,
139
+ "residual_norm": residual, "residual_history": [residual],
140
+ }
141
+
142
+ def _solve_linear_pcr(self, H, B_top, B_bot, C, g, h_c, precond_diag):
143
+ """Solve via Schur complement S = C - B_bot H^{-1} B_top with PCR."""
144
+ Bt = self._to_dense(B_top)
145
+ Bb = self._to_dense(B_bot)
146
+ C_arr = self._to_dense(C)
147
+
148
+ H_inv = self._compute_H_inv_action(H)
149
+ schur_rhs = h_c - Bb @ H_inv(g)
150
+
151
+ def schur_matvec(v):
152
+ return C_arr @ v - Bb @ H_inv(Bt @ v)
153
+
154
+ precond = None
155
+ if self.use_preconditioner and precond_diag is not None:
156
+ safe = np.where(np.abs(precond_diag) > 1e-30, precond_diag, 1.0)
157
+ precond = lambda v: v / safe
158
+
159
+ eps = self.diagonal_regularization
160
+ if eps > 0.0:
161
+ _base = schur_matvec
162
+ def schur_matvec(v, _m=_base, _e=eps):
163
+ return _m(v) + _e * v
164
+
165
+ delta_lam, pcr_info = pcr_solve(
166
+ schur_matvec, schur_rhs,
167
+ maxiter=self.pcr_maxiter, tol=self.pcr_tol,
168
+ preconditioner=precond,
169
+ )
170
+ delta_u = H_inv(g - Bt @ delta_lam)
171
+ return delta_u, delta_lam, pcr_info
172
+
173
+ def solve_linear(
174
+ self,
175
+ H: np.ndarray,
176
+ B_top: np.ndarray,
177
+ B_bot: np.ndarray,
178
+ C: np.ndarray,
179
+ g: np.ndarray,
180
+ h_c: np.ndarray,
181
+ precond_diag: Optional[np.ndarray] = None,
182
+ ) -> Tuple[np.ndarray, np.ndarray, dict]:
183
+ """Solve one linearized saddle-point system.
184
+
185
+ Parameters
186
+ ----------
187
+ H : (n_phys, n_phys)
188
+ B_top : (n_phys, n_react)
189
+ B_bot : (n_react, n_phys)
190
+ C : (n_react, n_react)
191
+ g : (n_phys,)
192
+ h_c : (n_react,)
193
+ precond_diag : (n_react,) or None
194
+
195
+ Returns
196
+ -------
197
+ delta_u, delta_lam, info
198
+ """
199
+ g = np.asarray(g, dtype=float).ravel()
200
+ h_c = np.asarray(h_c, dtype=float).ravel()
201
+
202
+ if self.linear_solver == "direct":
203
+ return self._solve_linear_direct(H, B_top, B_bot, C, g, h_c)
204
+ else:
205
+ return self._solve_linear_pcr(
206
+ H, B_top, B_bot, C, g, h_c, precond_diag,
207
+ )
208
+
209
+ def solve(
210
+ self,
211
+ block_system: BlockStructuredSystem,
212
+ y0: np.ndarray,
213
+ t: float,
214
+ h: float,
215
+ y_prev: np.ndarray,
216
+ ) -> Tuple[np.ndarray, float, bool, int]:
217
+ """Run the full Newton loop using block-structured assembly.
218
+
219
+ Parameters
220
+ ----------
221
+ block_system : BlockStructuredSystem
222
+ y0 : ndarray
223
+ Initial guess [u; lambda].
224
+ t, h, y_prev : float, float, ndarray
225
+ """
226
+ n_p = block_system.n_phys
227
+ n_r = block_system.n_react
228
+ y = np.asarray(y0, dtype=float).ravel().copy()
229
+
230
+ for iteration in range(1, self.maxiter + 1):
231
+ blocks = block_system.assemble_blocks(y, t, h, y_prev)
232
+ H = blocks["H"]
233
+ B_top = blocks["B_top"]
234
+ B_bot = blocks["B_bot"]
235
+ C = blocks["C"]
236
+ g = blocks["g"]
237
+ h_c = blocks["h_c"]
238
+ precond_diag = blocks.get("precond_diag", None)
239
+
240
+ err = max(
241
+ float(np.linalg.norm(g)),
242
+ float(np.linalg.norm(h_c)),
243
+ )
244
+ if err < self.tol:
245
+ return y.copy(), err, True, iteration
246
+
247
+ delta_u, delta_lam, info = self.solve_linear(
248
+ H, B_top, B_bot, C, g, h_c, precond_diag,
249
+ )
250
+
251
+ alpha = self.damped_step_fraction
252
+ y[:n_p] += alpha * delta_u
253
+ y[n_p:n_p + n_r] += alpha * delta_lam
254
+
255
+ blocks = block_system.assemble_blocks(y, t, h, y_prev)
256
+ err = max(
257
+ float(np.linalg.norm(blocks["g"])),
258
+ float(np.linalg.norm(blocks["h_c"])),
259
+ )
260
+ return y.copy(), err, False, self.maxiter
261
+
262
+ # ------------------------------------------------------------------
263
+ # Multi-stage coupled Newton (RadauIIA)
264
+ # ------------------------------------------------------------------
265
+
266
+ def solve_coupled(
267
+ self,
268
+ block_system: BlockStructuredSystem,
269
+ y0: np.ndarray,
270
+ t: float,
271
+ h: float,
272
+ y_prev: np.ndarray,
273
+ rk_A: np.ndarray,
274
+ rk_c: np.ndarray,
275
+ A_phys: np.ndarray,
276
+ ) -> Tuple[np.ndarray, float, bool, int]:
277
+ """Coupled Newton for multi-stage Radau IIA with Schur reduction.
278
+
279
+ Stacks all *s* stages into a single Newton system. The velocity
280
+ block ``H_full`` carries Butcher cross-stage coupling; ``B_bot``
281
+ and ``C`` are block-diagonal (each stage's NCP depends only on
282
+ its own state). ``B_top`` has off-diagonal blocks because the
283
+ physical RHS includes the contact-force coupling ``B·λ``.
284
+
285
+ Parameters
286
+ ----------
287
+ block_system : BlockStructuredSystem
288
+ y0 : (n_aug,)
289
+ t : float
290
+ Time at beginning of step.
291
+ h : float
292
+ Full step size.
293
+ y_prev : (n_aug,)
294
+ rk_A : (s, s)
295
+ Butcher A matrix.
296
+ rk_c : (s,)
297
+ A_phys : (n_phys, n_phys)
298
+ Physical mass matrix.
299
+
300
+ Returns
301
+ -------
302
+ Y_s, err, converged, iterations
303
+ """
304
+ s = rk_A.shape[0]
305
+ n_p = block_system.n_phys
306
+ n_r = block_system.n_react
307
+ n_aug = n_p + n_r
308
+ sn_p = s * n_p
309
+ sn_r = s * n_r
310
+
311
+ A_arr = self._to_dense(A_phys)
312
+ Z = np.tile(y0, s)
313
+
314
+ for iteration in range(1, self.maxiter + 1):
315
+ # -- Per-stage block assembly --
316
+ per_stage = []
317
+ f_phys = []
318
+ for i in range(s):
319
+ Y_i = Z[i * n_aug:(i + 1) * n_aug]
320
+ aii = rk_A[i, i]
321
+ stage_h = aii * h
322
+ t_i = t + rk_c[i] * h
323
+
324
+ bi = block_system.assemble_blocks(Y_i, t_i, stage_h, y_prev)
325
+ per_stage.append(bi)
326
+
327
+ A_over_aih = A_arr / stage_h
328
+ f_i = bi["g"] + A_over_aih @ (Y_i[:n_p] - y_prev[:n_p])
329
+ f_phys.append(f_i)
330
+
331
+ # -- Stacked residual --
332
+ g_full = np.empty(sn_p)
333
+ hc_full = np.empty(sn_r)
334
+ for i in range(s):
335
+ aii = rk_A[i, i]
336
+ g_i = per_stage[i]["g"].copy()
337
+ for j in range(s):
338
+ if j != i:
339
+ g_i += (rk_A[i, j] / aii) * f_phys[j]
340
+ g_full[i * n_p:(i + 1) * n_p] = g_i
341
+ hc_full[i * n_r:(i + 1) * n_r] = per_stage[i]["h_c"]
342
+
343
+ err = max(float(np.linalg.norm(g_full)),
344
+ float(np.linalg.norm(hc_full)))
345
+ if err < self.tol:
346
+ Y_s = Z[(s - 1) * n_aug:s * n_aug]
347
+ return Y_s.copy(), err, True, iteration
348
+
349
+ # -- H_full: (s·n_p × s·n_p) with Butcher cross-stage coupling --
350
+ H_full = np.zeros((sn_p, sn_p))
351
+ for i in range(s):
352
+ aii = rk_A[i, i]
353
+ r0, r1 = i * n_p, (i + 1) * n_p
354
+ for j in range(s):
355
+ c0, c1 = j * n_p, (j + 1) * n_p
356
+ if i == j:
357
+ H_full[r0:r1, c0:c1] = self._to_dense(per_stage[i]["H"])
358
+ else:
359
+ H_j = self._to_dense(per_stage[j]["H"])
360
+ dfdu_j = A_arr / (rk_A[j, j] * h) - H_j
361
+ H_full[r0:r1, c0:c1] = -(rk_A[i, j] / aii) * dfdu_j
362
+
363
+ # -- B_top_full: (s·n_p × s·n_r) — NOT block-diagonal --
364
+ Bt_full = np.zeros((sn_p, sn_r))
365
+ for i in range(s):
366
+ aii = rk_A[i, i]
367
+ r0, r1 = i * n_p, (i + 1) * n_p
368
+ for j in range(s):
369
+ c0, c1 = j * n_r, (j + 1) * n_r
370
+ if i == j:
371
+ Bt_full[r0:r1, c0:c1] = self._to_dense(
372
+ per_stage[i]["B_top"],
373
+ )
374
+ else:
375
+ Bt_full[r0:r1, c0:c1] = (
376
+ (rk_A[i, j] / aii)
377
+ * self._to_dense(per_stage[j]["B_top"])
378
+ )
379
+
380
+ # -- B_bot_diag, C_diag: block-diagonal --
381
+ Bb_diag = np.zeros((sn_r, sn_p))
382
+ C_diag = np.zeros((sn_r, sn_r))
383
+ for i in range(s):
384
+ Bb_diag[i * n_r:(i + 1) * n_r,
385
+ i * n_p:(i + 1) * n_p] = self._to_dense(
386
+ per_stage[i]["B_bot"],
387
+ )
388
+ C_diag[i * n_r:(i + 1) * n_r,
389
+ i * n_r:(i + 1) * n_r] = self._to_dense(
390
+ per_stage[i]["C"],
391
+ )
392
+
393
+ # -- Solve the stacked saddle-point system --
394
+ delta_u, delta_lam, info = self.solve_linear(
395
+ H_full, Bt_full, Bb_diag, C_diag, g_full, hc_full,
396
+ )
397
+
398
+ alpha = self.damped_step_fraction
399
+ for i in range(s):
400
+ Z[i * n_aug:i * n_aug + n_p] += (
401
+ alpha * delta_u[i * n_p:(i + 1) * n_p]
402
+ )
403
+ Z[i * n_aug + n_p:(i + 1) * n_aug] += (
404
+ alpha * delta_lam[i * n_r:(i + 1) * n_r]
405
+ )
406
+
407
+ Y_s = Z[(s - 1) * n_aug:s * n_aug]
408
+ return Y_s.copy(), err, False, self.maxiter
@@ -26,7 +26,7 @@ The helper returns everything needed to call :func:`solve_nivp`::
26
26
 
27
27
  import numpy as np
28
28
  import scipy.sparse as sp
29
- from dataclasses import dataclass
29
+ from dataclasses import dataclass, field
30
30
  from typing import Any, Callable, Dict, List, Optional
31
31
 
32
32
  from .projections import (
@@ -80,6 +80,7 @@ class ContactSystem:
80
80
  n_phys: int
81
81
  B: Any = None
82
82
  rhs_jac: Optional[Callable] = None
83
+ solver_opts: Dict[str, Any] = field(default_factory=dict)
83
84
 
84
85
 
85
86
  # ──────────────────────────────────────────────────────────────────────