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.
Files changed (45) hide show
  1. cpax-2026.10.0/LICENSE +21 -0
  2. cpax-2026.10.0/PKG-INFO +80 -0
  3. cpax-2026.10.0/README.md +56 -0
  4. cpax-2026.10.0/pyproject.toml +109 -0
  5. cpax-2026.10.0/setup.cfg +4 -0
  6. cpax-2026.10.0/src/CPax.egg-info/PKG-INFO +80 -0
  7. cpax-2026.10.0/src/CPax.egg-info/SOURCES.txt +43 -0
  8. cpax-2026.10.0/src/CPax.egg-info/dependency_links.txt +1 -0
  9. cpax-2026.10.0/src/CPax.egg-info/requires.txt +6 -0
  10. cpax-2026.10.0/src/CPax.egg-info/top_level.txt +1 -0
  11. cpax-2026.10.0/src/cpax/__init__.py +9 -0
  12. cpax-2026.10.0/src/cpax/_backends/__init__.py +1 -0
  13. cpax-2026.10.0/src/cpax/_backends/_diffrax.py +105 -0
  14. cpax-2026.10.0/src/cpax/dynamics/__init__.py +6 -0
  15. cpax-2026.10.0/src/cpax/dynamics/_lorenz.py +43 -0
  16. cpax-2026.10.0/src/cpax/dynamics/_systems.py +55 -0
  17. cpax-2026.10.0/src/cpax/mechanics/__init__.py +8 -0
  18. cpax-2026.10.0/src/cpax/mechanics/_gravity.py +38 -0
  19. cpax-2026.10.0/src/cpax/mechanics/_hamiltonian.py +49 -0
  20. cpax-2026.10.0/src/cpax/mechanics/_newtonian.py +107 -0
  21. cpax-2026.10.0/src/cpax/mechanics/_state.py +47 -0
  22. cpax-2026.10.0/src/cpax/ode/__init__.py +5 -0
  23. cpax-2026.10.0/src/cpax/ode/integrators.py +92 -0
  24. cpax-2026.10.0/src/cpax/problems/__init__.py +5 -0
  25. cpax-2026.10.0/src/cpax/problems/_ivp.py +114 -0
  26. cpax-2026.10.0/src/cpax/results/__init__.py +5 -0
  27. cpax-2026.10.0/src/cpax/results/_solution.py +85 -0
  28. cpax-2026.10.0/src/cpax/solvers/__init__.py +7 -0
  29. cpax-2026.10.0/src/cpax/solvers/_auto.py +12 -0
  30. cpax-2026.10.0/src/cpax/solvers/_saveat.py +58 -0
  31. cpax-2026.10.0/src/cpax/solvers/_solve.py +109 -0
  32. cpax-2026.10.0/tests/test_dynamics.py +71 -0
  33. cpax-2026.10.0/tests/test_examples.py +64 -0
  34. cpax-2026.10.0/tests/test_gravity.py +29 -0
  35. cpax-2026.10.0/tests/test_leapfrog.py +44 -0
  36. cpax-2026.10.0/tests/test_lorenz.py +44 -0
  37. cpax-2026.10.0/tests/test_mechanics.py +207 -0
  38. cpax-2026.10.0/tests/test_mechanics_state.py +61 -0
  39. cpax-2026.10.0/tests/test_problems.py +113 -0
  40. cpax-2026.10.0/tests/test_pypi_tag.py +87 -0
  41. cpax-2026.10.0/tests/test_results.py +93 -0
  42. cpax-2026.10.0/tests/test_saveat.py +116 -0
  43. cpax-2026.10.0/tests/test_solve.py +366 -0
  44. cpax-2026.10.0/tests/test_solve_jax.py +158 -0
  45. 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.
@@ -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,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"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -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,6 @@
1
+ diffrax<0.8,>=0.7.2
2
+ equinox>=0.13.8
3
+ jax>=0.11
4
+ numpy>=2.0
5
+ matplotlib>=3.7
6
+ pillow>=13.0.0
@@ -0,0 +1 @@
1
+ cpax
@@ -0,0 +1,9 @@
1
+ """CPax public physics-first API."""
2
+
3
+ from importlib.metadata import version as _distribution_version
4
+
5
+ __version__ = _distribution_version("CPax")
6
+
7
+ from . import dynamics, mechanics, ode
8
+ from .problems import IVP
9
+ from .solvers import SaveAt, solve
@@ -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,6 @@
1
+ """Continuous-time dynamical systems and user-defined vector fields."""
2
+
3
+ from ._lorenz import LorenzSystem
4
+ from ._systems import DynamicalSystem, VectorFieldSystem
5
+
6
+ __all__ = ["DynamicalSystem", "LorenzSystem", "VectorFieldSystem"]
@@ -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"]