nornax 0.0.1__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.
nornax-0.0.1/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 AstroAI Lab
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
nornax-0.0.1/PKG-INFO ADDED
@@ -0,0 +1,208 @@
1
+ Metadata-Version: 2.4
2
+ Name: nornax
3
+ Version: 0.0.1
4
+ Summary: JAX-native Hermite integrators for N-body dynamics with Diffrax-facing solver hooks
5
+ Author-email: AstroAI Lab <astroai@iwr.uni-heidelberg.de>, Tobias Buck <tobias.buck@iwr.uni-heidelberg.de>
6
+ Maintainer-email: AstroAI Lab <astroai@iwr.uni-heidelberg.de>
7
+ License-Expression: MIT
8
+ Project-URL: Homepage, https://github.com/TobiBu/nornax
9
+ Project-URL: Repository, https://github.com/TobiBu/nornax
10
+ Project-URL: Issues, https://github.com/TobiBu/nornax/issues
11
+ Project-URL: Documentation, https://github.com/TobiBu/nornax#readme
12
+ Keywords: hermite,n-body,jax,gravity,integrator,fmm
13
+ Classifier: Development Status :: 3 - Alpha
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Scientific/Engineering :: Physics
20
+ Requires-Python: >=3.11
21
+ Description-Content-Type: text/markdown
22
+ License-File: LICENSE
23
+ Requires-Dist: jax>=0.8.0
24
+ Requires-Dist: jaxlib>=0.8.0
25
+ Requires-Dist: diffrax>=0.6.0
26
+ Requires-Dist: equinox>=0.11.10
27
+ Requires-Dist: jaxtyping>=0.2.34
28
+ Requires-Dist: beartype>=0.14.0
29
+ Provides-Extra: dev
30
+ Requires-Dist: black>=24.8.0; extra == "dev"
31
+ Requires-Dist: isort>=5.13.2; extra == "dev"
32
+ Requires-Dist: pydoclint>=0.6.0; extra == "dev"
33
+ Requires-Dist: pytest>=8.3.0; extra == "dev"
34
+ Requires-Dist: pytest-cov>=5.0.0; extra == "dev"
35
+ Requires-Dist: pre-commit>=3.8.0; extra == "dev"
36
+ Requires-Dist: build>=1.2.2; extra == "dev"
37
+ Requires-Dist: twine>=5.1.1; extra == "dev"
38
+ Dynamic: license-file
39
+
40
+ # nornax
41
+
42
+ <p align="center">
43
+ <img src="./nornax.png" alt="nornax Logo" width="420" />
44
+ </p>
45
+
46
+ `nornax` is a JAX-native Hermite integrator package for gravitational N-body
47
+ dynamics.
48
+
49
+ The current direction is:
50
+
51
+ - standalone and general-purpose first
52
+ - Diffrax-facing solver design
53
+ - GPU-efficient JAX kernels
54
+ - clean backend adapters so `jaccpot` can plug in later
55
+
56
+ ## Current Scope
57
+
58
+ The current implementation includes:
59
+
60
+ - immutable `NBodyState` / `ForceDerivatives` PyTrees
61
+ - backend-agnostic `ForceModel` protocol
62
+ - standalone `DirectSumGravity` reference backend (O(N^2) all-pairs; the
63
+ default path materializes N×N pair tensors, and an opt-in `block_size`
64
+ caps peak memory at O(block_size × N) for larger N — `jaccpot` remains the
65
+ scalable backend)
66
+ - standalone and Diffrax-backed Hermite-4, Hermite-6, and Hermite-8
67
+ - adaptive global timestep control through Diffrax
68
+ - diagnostics for total energy and angular momentum
69
+ - convergence and long-run conservation tests
70
+
71
+ ## Installation
72
+
73
+ Install from source:
74
+
75
+ ```bash
76
+ pip install -e .
77
+ ```
78
+
79
+ Install with development tooling:
80
+
81
+ ```bash
82
+ pip install -e ".[dev]"
83
+ ```
84
+
85
+ ## Quick Start
86
+
87
+ ```python
88
+ import jax
89
+ import jax.numpy as jnp
90
+
91
+ from nornax import AarsethController, solve_adaptive_to_time, total_energy
92
+ from nornax.forces import DirectSumGravity
93
+
94
+ jax.config.update("jax_enable_x64", True)
95
+
96
+ force_model = DirectSumGravity()
97
+ result = solve_adaptive_to_time(
98
+ jnp.asarray([[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]),
99
+ jnp.asarray([[0.0, 0.5, 0.0], [0.0, -0.5, 0.0]]),
100
+ jnp.asarray([1.0, 1.0]),
101
+ force_model,
102
+ t_final=1.0,
103
+ order=8,
104
+ controller=AarsethController(eta=0.03, min_dt=1.0e-4, max_dt=5.0e-2),
105
+ atol=1.0e-8,
106
+ )
107
+ print(result.final_state.time, total_energy(result.final_state))
108
+ ```
109
+
110
+ Examples:
111
+
112
+ - [examples/two_body_diffrax.py](examples/two_body_diffrax.py): adaptive two-body solve
113
+ - [examples/compare_hermite_orders.py](examples/compare_hermite_orders.py): compare Hermite-4/6/8 energy and angular-momentum drift
114
+ - [examples/plummer_sphere_gpu.py](examples/plummer_sphere_gpu.py): small Plummer sphere run shaped for GPU/server use
115
+
116
+ Notebooks:
117
+
118
+ - [notebooks/hermite_orders_accuracy.ipynb](notebooks/hermite_orders_accuracy.ipynb): compare adaptive Hermite-4/6/8 drift and step counts
119
+ - [notebooks/hermite_vs_diffrax_rk.ipynb](notebooks/hermite_vs_diffrax_rk.ipynb): compare Hermite against a standard Diffrax RK solver
120
+ - [notebooks/direct_vs_jaccpot_backend.ipynb](notebooks/direct_vs_jaccpot_backend.ipynb): compare direct-sum and `jaccpot` backends for Hermite-4
121
+ - [notebooks/plummer_sphere_gpu.ipynb](notebooks/plummer_sphere_gpu.ipynb): small Plummer sphere application notebook for GPU/server runs
122
+
123
+ ## Jaccpot Adapter
124
+
125
+ `nornax` now includes a first adapter for `jaccpot` via `JaccpotForceModel`.
126
+ This is useful today for Hermite-4, because the current `jaccpot` runtime
127
+ exposes acceleration and jerk, but not higher time derivatives.
128
+
129
+ ```python
130
+ import jax.numpy as jnp
131
+
132
+ from jaccpot import FastMultipoleMethod
133
+ from nornax import AarsethController, JaccpotForceModel, solve_adaptive_to_time
134
+
135
+ solver = FastMultipoleMethod(preset="fast", basis="solidfmm")
136
+ force_model = JaccpotForceModel(solver)
137
+
138
+ result = solve_adaptive_to_time(
139
+ jnp.asarray([[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]]),
140
+ jnp.asarray([[0.0, 0.2, 0.0], [0.0, -0.2, 0.0]]),
141
+ jnp.asarray([1.0, 1.0]),
142
+ force_model,
143
+ t_final=0.1,
144
+ order=4,
145
+ controller=AarsethController(eta=0.03, min_dt=1.0e-4, max_dt=5.0e-2),
146
+ atol=1.0e-6,
147
+ args={"leaf_size": 8, "max_order": 2, "jerk_mode": "fast_approx"},
148
+ )
149
+ ```
150
+
151
+ Current limitation:
152
+
153
+ - use `JaccpotForceModel` with Hermite-4 for now
154
+ - Hermite-6 and Hermite-8 still require higher time derivatives than `jaccpot` currently provides
155
+
156
+ ## Planned Architecture
157
+
158
+ - `nornax.state`: particle state and cached force derivatives
159
+ - `nornax.forces`: standalone direct-sum backend plus future adapters
160
+ - `nornax.solvers`: Hermite kernels and Diffrax-facing solver classes
161
+ - `nornax.terms`: thin Diffrax integration hooks
162
+
163
+ The scientific target is the family of higher-order Hermite methods discussed
164
+ by Nitadori, Iwasawa, and Makino. The current implementation follows that paper
165
+ through Hermite-8, including direct derivatives through `crackle` and
166
+ higher-order adaptive timestep criteria.
167
+
168
+ ## Validation
169
+
170
+ The test suite currently covers:
171
+
172
+ - direct-force derivatives through `crackle`
173
+ - Hermite-4/6/8 kernel behavior
174
+ - Diffrax custom solver smoke tests for 4/6/8
175
+ - convergence checks for 4/6/8
176
+ - adaptive public solve APIs
177
+ - long-run two-body energy-drift ordering across 4/6/8
178
+
179
+ Benchmarks and examples live in [bench/bench_direct_sum.py](bench/bench_direct_sum.py)
180
+ and [examples/compare_hermite_orders.py](examples/compare_hermite_orders.py).
181
+
182
+ ## Development
183
+
184
+ Run quality gates locally:
185
+
186
+ ```bash
187
+ black --check .
188
+ isort --check-only .
189
+ pytest
190
+ ```
191
+
192
+ Or run pre-commit hooks:
193
+
194
+ ```bash
195
+ pre-commit run --all-files
196
+ ```
197
+
198
+ ## Runtime Type Checking
199
+
200
+ Enable package-wide runtime checks (`jaxtyping` + `beartype`) at import time:
201
+
202
+ ```bash
203
+ export NORNAX_RUNTIME_TYPECHECK=1
204
+ ```
205
+
206
+ ## License
207
+
208
+ MIT.
nornax-0.0.1/README.md ADDED
@@ -0,0 +1,169 @@
1
+ # nornax
2
+
3
+ <p align="center">
4
+ <img src="./nornax.png" alt="nornax Logo" width="420" />
5
+ </p>
6
+
7
+ `nornax` is a JAX-native Hermite integrator package for gravitational N-body
8
+ dynamics.
9
+
10
+ The current direction is:
11
+
12
+ - standalone and general-purpose first
13
+ - Diffrax-facing solver design
14
+ - GPU-efficient JAX kernels
15
+ - clean backend adapters so `jaccpot` can plug in later
16
+
17
+ ## Current Scope
18
+
19
+ The current implementation includes:
20
+
21
+ - immutable `NBodyState` / `ForceDerivatives` PyTrees
22
+ - backend-agnostic `ForceModel` protocol
23
+ - standalone `DirectSumGravity` reference backend (O(N^2) all-pairs; the
24
+ default path materializes N×N pair tensors, and an opt-in `block_size`
25
+ caps peak memory at O(block_size × N) for larger N — `jaccpot` remains the
26
+ scalable backend)
27
+ - standalone and Diffrax-backed Hermite-4, Hermite-6, and Hermite-8
28
+ - adaptive global timestep control through Diffrax
29
+ - diagnostics for total energy and angular momentum
30
+ - convergence and long-run conservation tests
31
+
32
+ ## Installation
33
+
34
+ Install from source:
35
+
36
+ ```bash
37
+ pip install -e .
38
+ ```
39
+
40
+ Install with development tooling:
41
+
42
+ ```bash
43
+ pip install -e ".[dev]"
44
+ ```
45
+
46
+ ## Quick Start
47
+
48
+ ```python
49
+ import jax
50
+ import jax.numpy as jnp
51
+
52
+ from nornax import AarsethController, solve_adaptive_to_time, total_energy
53
+ from nornax.forces import DirectSumGravity
54
+
55
+ jax.config.update("jax_enable_x64", True)
56
+
57
+ force_model = DirectSumGravity()
58
+ result = solve_adaptive_to_time(
59
+ jnp.asarray([[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]),
60
+ jnp.asarray([[0.0, 0.5, 0.0], [0.0, -0.5, 0.0]]),
61
+ jnp.asarray([1.0, 1.0]),
62
+ force_model,
63
+ t_final=1.0,
64
+ order=8,
65
+ controller=AarsethController(eta=0.03, min_dt=1.0e-4, max_dt=5.0e-2),
66
+ atol=1.0e-8,
67
+ )
68
+ print(result.final_state.time, total_energy(result.final_state))
69
+ ```
70
+
71
+ Examples:
72
+
73
+ - [examples/two_body_diffrax.py](examples/two_body_diffrax.py): adaptive two-body solve
74
+ - [examples/compare_hermite_orders.py](examples/compare_hermite_orders.py): compare Hermite-4/6/8 energy and angular-momentum drift
75
+ - [examples/plummer_sphere_gpu.py](examples/plummer_sphere_gpu.py): small Plummer sphere run shaped for GPU/server use
76
+
77
+ Notebooks:
78
+
79
+ - [notebooks/hermite_orders_accuracy.ipynb](notebooks/hermite_orders_accuracy.ipynb): compare adaptive Hermite-4/6/8 drift and step counts
80
+ - [notebooks/hermite_vs_diffrax_rk.ipynb](notebooks/hermite_vs_diffrax_rk.ipynb): compare Hermite against a standard Diffrax RK solver
81
+ - [notebooks/direct_vs_jaccpot_backend.ipynb](notebooks/direct_vs_jaccpot_backend.ipynb): compare direct-sum and `jaccpot` backends for Hermite-4
82
+ - [notebooks/plummer_sphere_gpu.ipynb](notebooks/plummer_sphere_gpu.ipynb): small Plummer sphere application notebook for GPU/server runs
83
+
84
+ ## Jaccpot Adapter
85
+
86
+ `nornax` now includes a first adapter for `jaccpot` via `JaccpotForceModel`.
87
+ This is useful today for Hermite-4, because the current `jaccpot` runtime
88
+ exposes acceleration and jerk, but not higher time derivatives.
89
+
90
+ ```python
91
+ import jax.numpy as jnp
92
+
93
+ from jaccpot import FastMultipoleMethod
94
+ from nornax import AarsethController, JaccpotForceModel, solve_adaptive_to_time
95
+
96
+ solver = FastMultipoleMethod(preset="fast", basis="solidfmm")
97
+ force_model = JaccpotForceModel(solver)
98
+
99
+ result = solve_adaptive_to_time(
100
+ jnp.asarray([[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]]),
101
+ jnp.asarray([[0.0, 0.2, 0.0], [0.0, -0.2, 0.0]]),
102
+ jnp.asarray([1.0, 1.0]),
103
+ force_model,
104
+ t_final=0.1,
105
+ order=4,
106
+ controller=AarsethController(eta=0.03, min_dt=1.0e-4, max_dt=5.0e-2),
107
+ atol=1.0e-6,
108
+ args={"leaf_size": 8, "max_order": 2, "jerk_mode": "fast_approx"},
109
+ )
110
+ ```
111
+
112
+ Current limitation:
113
+
114
+ - use `JaccpotForceModel` with Hermite-4 for now
115
+ - Hermite-6 and Hermite-8 still require higher time derivatives than `jaccpot` currently provides
116
+
117
+ ## Planned Architecture
118
+
119
+ - `nornax.state`: particle state and cached force derivatives
120
+ - `nornax.forces`: standalone direct-sum backend plus future adapters
121
+ - `nornax.solvers`: Hermite kernels and Diffrax-facing solver classes
122
+ - `nornax.terms`: thin Diffrax integration hooks
123
+
124
+ The scientific target is the family of higher-order Hermite methods discussed
125
+ by Nitadori, Iwasawa, and Makino. The current implementation follows that paper
126
+ through Hermite-8, including direct derivatives through `crackle` and
127
+ higher-order adaptive timestep criteria.
128
+
129
+ ## Validation
130
+
131
+ The test suite currently covers:
132
+
133
+ - direct-force derivatives through `crackle`
134
+ - Hermite-4/6/8 kernel behavior
135
+ - Diffrax custom solver smoke tests for 4/6/8
136
+ - convergence checks for 4/6/8
137
+ - adaptive public solve APIs
138
+ - long-run two-body energy-drift ordering across 4/6/8
139
+
140
+ Benchmarks and examples live in [bench/bench_direct_sum.py](bench/bench_direct_sum.py)
141
+ and [examples/compare_hermite_orders.py](examples/compare_hermite_orders.py).
142
+
143
+ ## Development
144
+
145
+ Run quality gates locally:
146
+
147
+ ```bash
148
+ black --check .
149
+ isort --check-only .
150
+ pytest
151
+ ```
152
+
153
+ Or run pre-commit hooks:
154
+
155
+ ```bash
156
+ pre-commit run --all-files
157
+ ```
158
+
159
+ ## Runtime Type Checking
160
+
161
+ Enable package-wide runtime checks (`jaxtyping` + `beartype`) at import time:
162
+
163
+ ```bash
164
+ export NORNAX_RUNTIME_TYPECHECK=1
165
+ ```
166
+
167
+ ## License
168
+
169
+ MIT.
@@ -0,0 +1,44 @@
1
+ """Nornax: JAX-native Hermite integrators for N-body dynamics."""
2
+
3
+ from ._typecheck import enable_runtime_typecheck
4
+
5
+ # Install the optional runtime type-checking hook before importing any nornax
6
+ # submodules; the jaxtyping hook only instruments modules imported after it.
7
+ enable_runtime_typecheck()
8
+
9
+ from .adapters import JaccpotForceModel, JaccpotOptions # noqa: E402
10
+ from .controllers import AarsethController, AdaptiveStepPolicy # noqa: E402
11
+ from .diagnostics import ( # noqa: E402
12
+ gravitational_potential_energy,
13
+ total_angular_momentum,
14
+ total_energy,
15
+ )
16
+ from .initial_conditions import sample_plummer_sphere # noqa: E402
17
+ from .initialize import initialize_state # noqa: E402
18
+ from .solve import ( # noqa: E402
19
+ solve_adaptive_hermite4,
20
+ solve_adaptive_hermite4_to_time,
21
+ solve_adaptive_hermite6_to_time,
22
+ solve_adaptive_hermite8_to_time,
23
+ solve_adaptive_to_time,
24
+ )
25
+ from .state import ForceDerivatives, NBodyState # noqa: E402
26
+
27
+ __all__ = [
28
+ "AarsethController",
29
+ "AdaptiveStepPolicy",
30
+ "ForceDerivatives",
31
+ "JaccpotForceModel",
32
+ "JaccpotOptions",
33
+ "NBodyState",
34
+ "gravitational_potential_energy",
35
+ "initialize_state",
36
+ "sample_plummer_sphere",
37
+ "solve_adaptive_to_time",
38
+ "solve_adaptive_hermite4",
39
+ "solve_adaptive_hermite4_to_time",
40
+ "solve_adaptive_hermite6_to_time",
41
+ "solve_adaptive_hermite8_to_time",
42
+ "total_angular_momentum",
43
+ "total_energy",
44
+ ]
@@ -0,0 +1,24 @@
1
+ """Runtime type-checking opt-in for nornax.
2
+
3
+ Set the environment variable ``NORNAX_RUNTIME_TYPECHECK=1`` before importing
4
+ ``nornax`` to enable package-wide ``jaxtyping`` + ``beartype`` instrumentation.
5
+
6
+ The jaxtyping import hook only instruments modules imported *after* it is
7
+ installed, so ``nornax.__init__`` calls :func:`enable_runtime_typecheck` before
8
+ importing any nornax submodules.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import os
14
+
15
+
16
+ def enable_runtime_typecheck() -> None:
17
+ """Install the jaxtyping + beartype import hook when requested via env var."""
18
+ if os.environ.get("NORNAX_RUNTIME_TYPECHECK", "0") != "1":
19
+ return
20
+ try:
21
+ from jaxtyping import install_import_hook
22
+ except ImportError: # pragma: no cover - jaxtyping is a hard dependency
23
+ return
24
+ install_import_hook("nornax", "beartype.beartype")
@@ -0,0 +1,26 @@
1
+ """Shared jaxtyping shape aliases for nornax array annotations.
2
+
3
+ These aliases document (and, when ``NORNAX_RUNTIME_TYPECHECK=1``, enforce) the
4
+ array shapes flowing through the N-body kernels. Axis ``n`` is the particle
5
+ count; jaxtyping binds it consistently within a single function call, so a
6
+ mismatch between, say, ``positions`` and ``masses`` is caught.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from jax import Array
12
+ from jaxtyping import Float
13
+
14
+ # Per-particle 3-vector field: positions, velocities, and every acceleration
15
+ # time derivative (jerk, snap, crackle, ...).
16
+ Vec3 = Float[Array, "n 3"]
17
+
18
+ # Per-particle scalar field (e.g. masses).
19
+ PerParticle = Float[Array, "n"]
20
+
21
+ # A 0-d array scalar (e.g. cached time). Python floats are accepted separately
22
+ # at call boundaries via ``ScalarLike``.
23
+ Scalar = Float[Array, ""]
24
+
25
+ # A scalar accepted at an API boundary: either a 0-d array or a Python float.
26
+ ScalarLike = Scalar | float
@@ -0,0 +1,5 @@
1
+ """Optional backend adapters for sibling projects."""
2
+
3
+ from .jaccpot import JaccpotForceModel, JaccpotOptions
4
+
5
+ __all__ = ["JaccpotForceModel", "JaccpotOptions"]
@@ -0,0 +1,115 @@
1
+ """Adapter from ``jaccpot`` FMM solvers to the Nornax force-model protocol."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from collections.abc import Mapping
6
+ from dataclasses import asdict, dataclass
7
+ from typing import Any
8
+
9
+ import jax.numpy as jnp
10
+
11
+ from nornax._typing import PerParticle, ScalarLike, Vec3
12
+ from nornax.state import ForceDerivatives
13
+
14
+
15
+ @dataclass(frozen=True)
16
+ class JaccpotOptions:
17
+ """Execution options forwarded to ``jaccpot`` runtime calls.
18
+
19
+ Note that ``max_order`` here is the *FMM multipole expansion order* passed
20
+ to the ``jaccpot`` solver. It is unrelated to the ``max_order`` argument of
21
+ :meth:`ForceModel.derivatives`, which selects how many *time* derivatives of
22
+ the acceleration (jerk, snap, ...) the Hermite solver requests. The name is
23
+ kept to match the ``jaccpot`` runtime kwarg it maps onto.
24
+ """
25
+
26
+ target_indices: jnp.ndarray | None = None
27
+ bounds: tuple[jnp.ndarray, jnp.ndarray] | None = None
28
+ leaf_size: int = 16
29
+ max_order: int = 2 # FMM expansion order (see class docstring)
30
+ theta: float | None = None
31
+ jit_tree: bool | None = None
32
+ refine_local: bool | None = None
33
+ max_refine_levels: int | None = None
34
+ aspect_threshold: float | None = None
35
+ jit_traversal: bool | None = None
36
+ reuse_prepared_state: bool = False
37
+ jerk_mode: str = "accurate"
38
+ jerk_fd_dt: float = 1.0e-3
39
+
40
+ def to_kwargs(self) -> dict[str, Any]:
41
+ """Convert options to runtime kwargs, dropping ``None`` entries."""
42
+ return {key: value for key, value in asdict(self).items() if value is not None}
43
+
44
+
45
+ @dataclass(frozen=True)
46
+ class JaccpotForceModel:
47
+ """Thin adapter exposing a ``jaccpot`` solver via the Nornax protocol.
48
+
49
+ Current `jaccpot` support covers acceleration and jerk only, so this
50
+ adapter is suitable for Hermite-4 today. Requests for higher time
51
+ derivatives raise ``NotImplementedError``.
52
+ """
53
+
54
+ solver: Any
55
+ options: JaccpotOptions = JaccpotOptions()
56
+
57
+ def derivatives(
58
+ self,
59
+ t: ScalarLike,
60
+ positions: Vec3,
61
+ velocities: Vec3,
62
+ masses: PerParticle,
63
+ *,
64
+ max_order: int,
65
+ args: object = None,
66
+ ) -> ForceDerivatives:
67
+ """Return acceleration derivatives supported by the underlying solver."""
68
+ del t
69
+ runtime_kwargs = self._resolve_runtime_kwargs(args)
70
+ if max_order < 1:
71
+ raise ValueError("max_order must be >= 1")
72
+ if max_order > 2:
73
+ raise NotImplementedError(
74
+ "jaccpot currently exposes time derivatives only through jerk; "
75
+ "use this adapter with Hermite-4 for now"
76
+ )
77
+ if max_order == 1:
78
+ accelerations = self.solver.compute_accelerations(
79
+ positions,
80
+ masses,
81
+ **self._acceleration_kwargs(runtime_kwargs),
82
+ )
83
+ return ForceDerivatives(acc=accelerations)
84
+ accelerations, jerk = self.solver.compute_accelerations_and_jerk(
85
+ positions,
86
+ masses,
87
+ velocities,
88
+ **runtime_kwargs,
89
+ )
90
+ return ForceDerivatives(acc=accelerations, jerk=jerk)
91
+
92
+ def _resolve_runtime_kwargs(self, args: object) -> dict[str, Any]:
93
+ """Merge default adapter options with per-call overrides."""
94
+ runtime_kwargs = self.options.to_kwargs()
95
+ if args is None:
96
+ return runtime_kwargs
97
+ if isinstance(args, JaccpotOptions):
98
+ runtime_kwargs.update(args.to_kwargs())
99
+ return runtime_kwargs
100
+ if isinstance(args, Mapping):
101
+ runtime_kwargs.update(
102
+ {key: value for key, value in args.items() if value is not None}
103
+ )
104
+ return runtime_kwargs
105
+ raise TypeError(
106
+ "jaccpot adapter args must be None, JaccpotOptions, or a mapping of runtime kwargs"
107
+ )
108
+
109
+ @staticmethod
110
+ def _acceleration_kwargs(runtime_kwargs: Mapping[str, Any]) -> dict[str, Any]:
111
+ """Drop jerk-only options for acceleration-only calls."""
112
+ jerk_only = {"jerk_mode", "jerk_fd_dt"}
113
+ return {
114
+ key: value for key, value in runtime_kwargs.items() if key not in jerk_only
115
+ }
@@ -0,0 +1,17 @@
1
+ """Adaptive timestep controllers for Nornax."""
2
+
3
+ from .aarseth import (
4
+ AarsethController,
5
+ AdaptiveStepPolicy,
6
+ aarseth_timestep,
7
+ aarseth_timestep_6th_order,
8
+ aarseth_timestep_8th_order,
9
+ )
10
+
11
+ __all__ = [
12
+ "AarsethController",
13
+ "AdaptiveStepPolicy",
14
+ "aarseth_timestep",
15
+ "aarseth_timestep_6th_order",
16
+ "aarseth_timestep_8th_order",
17
+ ]