5 rk45 jax solver returning nan solution

README.md

@@ -14,6 +14,34 @@ The following command will install `safe-autonomy-simulation` into your local en
 pip install safe-autonomy-simulation
 ```
 
+### Installing with JAX (experimental)
+
+The `safe-autonomy-simulation` package supports numerical computation acceleration via the [JAX](https://jax.readthedocs.io/en/latest/index.html) library. This is an experimental feature.
+
+If you want to use JAX acceleration you can install JAX with `safe-autonomy-simulation` using `pip`:
+
+```shell
+pip install safe-autonomy-simulation[jax]
+```
+
+This will install JAX into your local environment for CPU acceleration. JAX also provides GPU acceleration for numerical computing. If you'd like this feature the easiest way to install it is using `pip`:
+
+```shell
+pip install jax[cuda12]
+```
+
+For more information on installing JAX see the [official documentation](https://jax.readthedocs.io/en/latest/installation.html).
+
+JAX can be enabled on a per-class basis for classes that support it. As JAX is an experimental feature it is not enabled by default.
+
+You can additionally force `safe-autonomy-simulation` to not use JAX at all by setting the following environment variable:
+
+```shell
+USE_JAX="0"
+```
+
+This will turn off JAX globally for `safe-autonomy-simulation` even if a class requests it.
+
 ### Installing from source
 
 Alternatively, `safe-autonomy-simulation` can be installed from source using any of the following methods. Again, it is recommended to install this package in an isolated virtual environment. The following sections describe how to install `safe-autonomy-simulation` from source in an isolated virtual environment using `poetry`, `conda`, and `pip + virtualenv`.

pyproject.toml

@@ -26,9 +26,10 @@ scipy = "^1.13.0"
 scikit-learn = "^1.5.0"
 typing-extensions = "^4.12.2"
 jax = {extras = ["cpu"], version = "^0.4.30", optional = true}
+diffrax = {version = "^0.6.0", optional = true}
 
 [tool.poetry.extras]
-jax = ["jax"]
+jax = ["jax", "diffrax"]
 
 [tool.poetry.group.test]
 optional = true

safe_autonomy_simulation/__init__.py

@@ -3,6 +3,7 @@
 continuous simulations in python.
 """
 
+from safe_autonomy_simulation.opts import jax_available
 from safe_autonomy_simulation import (
     dynamics,
     entities,
@@ -32,4 +33,10 @@
     "controls",
     "sims",
     "utils",
+    "jax_available",
 ]
+
+
+if jax_available():
+    from safe_autonomy_simulation import jax
+    __all__.append("jax")

safe_autonomy_simulation/controls/control_queue.py

@@ -3,11 +3,7 @@
 import typing
 import queue
 import warnings
-
-try:
-    import jax.numpy as np
-except ImportError:
-    import numpy as np
+import numpy as np
 
 
 class ControlQueue:

safe_autonomy_simulation/controls/no_control.py

@@ -1,12 +1,8 @@
 """A default control queue with an empty default control vector."""
 
+import numpy as np
 import safe_autonomy_simulation.controls.control_queue as c
 
-try:
-    import jax.numpy as np
-except ImportError:
-    import numpy as np
-
 
 class NoControl(c.ControlQueue):
     """A control queue with an empty default control vector.

safe_autonomy_simulation/dynamics/__init__.py

@@ -15,11 +15,13 @@
     ControlAffineODEDynamics,
     LinearODEDynamics,
 )
+import safe_autonomy_simulation.dynamics.utils as utils
 
 __all__ = [
     "Dynamics",  # base dynamics model
     "PassThroughDynamics",  # dynamics model that passes through control inputs
     "ODEDynamics",  # base dynamics model that integrates ordinary differential equations
     "ControlAffineODEDynamics",  # control-affine ODE dynamics model
     "LinearODEDynamics",  # linear ODE dynamics model
+    "utils",  # utility functions for dynamics models
 ]

safe_autonomy_simulation/dynamics/dynamics.py

@@ -1,11 +1,12 @@
 """Base class for state transition dynamics models of entities in the simulation environment."""
 
 import typing
+import numpy as np
+import safe_autonomy_simulation
+from safe_autonomy_simulation.utils import cast_jax
 
-try:
-    import jax.numpy as np
-except ImportError:
-    import numpy as np
+if safe_autonomy_simulation.jax_available():
+    import jax.numpy as jnp
 
 
 class Dynamics:
@@ -25,19 +26,26 @@ class Dynamics:
         When a float, represents single limit applied to entire state vector.
         When an ndarray, each element represents the limit to the corresponding state vector element.
         By default, np.inf
+    use_jax : bool, optional
+        EXPERIMENTAL: Use JAX to accelerate state transition computation, by default False.
     """
 
     def __init__(
         self,
         state_min: typing.Union[float, np.ndarray] = -np.inf,
         state_max: typing.Union[float, np.ndarray] = np.inf,
+        use_jax: bool = False,
     ):
-        self.np = np
-        self.state_min = state_min
-        self.state_max = state_max
+        self.use_jax = safe_autonomy_simulation.jax_available and use_jax
+        self.np = jnp if self.use_jax else np
+        self.state_min = self.np.copy(state_min)
+        self.state_max = self.np.copy(state_max)
 
     def step(
-        self, step_size: float, state: np.ndarray, control: np.ndarray
+        self,
+        step_size: float,
+        state: np.ndarray,
+        control: np.ndarray,
     ) -> typing.Tuple[np.ndarray, np.ndarray]:
         """
         Computes the dynamics state transition from the current state and control input.
@@ -56,12 +64,22 @@ def step(
         Tuple[np.ndarray, np.ndarray]
             Tuple of the system's next state and the state's instantaneous time derivative at the end of the step
         """
+        state = cast_jax(state, use_jax=self.use_jax)
+        control = cast_jax(control, use_jax=self.use_jax)
         next_state, state_dot = self._step(step_size, state, control)
-        next_state = self.np.clip(next_state, self.state_min, self.state_max)
+
+        next_state = self.np.clip(
+            cast_jax(next_state, use_jax=self.use_jax),
+            self.state_min,
+            self.state_max,
+        )
         return next_state, state_dot
 
     def _step(
-        self, step_size: float, state: np.ndarray, control: np.ndarray
+        self,
+        step_size: float,
+        state: np.ndarray,
+        control: np.ndarray,
     ) -> typing.Tuple[np.ndarray, np.ndarray]:
         """Computes the next state and state derivative of the system