CPax 2026.10.0__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.
- cpax-2026.10.0/LICENSE +21 -0
- cpax-2026.10.0/PKG-INFO +80 -0
- cpax-2026.10.0/README.md +56 -0
- cpax-2026.10.0/pyproject.toml +109 -0
- cpax-2026.10.0/setup.cfg +4 -0
- cpax-2026.10.0/src/CPax.egg-info/PKG-INFO +80 -0
- cpax-2026.10.0/src/CPax.egg-info/SOURCES.txt +43 -0
- cpax-2026.10.0/src/CPax.egg-info/dependency_links.txt +1 -0
- cpax-2026.10.0/src/CPax.egg-info/requires.txt +6 -0
- cpax-2026.10.0/src/CPax.egg-info/top_level.txt +1 -0
- cpax-2026.10.0/src/cpax/__init__.py +9 -0
- cpax-2026.10.0/src/cpax/_backends/__init__.py +1 -0
- cpax-2026.10.0/src/cpax/_backends/_diffrax.py +105 -0
- cpax-2026.10.0/src/cpax/dynamics/__init__.py +6 -0
- cpax-2026.10.0/src/cpax/dynamics/_lorenz.py +43 -0
- cpax-2026.10.0/src/cpax/dynamics/_systems.py +55 -0
- cpax-2026.10.0/src/cpax/mechanics/__init__.py +8 -0
- cpax-2026.10.0/src/cpax/mechanics/_gravity.py +38 -0
- cpax-2026.10.0/src/cpax/mechanics/_hamiltonian.py +49 -0
- cpax-2026.10.0/src/cpax/mechanics/_newtonian.py +107 -0
- cpax-2026.10.0/src/cpax/mechanics/_state.py +47 -0
- cpax-2026.10.0/src/cpax/ode/__init__.py +5 -0
- cpax-2026.10.0/src/cpax/ode/integrators.py +92 -0
- cpax-2026.10.0/src/cpax/problems/__init__.py +5 -0
- cpax-2026.10.0/src/cpax/problems/_ivp.py +114 -0
- cpax-2026.10.0/src/cpax/results/__init__.py +5 -0
- cpax-2026.10.0/src/cpax/results/_solution.py +85 -0
- cpax-2026.10.0/src/cpax/solvers/__init__.py +7 -0
- cpax-2026.10.0/src/cpax/solvers/_auto.py +12 -0
- cpax-2026.10.0/src/cpax/solvers/_saveat.py +58 -0
- cpax-2026.10.0/src/cpax/solvers/_solve.py +109 -0
- cpax-2026.10.0/tests/test_dynamics.py +71 -0
- cpax-2026.10.0/tests/test_examples.py +64 -0
- cpax-2026.10.0/tests/test_gravity.py +29 -0
- cpax-2026.10.0/tests/test_leapfrog.py +44 -0
- cpax-2026.10.0/tests/test_lorenz.py +44 -0
- cpax-2026.10.0/tests/test_mechanics.py +207 -0
- cpax-2026.10.0/tests/test_mechanics_state.py +61 -0
- cpax-2026.10.0/tests/test_problems.py +113 -0
- cpax-2026.10.0/tests/test_pypi_tag.py +87 -0
- cpax-2026.10.0/tests/test_results.py +93 -0
- cpax-2026.10.0/tests/test_saveat.py +116 -0
- cpax-2026.10.0/tests/test_solve.py +366 -0
- cpax-2026.10.0/tests/test_solve_jax.py +158 -0
- cpax-2026.10.0/tests/test_versioning.py +110 -0
cpax-2026.10.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Jan Lukas Späh
|
|
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.
|
cpax-2026.10.0/PKG-INFO
ADDED
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: CPax
|
|
3
|
+
Version: 2026.10.0
|
|
4
|
+
Summary: CPax: A JAX-based computational physics library.
|
|
5
|
+
Author: Jan Lukas Späh
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Source, https://github.com/JaLuka98/CPax
|
|
8
|
+
Project-URL: Tracker, https://github.com/JaLuka98/CPax/issues
|
|
9
|
+
Keywords: physics,computational,jax,autograd,ODE,PDE
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Requires-Python: >=3.12
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
License-File: LICENSE
|
|
17
|
+
Requires-Dist: diffrax<0.8,>=0.7.2
|
|
18
|
+
Requires-Dist: equinox>=0.13.8
|
|
19
|
+
Requires-Dist: jax>=0.11
|
|
20
|
+
Requires-Dist: numpy>=2.0
|
|
21
|
+
Requires-Dist: matplotlib>=3.7
|
|
22
|
+
Requires-Dist: pillow>=13.0.0
|
|
23
|
+
Dynamic: license-file
|
|
24
|
+
|
|
25
|
+
# CPax
|
|
26
|
+
|
|
27
|
+
**CPax** is an open-source, JAX-based library for computational physics.
|
|
28
|
+
It separates physical **systems**, **problems**, **solvers** and
|
|
29
|
+
**solutions** in reusable workflows. The current implementation supports
|
|
30
|
+
continuous-time initial-value problems and classical mechanics; additional
|
|
31
|
+
physical domains are planned, not yet implemented.
|
|
32
|
+
|
|
33
|
+
**[Documentation](https://cpax.readthedocs.io/en/latest/)** ·
|
|
34
|
+
**[Source](https://github.com/JaLuka98/CPax)** ·
|
|
35
|
+
**[Issues](https://github.com/JaLuka98/CPax/issues)**
|
|
36
|
+
|
|
37
|
+
## Installation
|
|
38
|
+
|
|
39
|
+
Clone the repository, then run `uv sync` to create a reproducible
|
|
40
|
+
development environment. CPax currently requires Python 3.12 or newer.
|
|
41
|
+
|
|
42
|
+
## Quickstart
|
|
43
|
+
|
|
44
|
+
```python
|
|
45
|
+
import cpax
|
|
46
|
+
from cpax.mechanics import NewtonianSystem, PhaseSpaceState
|
|
47
|
+
|
|
48
|
+
system = NewtonianSystem(potential=lambda q: q**2 / 2, mass=1.0)
|
|
49
|
+
problem = cpax.IVP(system, PhaseSpaceState(1.0, 0.0), (0.0, 10.0))
|
|
50
|
+
solution = cpax.solve(problem, saveat=cpax.SaveAt.linspace(500))
|
|
51
|
+
|
|
52
|
+
print(solution.ys.q) # Position
|
|
53
|
+
print(solution.ys.p) # Canonical momentum
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
The same workflow supports the conservative Duffing oscillator, the
|
|
57
|
+
built-in `cpax.dynamics.LorenzSystem`, and multibody Newtonian gravity.
|
|
58
|
+
See the [examples](examples/) and documentation for complete scripts,
|
|
59
|
+
plots, and the animated Sun–Earth–Mars–Jupiter–comet showcase.
|
|
60
|
+
|
|
61
|
+
The default `solver="auto"` currently selects Diffrax Tsit5 for IVPs.
|
|
62
|
+
It does not infer stiffness or symplectic structure. A standalone
|
|
63
|
+
symplectic leapfrog integrator is preserved in `cpax.ode`.
|
|
64
|
+
|
|
65
|
+
## Development
|
|
66
|
+
|
|
67
|
+
CPax uses **uv**, **Ruff**, and **ty**:
|
|
68
|
+
|
|
69
|
+
```shell
|
|
70
|
+
uv sync
|
|
71
|
+
uv run pytest
|
|
72
|
+
uv run ruff check .
|
|
73
|
+
uv run ruff format --check .
|
|
74
|
+
uv run ty check
|
|
75
|
+
uv run sphinx-build -W -b html docs/source docs/_build/html
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
Releases use calendar versions `YYYY.M.N`. The manual GitHub Actions
|
|
79
|
+
release workflow supports non-publishing previews and publishing from
|
|
80
|
+
`main`. See the workflow for its latest controls and validation. Publishing to PyPI\nis a separate, manually approved process described in\n[the release guide](docs/publishing.md).
|
cpax-2026.10.0/README.md
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
# CPax
|
|
2
|
+
|
|
3
|
+
**CPax** is an open-source, JAX-based library for computational physics.
|
|
4
|
+
It separates physical **systems**, **problems**, **solvers** and
|
|
5
|
+
**solutions** in reusable workflows. The current implementation supports
|
|
6
|
+
continuous-time initial-value problems and classical mechanics; additional
|
|
7
|
+
physical domains are planned, not yet implemented.
|
|
8
|
+
|
|
9
|
+
**[Documentation](https://cpax.readthedocs.io/en/latest/)** ·
|
|
10
|
+
**[Source](https://github.com/JaLuka98/CPax)** ·
|
|
11
|
+
**[Issues](https://github.com/JaLuka98/CPax/issues)**
|
|
12
|
+
|
|
13
|
+
## Installation
|
|
14
|
+
|
|
15
|
+
Clone the repository, then run `uv sync` to create a reproducible
|
|
16
|
+
development environment. CPax currently requires Python 3.12 or newer.
|
|
17
|
+
|
|
18
|
+
## Quickstart
|
|
19
|
+
|
|
20
|
+
```python
|
|
21
|
+
import cpax
|
|
22
|
+
from cpax.mechanics import NewtonianSystem, PhaseSpaceState
|
|
23
|
+
|
|
24
|
+
system = NewtonianSystem(potential=lambda q: q**2 / 2, mass=1.0)
|
|
25
|
+
problem = cpax.IVP(system, PhaseSpaceState(1.0, 0.0), (0.0, 10.0))
|
|
26
|
+
solution = cpax.solve(problem, saveat=cpax.SaveAt.linspace(500))
|
|
27
|
+
|
|
28
|
+
print(solution.ys.q) # Position
|
|
29
|
+
print(solution.ys.p) # Canonical momentum
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
The same workflow supports the conservative Duffing oscillator, the
|
|
33
|
+
built-in `cpax.dynamics.LorenzSystem`, and multibody Newtonian gravity.
|
|
34
|
+
See the [examples](examples/) and documentation for complete scripts,
|
|
35
|
+
plots, and the animated Sun–Earth–Mars–Jupiter–comet showcase.
|
|
36
|
+
|
|
37
|
+
The default `solver="auto"` currently selects Diffrax Tsit5 for IVPs.
|
|
38
|
+
It does not infer stiffness or symplectic structure. A standalone
|
|
39
|
+
symplectic leapfrog integrator is preserved in `cpax.ode`.
|
|
40
|
+
|
|
41
|
+
## Development
|
|
42
|
+
|
|
43
|
+
CPax uses **uv**, **Ruff**, and **ty**:
|
|
44
|
+
|
|
45
|
+
```shell
|
|
46
|
+
uv sync
|
|
47
|
+
uv run pytest
|
|
48
|
+
uv run ruff check .
|
|
49
|
+
uv run ruff format --check .
|
|
50
|
+
uv run ty check
|
|
51
|
+
uv run sphinx-build -W -b html docs/source docs/_build/html
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
Releases use calendar versions `YYYY.M.N`. The manual GitHub Actions
|
|
55
|
+
release workflow supports non-publishing previews and publishing from
|
|
56
|
+
`main`. See the workflow for its latest controls and validation. Publishing to PyPI\nis a separate, manually approved process described in\n[the release guide](docs/publishing.md).
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=65.5", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "CPax"
|
|
7
|
+
version = "2026.10.0"
|
|
8
|
+
description = "CPax: A JAX-based computational physics library."
|
|
9
|
+
requires-python = ">=3.12"
|
|
10
|
+
authors = [
|
|
11
|
+
{ name = "Jan Lukas Späh"}
|
|
12
|
+
]
|
|
13
|
+
license = { text = "MIT" }
|
|
14
|
+
readme = "README.md"
|
|
15
|
+
keywords = ["physics", "computational", "jax", "autograd", "ODE", "PDE"]
|
|
16
|
+
classifiers = [
|
|
17
|
+
"Programming Language :: Python :: 3.12",
|
|
18
|
+
"Programming Language :: Python :: 3.13",
|
|
19
|
+
"Programming Language :: Python :: 3.14",
|
|
20
|
+
"License :: OSI Approved :: MIT License"
|
|
21
|
+
]
|
|
22
|
+
|
|
23
|
+
dependencies = [
|
|
24
|
+
"diffrax>=0.7.2,<0.8",
|
|
25
|
+
"equinox>=0.13.8",
|
|
26
|
+
"jax>=0.11",
|
|
27
|
+
"numpy>=2.0",
|
|
28
|
+
"matplotlib>=3.7",
|
|
29
|
+
"pillow>=13.0.0",
|
|
30
|
+
]
|
|
31
|
+
|
|
32
|
+
[dependency-groups]
|
|
33
|
+
test = [
|
|
34
|
+
"bump-my-version",
|
|
35
|
+
"packaging>=24",
|
|
36
|
+
"pytest",
|
|
37
|
+
]
|
|
38
|
+
lint = [
|
|
39
|
+
"ruff==0.17.0",
|
|
40
|
+
"ty==0.0.86",
|
|
41
|
+
]
|
|
42
|
+
docs = [
|
|
43
|
+
"sphinx>=5.0",
|
|
44
|
+
"sphinx-rtd-theme",
|
|
45
|
+
"sphinx-autodoc-typehints",
|
|
46
|
+
"sphinx-autoapi",
|
|
47
|
+
"myst-parser",
|
|
48
|
+
]
|
|
49
|
+
dev = [
|
|
50
|
+
{ include-group = "test" },
|
|
51
|
+
{ include-group = "lint" },
|
|
52
|
+
{ include-group = "docs" },
|
|
53
|
+
]
|
|
54
|
+
|
|
55
|
+
[project.urls]
|
|
56
|
+
"Source" = "https://github.com/JaLuka98/CPax"
|
|
57
|
+
"Tracker" = "https://github.com/JaLuka98/CPax/issues"
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
[tool.bumpversion]
|
|
61
|
+
parse = '^(?P<release>[1-9][0-9]{3}\.(?:1[0-2]|0?[1-9]))\.(?P<patch>[0-9]+)$'
|
|
62
|
+
serialize = ["{release}.{patch}"]
|
|
63
|
+
commit = false
|
|
64
|
+
tag = false
|
|
65
|
+
tag_name = "v{new_version}"
|
|
66
|
+
|
|
67
|
+
[tool.bumpversion.parts.release]
|
|
68
|
+
calver_format = "{YYYY}.{MM}"
|
|
69
|
+
|
|
70
|
+
[[tool.bumpversion.files]]
|
|
71
|
+
filename = "pyproject.toml"
|
|
72
|
+
search = 'version = "{current_version}"'
|
|
73
|
+
replace = 'version = "{new_version}"'
|
|
74
|
+
[tool.uv]
|
|
75
|
+
required-version = ">=0.13.4"
|
|
76
|
+
[tool.ruff]
|
|
77
|
+
target-version = "py312"
|
|
78
|
+
line-length = 119
|
|
79
|
+
|
|
80
|
+
[tool.ruff.lint]
|
|
81
|
+
select = [
|
|
82
|
+
"E4", # Import conventions
|
|
83
|
+
"E7", # Statement conventions
|
|
84
|
+
"E9", # Syntax errors
|
|
85
|
+
"F", # Pyflakes
|
|
86
|
+
"I", # isort
|
|
87
|
+
"UP", # Modern Python syntax
|
|
88
|
+
"B", # flake8-bugbear
|
|
89
|
+
"SIM", # flake8-simplify
|
|
90
|
+
"C4", # flake8-comprehensions
|
|
91
|
+
"PIE", # flake8-pie
|
|
92
|
+
"RUF", # Ruff-specific correctness checks
|
|
93
|
+
"NPY", # NumPy correctness
|
|
94
|
+
"PERF", # Perflint
|
|
95
|
+
]
|
|
96
|
+
|
|
97
|
+
[tool.ruff.lint.per-file-ignores]
|
|
98
|
+
"src/cpax/__init__.py" = ["F401"]
|
|
99
|
+
"src/cpax/ode/__init__.py" = ["F401"]
|
|
100
|
+
|
|
101
|
+
[tool.ruff.format]
|
|
102
|
+
quote-style = "double"
|
|
103
|
+
indent-style = "space"
|
|
104
|
+
docstring-code-format = true
|
|
105
|
+
|
|
106
|
+
# Type-check the public library. Tests and examples can be added after the
|
|
107
|
+
# scientific API has been fully annotated.
|
|
108
|
+
[tool.ty.src]
|
|
109
|
+
include = ["src"]
|
cpax-2026.10.0/setup.cfg
ADDED
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: CPax
|
|
3
|
+
Version: 2026.10.0
|
|
4
|
+
Summary: CPax: A JAX-based computational physics library.
|
|
5
|
+
Author: Jan Lukas Späh
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Source, https://github.com/JaLuka98/CPax
|
|
8
|
+
Project-URL: Tracker, https://github.com/JaLuka98/CPax/issues
|
|
9
|
+
Keywords: physics,computational,jax,autograd,ODE,PDE
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Requires-Python: >=3.12
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
License-File: LICENSE
|
|
17
|
+
Requires-Dist: diffrax<0.8,>=0.7.2
|
|
18
|
+
Requires-Dist: equinox>=0.13.8
|
|
19
|
+
Requires-Dist: jax>=0.11
|
|
20
|
+
Requires-Dist: numpy>=2.0
|
|
21
|
+
Requires-Dist: matplotlib>=3.7
|
|
22
|
+
Requires-Dist: pillow>=13.0.0
|
|
23
|
+
Dynamic: license-file
|
|
24
|
+
|
|
25
|
+
# CPax
|
|
26
|
+
|
|
27
|
+
**CPax** is an open-source, JAX-based library for computational physics.
|
|
28
|
+
It separates physical **systems**, **problems**, **solvers** and
|
|
29
|
+
**solutions** in reusable workflows. The current implementation supports
|
|
30
|
+
continuous-time initial-value problems and classical mechanics; additional
|
|
31
|
+
physical domains are planned, not yet implemented.
|
|
32
|
+
|
|
33
|
+
**[Documentation](https://cpax.readthedocs.io/en/latest/)** ·
|
|
34
|
+
**[Source](https://github.com/JaLuka98/CPax)** ·
|
|
35
|
+
**[Issues](https://github.com/JaLuka98/CPax/issues)**
|
|
36
|
+
|
|
37
|
+
## Installation
|
|
38
|
+
|
|
39
|
+
Clone the repository, then run `uv sync` to create a reproducible
|
|
40
|
+
development environment. CPax currently requires Python 3.12 or newer.
|
|
41
|
+
|
|
42
|
+
## Quickstart
|
|
43
|
+
|
|
44
|
+
```python
|
|
45
|
+
import cpax
|
|
46
|
+
from cpax.mechanics import NewtonianSystem, PhaseSpaceState
|
|
47
|
+
|
|
48
|
+
system = NewtonianSystem(potential=lambda q: q**2 / 2, mass=1.0)
|
|
49
|
+
problem = cpax.IVP(system, PhaseSpaceState(1.0, 0.0), (0.0, 10.0))
|
|
50
|
+
solution = cpax.solve(problem, saveat=cpax.SaveAt.linspace(500))
|
|
51
|
+
|
|
52
|
+
print(solution.ys.q) # Position
|
|
53
|
+
print(solution.ys.p) # Canonical momentum
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
The same workflow supports the conservative Duffing oscillator, the
|
|
57
|
+
built-in `cpax.dynamics.LorenzSystem`, and multibody Newtonian gravity.
|
|
58
|
+
See the [examples](examples/) and documentation for complete scripts,
|
|
59
|
+
plots, and the animated Sun–Earth–Mars–Jupiter–comet showcase.
|
|
60
|
+
|
|
61
|
+
The default `solver="auto"` currently selects Diffrax Tsit5 for IVPs.
|
|
62
|
+
It does not infer stiffness or symplectic structure. A standalone
|
|
63
|
+
symplectic leapfrog integrator is preserved in `cpax.ode`.
|
|
64
|
+
|
|
65
|
+
## Development
|
|
66
|
+
|
|
67
|
+
CPax uses **uv**, **Ruff**, and **ty**:
|
|
68
|
+
|
|
69
|
+
```shell
|
|
70
|
+
uv sync
|
|
71
|
+
uv run pytest
|
|
72
|
+
uv run ruff check .
|
|
73
|
+
uv run ruff format --check .
|
|
74
|
+
uv run ty check
|
|
75
|
+
uv run sphinx-build -W -b html docs/source docs/_build/html
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
Releases use calendar versions `YYYY.M.N`. The manual GitHub Actions
|
|
79
|
+
release workflow supports non-publishing previews and publishing from
|
|
80
|
+
`main`. See the workflow for its latest controls and validation. Publishing to PyPI\nis a separate, manually approved process described in\n[the release guide](docs/publishing.md).
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
src/CPax.egg-info/PKG-INFO
|
|
5
|
+
src/CPax.egg-info/SOURCES.txt
|
|
6
|
+
src/CPax.egg-info/dependency_links.txt
|
|
7
|
+
src/CPax.egg-info/requires.txt
|
|
8
|
+
src/CPax.egg-info/top_level.txt
|
|
9
|
+
src/cpax/__init__.py
|
|
10
|
+
src/cpax/_backends/__init__.py
|
|
11
|
+
src/cpax/_backends/_diffrax.py
|
|
12
|
+
src/cpax/dynamics/__init__.py
|
|
13
|
+
src/cpax/dynamics/_lorenz.py
|
|
14
|
+
src/cpax/dynamics/_systems.py
|
|
15
|
+
src/cpax/mechanics/__init__.py
|
|
16
|
+
src/cpax/mechanics/_gravity.py
|
|
17
|
+
src/cpax/mechanics/_hamiltonian.py
|
|
18
|
+
src/cpax/mechanics/_newtonian.py
|
|
19
|
+
src/cpax/mechanics/_state.py
|
|
20
|
+
src/cpax/ode/__init__.py
|
|
21
|
+
src/cpax/ode/integrators.py
|
|
22
|
+
src/cpax/problems/__init__.py
|
|
23
|
+
src/cpax/problems/_ivp.py
|
|
24
|
+
src/cpax/results/__init__.py
|
|
25
|
+
src/cpax/results/_solution.py
|
|
26
|
+
src/cpax/solvers/__init__.py
|
|
27
|
+
src/cpax/solvers/_auto.py
|
|
28
|
+
src/cpax/solvers/_saveat.py
|
|
29
|
+
src/cpax/solvers/_solve.py
|
|
30
|
+
tests/test_dynamics.py
|
|
31
|
+
tests/test_examples.py
|
|
32
|
+
tests/test_gravity.py
|
|
33
|
+
tests/test_leapfrog.py
|
|
34
|
+
tests/test_lorenz.py
|
|
35
|
+
tests/test_mechanics.py
|
|
36
|
+
tests/test_mechanics_state.py
|
|
37
|
+
tests/test_problems.py
|
|
38
|
+
tests/test_pypi_tag.py
|
|
39
|
+
tests/test_results.py
|
|
40
|
+
tests/test_saveat.py
|
|
41
|
+
tests/test_solve.py
|
|
42
|
+
tests/test_solve_jax.py
|
|
43
|
+
tests/test_versioning.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
cpax
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Private numerical backend adapters. Not a public CPax namespace."""
|
|
@@ -0,0 +1,105 @@
|
|
|
1
|
+
"""Private Diffrax adapter for continuous-time initial-value problems."""
|
|
2
|
+
|
|
3
|
+
from typing import Any
|
|
4
|
+
|
|
5
|
+
import diffrax
|
|
6
|
+
import jax
|
|
7
|
+
import jax.numpy as jnp
|
|
8
|
+
import numpy as np
|
|
9
|
+
from jax.core import Tracer
|
|
10
|
+
|
|
11
|
+
from cpax.dynamics._systems import Time
|
|
12
|
+
from cpax.problems import IVP
|
|
13
|
+
from cpax.results import Solution, SolutionStatus
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def _resolve_saveat(policy: Any, t_span: tuple[Time, Time]) -> diffrax.SaveAt:
|
|
17
|
+
"""Resolve a CPax sampling policy against its problem's interval."""
|
|
18
|
+
t0, t1 = t_span
|
|
19
|
+
|
|
20
|
+
if policy.kind == "final":
|
|
21
|
+
return diffrax.SaveAt(t1=True)
|
|
22
|
+
|
|
23
|
+
if policy.kind == "linspace":
|
|
24
|
+
return diffrax.SaveAt(ts=jnp.linspace(t0, t1, policy.n))
|
|
25
|
+
|
|
26
|
+
if policy.kind == "at":
|
|
27
|
+
times = policy.times
|
|
28
|
+
if times is None:
|
|
29
|
+
raise ValueError("SaveAt.at requires an array of times.")
|
|
30
|
+
# Concrete checks provide an immediate, intelligible error. Diffrax
|
|
31
|
+
# itself handles the admissibility of traced times during solving.
|
|
32
|
+
if not any(isinstance(value, Tracer) for value in (t0, t1, times)):
|
|
33
|
+
interval_start = float(np.asarray(t0))
|
|
34
|
+
interval_end = float(np.asarray(t1))
|
|
35
|
+
concrete_times = np.asarray(times)
|
|
36
|
+
if np.any(concrete_times < interval_start) or np.any(concrete_times > interval_end):
|
|
37
|
+
raise ValueError("SaveAt.at times must lie within the IVP time interval.")
|
|
38
|
+
return diffrax.SaveAt(ts=times)
|
|
39
|
+
|
|
40
|
+
raise ValueError(f"Unsupported SaveAt policy: {policy.kind!r}.")
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def solve_diffrax(
|
|
44
|
+
problem: IVP,
|
|
45
|
+
*,
|
|
46
|
+
solver: diffrax.AbstractSolver,
|
|
47
|
+
solver_selection: str,
|
|
48
|
+
saveat: Any,
|
|
49
|
+
dt0: float | jax.Array | None,
|
|
50
|
+
stepsize_controller: diffrax.AbstractStepSizeController,
|
|
51
|
+
max_steps: int,
|
|
52
|
+
throw: bool,
|
|
53
|
+
) -> Solution:
|
|
54
|
+
"""Translate a CPax IVP into a Diffrax solve and back into CPax."""
|
|
55
|
+
|
|
56
|
+
def vector_field(
|
|
57
|
+
t: float | jax.Array | np.ndarray,
|
|
58
|
+
state: Any,
|
|
59
|
+
args: Any,
|
|
60
|
+
/,
|
|
61
|
+
) -> Any:
|
|
62
|
+
"""Adapt Diffrax's scalar-like time input to CPax's time interface."""
|
|
63
|
+
return problem.system.vector_field(jnp.asarray(t), state, args)
|
|
64
|
+
|
|
65
|
+
t0, t1 = problem.t_span
|
|
66
|
+
backend_solution = diffrax.diffeqsolve(
|
|
67
|
+
terms=diffrax.ODETerm(vector_field),
|
|
68
|
+
solver=solver,
|
|
69
|
+
t0=t0,
|
|
70
|
+
t1=t1,
|
|
71
|
+
dt0=dt0,
|
|
72
|
+
y0=problem.y0,
|
|
73
|
+
args=problem.args,
|
|
74
|
+
saveat=_resolve_saveat(saveat, problem.t_span),
|
|
75
|
+
stepsize_controller=stepsize_controller,
|
|
76
|
+
max_steps=max_steps,
|
|
77
|
+
throw=throw,
|
|
78
|
+
)
|
|
79
|
+
# Keep the numerical status JAX-traceable: never convert the result
|
|
80
|
+
# to a Python boolean/string inside a compiled integration.
|
|
81
|
+
# Diffrax returns an Equinox EnumerationItem, not a numeric JAX array.
|
|
82
|
+
# Its comparison operator yields a JAX-compatible scalar. Preserve the
|
|
83
|
+
# original result in metadata to retain precise backend failure reasons.
|
|
84
|
+
status = jnp.where(
|
|
85
|
+
backend_solution.result == diffrax.RESULTS.successful,
|
|
86
|
+
int(SolutionStatus.SUCCESS),
|
|
87
|
+
int(SolutionStatus.FAILED),
|
|
88
|
+
)
|
|
89
|
+
return Solution(
|
|
90
|
+
ts=backend_solution.ts,
|
|
91
|
+
ys=backend_solution.ys,
|
|
92
|
+
status=status,
|
|
93
|
+
stats=dict(backend_solution.stats),
|
|
94
|
+
metadata={
|
|
95
|
+
"backend": "diffrax",
|
|
96
|
+
"solver": type(solver).__name__,
|
|
97
|
+
"solver_selection": solver_selection,
|
|
98
|
+
"stepsize_controller": type(stepsize_controller).__name__,
|
|
99
|
+
"dt0": dt0,
|
|
100
|
+
"rtol": (stepsize_controller.rtol if isinstance(stepsize_controller, diffrax.PIDController) else None),
|
|
101
|
+
"atol": (stepsize_controller.atol if isinstance(stepsize_controller, diffrax.PIDController) else None),
|
|
102
|
+
"max_steps": max_steps,
|
|
103
|
+
"backend_result": backend_solution.result,
|
|
104
|
+
},
|
|
105
|
+
)
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
"""The autonomous Lorenz system."""
|
|
2
|
+
|
|
3
|
+
from typing import Any
|
|
4
|
+
|
|
5
|
+
import jax
|
|
6
|
+
import jax.numpy as jnp
|
|
7
|
+
|
|
8
|
+
from ._systems import DynamicalSystem
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class LorenzSystem(DynamicalSystem):
|
|
12
|
+
"""The three-dimensional Lorenz equations.
|
|
13
|
+
|
|
14
|
+
Parameters default to the classical chaotic regime. The state is a
|
|
15
|
+
floating-point array of shape (3,). Parameters are JAX arrays and may
|
|
16
|
+
be differentiated when used in a supported solver configuration.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
sigma: jax.Array
|
|
20
|
+
rho: jax.Array
|
|
21
|
+
beta: jax.Array
|
|
22
|
+
|
|
23
|
+
def __init__(
|
|
24
|
+
self,
|
|
25
|
+
sigma: float | jax.Array = 10.0,
|
|
26
|
+
rho: float | jax.Array = 28.0,
|
|
27
|
+
beta: float | jax.Array = 8.0 / 3.0,
|
|
28
|
+
) -> None:
|
|
29
|
+
for name, value in (("sigma", sigma), ("rho", rho), ("beta", beta)):
|
|
30
|
+
array = jnp.asarray(value)
|
|
31
|
+
if array.shape != () or array.dtype.kind not in "iuf":
|
|
32
|
+
raise TypeError(f"{name} must be a real numeric scalar.")
|
|
33
|
+
object.__setattr__(self, name, jnp.asarray(value, dtype=jnp.result_type(float, array)))
|
|
34
|
+
|
|
35
|
+
def vector_field(self, t: Any, state: jax.Array, args: Any = None) -> jax.Array:
|
|
36
|
+
"""Return the Lorenz vector field with the shape of `state`."""
|
|
37
|
+
del t
|
|
38
|
+
if args is not None:
|
|
39
|
+
raise ValueError("LorenzSystem does not support IVP.args; set system parameters instead.")
|
|
40
|
+
if not hasattr(state, "shape") or state.shape != (3,):
|
|
41
|
+
raise ValueError("LorenzSystem requires an array state of shape (3,).")
|
|
42
|
+
x, y, z = state
|
|
43
|
+
return jnp.stack((self.sigma * (y - x), x * (self.rho - z) - y, x * y - self.beta * z))
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
"""Shared continuous-time dynamical-system interfaces.
|
|
2
|
+
|
|
3
|
+
These systems specify equations of motion, not numerical solution methods.
|
|
4
|
+
They deliberately do not depend on physical domains or integration backends.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from abc import abstractmethod
|
|
8
|
+
from collections.abc import Callable
|
|
9
|
+
from typing import Any
|
|
10
|
+
|
|
11
|
+
import equinox as eqx
|
|
12
|
+
import jax
|
|
13
|
+
|
|
14
|
+
type Time = float | jax.Array
|
|
15
|
+
type State = Any # Array-valued PyTree; concrete state types are model-specific.
|
|
16
|
+
type Args = Any
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class DynamicalSystem(eqx.Module):
|
|
20
|
+
"""Base interface for continuous-time evolution.
|
|
21
|
+
|
|
22
|
+
Subclasses implement ``vector_field(t, state, args=None)`` and return
|
|
23
|
+
derivatives with the same PyTree structure and leaf shapes as ``state``.
|
|
24
|
+
Time is the first argument, including for autonomous systems.
|
|
25
|
+
"""
|
|
26
|
+
|
|
27
|
+
@abstractmethod
|
|
28
|
+
def vector_field(
|
|
29
|
+
self,
|
|
30
|
+
t: Time,
|
|
31
|
+
state: State,
|
|
32
|
+
args: Args = None,
|
|
33
|
+
) -> State:
|
|
34
|
+
"""Evaluate the instantaneous derivative of ``state`` at ``t``."""
|
|
35
|
+
raise NotImplementedError
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
class VectorFieldSystem(DynamicalSystem):
|
|
39
|
+
"""Wrap a user-supplied vector field as a continuous-time system.
|
|
40
|
+
|
|
41
|
+
``rhs`` has signature ``rhs(t, state, args)``. It may be a plain
|
|
42
|
+
Python function or a callable Equinox module with differentiable array
|
|
43
|
+
parameters. It is intentionally not marked as a static field.
|
|
44
|
+
"""
|
|
45
|
+
|
|
46
|
+
rhs: Callable[[Time, State, Args], State]
|
|
47
|
+
|
|
48
|
+
def vector_field(
|
|
49
|
+
self,
|
|
50
|
+
t: Time,
|
|
51
|
+
state: State,
|
|
52
|
+
args: Args = None,
|
|
53
|
+
) -> State:
|
|
54
|
+
"""Evaluate the wrapped vector field without modifying the state."""
|
|
55
|
+
return self.rhs(t, state, args)
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
"""Physics-first classical mechanics systems and potential helpers."""
|
|
2
|
+
|
|
3
|
+
from ._gravity import gravitational_potential
|
|
4
|
+
from ._hamiltonian import HamiltonianSystem
|
|
5
|
+
from ._newtonian import NewtonianSystem
|
|
6
|
+
from ._state import PhaseSpaceState
|
|
7
|
+
|
|
8
|
+
__all__ = ["HamiltonianSystem", "NewtonianSystem", "PhaseSpaceState", "gravitational_potential"]
|