Code update

PiperOrigin-RevId: 554634099
google-research · Aug 7, 2023 · ad5b930 · ad5b930
1 parent 8a5f5f2
commit ad5b930
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/swirl_dynamics/lib/solvers/ode.py b/swirl_dynamics/lib/solvers/ode.py
@@ -103,6 +103,16 @@ def step(
     return x0 + dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6
 
 
+class OneStepDirect(ScanOdeSolver):
+  """Solver that directly returns function output as next time step."""
+
+  def step(
+      self, func: OdeDynamics, x0: Array, t0: Array, dt: Array, params: PyTree
+  ) -> Array:
+    """Performs a single prediction step: `x_{n+1} = f(x_n, t_n)`."""
+    return func(x0, t0, params)
+
+
 @flax.struct.dataclass
 class DoPri45:
   """Adaptive-step Dormand-Prince solver."""