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 +21 -0
- nornax-0.0.1/PKG-INFO +208 -0
- nornax-0.0.1/README.md +169 -0
- nornax-0.0.1/nornax/__init__.py +44 -0
- nornax-0.0.1/nornax/_typecheck.py +24 -0
- nornax-0.0.1/nornax/_typing.py +26 -0
- nornax-0.0.1/nornax/adapters/__init__.py +5 -0
- nornax-0.0.1/nornax/adapters/jaccpot.py +115 -0
- nornax-0.0.1/nornax/controllers/__init__.py +17 -0
- nornax-0.0.1/nornax/controllers/aarseth.py +196 -0
- nornax-0.0.1/nornax/diagnostics.py +48 -0
- nornax-0.0.1/nornax/forces/__init__.py +6 -0
- nornax-0.0.1/nornax/forces/base.py +30 -0
- nornax-0.0.1/nornax/forces/direct.py +259 -0
- nornax-0.0.1/nornax/initial_conditions.py +113 -0
- nornax-0.0.1/nornax/initialize.py +37 -0
- nornax-0.0.1/nornax/py.typed +0 -0
- nornax-0.0.1/nornax/solve.py +289 -0
- nornax-0.0.1/nornax/solvers/__init__.py +23 -0
- nornax-0.0.1/nornax/solvers/hermite4.py +351 -0
- nornax-0.0.1/nornax/solvers/hermite6.py +235 -0
- nornax-0.0.1/nornax/solvers/hermite8.py +295 -0
- nornax-0.0.1/nornax/state.py +48 -0
- nornax-0.0.1/nornax/terms.py +86 -0
- nornax-0.0.1/nornax.egg-info/PKG-INFO +208 -0
- nornax-0.0.1/nornax.egg-info/SOURCES.txt +29 -0
- nornax-0.0.1/nornax.egg-info/dependency_links.txt +1 -0
- nornax-0.0.1/nornax.egg-info/requires.txt +16 -0
- nornax-0.0.1/nornax.egg-info/top_level.txt +1 -0
- nornax-0.0.1/pyproject.toml +99 -0
- nornax-0.0.1/setup.cfg +4 -0
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,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
|
+
]
|