torchode 0.1.8__tar.gz → 0.2.0__tar.gz

{torchode-0.1.8 → torchode-0.2.0}/CHANGELOG.md

@@ -5,6 +5,18 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## 0.2.0 - 2023-11-10
+
+### Changed
+
+- Removed compatibility with pytorch 1.x
+
+## 0.1.9 - 2023-08-29
+
+### Added
+
+- Allow passing `dt_0` for the forward and backward pass in the backsolve adjoint
+
 ## 0.1.8 - 2023-03-25
 
 ### Changed

{torchode-0.1.8 → torchode-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchode
-Version: 0.1.8
+Version: 0.2.0
 Summary: A parallel ODE solver for PyTorch
 Author-email: Marten Lienen <m.lienen@tum.de>
 Requires-Python: >= 3.8
@@ -12,9 +12,8 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Requires-Dist: sympy ~= 1.10
-Requires-Dist: torch >= 1.11
+Requires-Dist: torch >= 2.0
 Requires-Dist: torchtyping ~= 0.1.4
-Requires-Dist: functorch
 Requires-Dist: pytest-benchmark ; extra == "benchmark"
 Requires-Dist: pytest ; extra == "test"
 Requires-Dist: numpy ~= 1.22 ; extra == "test"
@@ -77,9 +76,6 @@ step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
 solver = to.AutoDiffAdjoint(step_method, step_size_controller)
 jit_solver = torch.compile(solver)
 
-# For pytorch versions < 2.0, use the older TorchScript compiler
-#jit_solver = torch.jit.script(solver)
-
 sol = jit_solver.solve(to.InitialValueProblem(y0=y0, t_eval=t_eval))
 print(sol.stats)
 # => {'n_f_evals': tensor([26, 26]), 'n_steps': tensor([4, 2]),

{torchode-0.1.8 → torchode-0.2.0}/README.md

@@ -53,9 +53,6 @@ step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
 solver = to.AutoDiffAdjoint(step_method, step_size_controller)
 jit_solver = torch.compile(solver)
 
-# For pytorch versions < 2.0, use the older TorchScript compiler
-#jit_solver = torch.jit.script(solver)
-
 sol = jit_solver.solve(to.InitialValueProblem(y0=y0, t_eval=t_eval))
 print(sol.stats)
 # => {'n_f_evals': tensor([26, 26]), 'n_steps': tensor([4, 2]),

torchode-0.2.0/docs/extra-stats.ipynb

@@ -0,0 +1,271 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "46327300-b000-4a93-b42a-7caef72e0a59",
+   "metadata": {},
+   "source": [
+    "# Logging Solver Information"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "45feb98e-1972-4600-a3a8-13cc1661a277",
+   "metadata": {},
+   "source": [
+    "One of the key features of torchode is that all components are replaceable and any components can log its own outputs (captured in a dictionary called `stats`). This means that you can inject your own code and log anything information that is relevant for your usecase. In this example, we will create a step size controller wrapper that logs the step times `t` and all accept decisions, i.e. if each step was accepted or rejected by the step size controller."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5f35a722-6c14-49cb-9768-954e2ef5eb61",
+   "metadata": {},
+   "source": [
+    "We begin by importing relevant modules and defining a generic model class."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c717d8d2-0564-424b-8b63-0bf74395c5fc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torchode as to\n",
+    "from torchode.step_size_controllers import StepSizeController\n",
+    "\n",
+    "torch.random.manual_seed(180819023);"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "98820150-ecf4-4ab4-88b7-fa1dcee681c5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Model(nn.Module):\n",
+    "    def __init__(self, n_features, n_hidden):\n",
+    "        super().__init__()\n",
+    "        self.layers = nn.Sequential(\n",
+    "            nn.Linear(n_features, n_hidden),\n",
+    "            nn.Softplus(),\n",
+    "            nn.Linear(n_hidden, n_hidden),\n",
+    "            nn.Softplus(),\n",
+    "            nn.Linear(n_hidden, n_features)\n",
+    "        )\n",
+    "    \n",
+    "    def forward(self, t, y):\n",
+    "        return self.layers(y)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9bd96be4-0a78-4e5f-9815-afda3953aa11",
+   "metadata": {},
+   "source": [
+    "Now we define the wrapper that will track the step size data. By deferring the actual functionality to another controller, we can re-use the existing controller implementations and focus on collecting the information that we care about, in this case the integration time points `t`, the step size `dt` and whether each step was accepted.\n",
+    "\n",
+    "To define this custom controller, we just have to satisfy the `StepSizeController` interface. For the actual functionality we defer to another controller instance. In `init` we additionally initialize fields in the statistics dictionary for the current solve to capture `t` and so on. The `adapt_step_size` method then records the information into those fields.\n",
+    "\n",
+    "Note that you could proceed in a similar way to track information about the stepping methods, e.g. dopri5, by defining a `SingleStepMethod`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "b8aed13a-a135-4ad7-aa37-f28a9837ded6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class StepSizeControllerTracker(StepSizeController):\n",
+    "    \"\"\"A wrapper that collects time step and step acceptance information.\"\"\"\n",
+    "\n",
+    "    def __init__(self, controller: StepSizeController):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.controller = controller\n",
+    "\n",
+    "    def init(self, term, problem, method_order: int, dt0, *, stats, args):\n",
+    "        stats[\"all_t\"] = []\n",
+    "        stats[\"all_dt\"] = []\n",
+    "        stats[\"all_accept\"] = []\n",
+    "\n",
+    "        return self.controller.init(\n",
+    "            term, problem, method_order, dt0, stats=stats, args=args\n",
+    "        )\n",
+    "\n",
+    "    def adapt_step_size(self, t0, dt, y0, step, state, stats):\n",
+    "        accept, dt_next, state, status = self.controller.adapt_step_size(\n",
+    "            t0, dt, y0, step, state, stats\n",
+    "        )\n",
+    "\n",
+    "        stats[\"all_t\"].append(t0)\n",
+    "        stats[\"all_dt\"].append(dt)\n",
+    "        stats[\"all_accept\"].append(accept)\n",
+    "\n",
+    "        return accept, dt_next, state, status\n",
+    "\n",
+    "    def merge_states(self, running, current, previous):\n",
+    "        return self.controller.merge_states(running, current, previous)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ebda2a83-81ee-48c8-93a4-dd6c8c73aa06",
+   "metadata": {},
+   "source": [
+    "Next, we construct a solver and wrap the step size controller with our tracker."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "3720e3bf-dc2a-4d91-a56f-b9ca93d77392",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "n_features = 5\n",
+    "batch_size = 3\n",
+    "\n",
+    "model = Model(n_features=n_features, n_hidden=32)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "d8a40685-5df1-4f56-b317-cd2e419cbee9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dev = torch.device(\"cpu\")\n",
+    "term = to.ODETerm(model)\n",
+    "step_method = to.Dopri5(term=term)\n",
+    "step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)\n",
+    "step_size_controller = StepSizeControllerTracker(step_size_controller)\n",
+    "adjoint = to.AutoDiffAdjoint(step_method, step_size_controller).to(dev)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fe6e8b4e-331e-405a-8757-91d179e9da75",
+   "metadata": {},
+   "source": [
+    "Finally, we generate some example data and evaluate the ODE defined by a randomly initialized model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b6a004e0-4327-4d6f-b19e-fb44b21c9584",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "t_eval = torch.tile(torch.linspace(0.0, 3.0, 10), (batch_size, 1))\n",
+    "problem = to.InitialValueProblem(y0=torch.zeros((batch_size, n_features)).to(dev), t_eval=t_eval.to(dev))\n",
+    "\n",
+    "sol = adjoint.solve(problem)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dacffa2b-fed5-4a60-a41b-1c5027b2854f",
+   "metadata": {},
+   "source": [
+    "In the end, we can inspect the statistics recorded in the solution object and see that our custom step size controller collected the data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "960abf1b-16c0-497d-bf89-ad999734039e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[0.0000e+00, 0.0000e+00, 0.0000e+00],\n",
+      "        [1.0000e-04, 1.0000e-04, 1.0000e-04],\n",
+      "        [1.1000e-03, 1.1000e-03, 1.1000e-03],\n",
+      "        [1.1100e-02, 1.1100e-02, 1.1100e-02],\n",
+      "        [1.1110e-01, 1.1110e-01, 1.1110e-01],\n",
+      "        [9.3798e-01, 9.3798e-01, 9.3798e-01]], grad_fn=<StackBackward0>)\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(torch.stack(sol.stats[\"all_t\"]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "6994ebb9-85e3-460f-8542-b0a00722c820",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[1.0000e-04, 1.0000e-04, 1.0000e-04],\n",
+      "        [1.0000e-03, 1.0000e-03, 1.0000e-03],\n",
+      "        [1.0000e-02, 1.0000e-02, 1.0000e-02],\n",
+      "        [1.0000e-01, 1.0000e-01, 1.0000e-01],\n",
+      "        [8.2688e-01, 8.2688e-01, 8.2688e-01],\n",
+      "        [2.0620e+00, 2.0620e+00, 2.0620e+00]], grad_fn=<StackBackward0>)\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(torch.stack(sol.stats[\"all_dt\"]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "7d19d776-d44e-4e8f-b370-2deceb8d1020",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[True, True, True],\n",
+      "        [True, True, True],\n",
+      "        [True, True, True],\n",
+      "        [True, True, True],\n",
+      "        [True, True, True],\n",
+      "        [True, True, True]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(torch.stack(sol.stats[\"all_accept\"]))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

torchode-0.2.0/docs/nd-data.md

@@ -0,0 +1,34 @@
+# Images, Graphs and N-d Tensors
+
+torchode is written with a sharp focus on its target problem: being a high-performance ODE
+solver for PyTorch. To keep the code base as simple as possible, it neglects some niceties
+such as arbitrary data dimensions. torchode expects your data to be exactly 2-dimensional,
+`batch x features`. If your samples are not 1-dimensional vectors, but for example images,
+you have to manage the shapes yourself. This means that you have to flatten the inputs
+when you hand them over to torchode and reshape them, when you pass them to a
+convolutional layer, let's say.
+
+```python
+# Example with a batch of RGB images
+b, c, w, h = 8, 3, 72, 72
+y0 = torch.randn(b, w, h)
+
+# Setup solver etc.
+# ...
+
+def f(t, y, args)
+    c, w, h = args
+
+    # Reshape 2D to B x C x W x H tensor and apply convolutional neural network
+    y = some_cnn(y.reshape((-1, c, w, h)))
+
+    # Flatten back to 2D when we return control back to torchode
+    return y.flatten(start_dim=1)
+
+term = to.ODETerm(f, with_args=True)
+problem = to.InitialValueProblem(y0=y0.flatten(start_dim=1), t_eval=..)
+sol = solver.solve(problem, args=(c, w, h))
+```
+
+If you are willing to contribute a small utility that does this shape flattening and
+reshaping automatically, we would be happy about your contribution!

{torchode-0.1.8 → torchode-0.2.0}/mkdocs.yml

@@ -14,5 +14,7 @@ nav:
   - step-size-controllers.md
   - jit.ipynb
   - extra-args.md
+  - extra-stats.ipynb
+  - nd-data.md
   - torchdiffeq.ipynb
   - comparison.md

{torchode-0.1.8 → torchode-0.2.0}/pyproject.toml

@@ -19,11 +19,8 @@ classifiers = [
 ]
 dependencies = [
   "sympy ~= 1.10",
-  "torch >= 1.11",
+  "torch >= 2.0",
   "torchtyping ~= 0.1.4",
-  # functorch has been integrated into pytorch and this is just a dummy dependency for
-  # compatibility with pytorch <= 1.13
-  "functorch",
 ]
 
 [project.optional-dependencies]

{torchode-0.1.8 → torchode-0.2.0}/torchode/__init__.py

@@ -1,6 +1,6 @@
 """A parallel ODE solver for PyTorch"""
 
-__version__ = "0.1.8"
+__version__ = "0.2.0"
 
 from .adjoints import AutoDiffAdjoint, BacksolveAdjoint, JointBacksolveAdjoint
 from .interface import register_method, solve_ivp

{torchode-0.1.8 → torchode-0.2.0}/torchode/adjoints.py

@@ -1,6 +1,5 @@
 from typing import Any, Dict, Optional
 
-import functorch
 import torch
 import torch.nn as nn
 
@@ -352,6 +351,8 @@ class BacksolveFunction(torch.autograd.Function):
         t_start,
         t_end,
         t_eval,
+        dt0,
+        backward_dt0,
         args,
         *adjoint_params,
     ):
@@ -361,10 +362,12 @@ class BacksolveFunction(torch.autograd.Function):
 
         problem = InitialValueProblem(y0, t_start, t_end, t_eval)
         with torch.no_grad():
-            solution = solver.solve(problem, dt0=None, args=args, term=term)
+            solution = solver.solve(problem, dt0=dt0, args=args, term=term)
 
         ctx.stats = solution.stats
-        ctx.save_for_backward(t_start, t_end, t_eval, solution.ys, *adjoint_params)
+        ctx.save_for_backward(
+            t_start, t_end, t_eval, backward_dt0, solution.ys, *adjoint_params
+        )
 
         return solution.ts, solution.ys, solution.stats, solution.status
 
@@ -373,7 +376,7 @@ class BacksolveFunction(torch.autograd.Function):
         aug_solver = ctx.aug_solver
         aug_term = ctx.aug_term
         args = ctx.args
-        t_start, t_end, t_eval, ys, *adjoint_params = ctx.saved_tensors
+        t_start, t_end, t_eval, dt0, ys, *adjoint_params = ctx.saved_tensors
 
         stats = ctx.stats
         stats.setdefault("backsolve", [])
@@ -391,7 +394,7 @@ class BacksolveFunction(torch.autograd.Function):
                     aug_state, t_start=t_end, t_end=t_start, t_eval=None
                 )
                 solution = aug_solver.solve(
-                    problem, dt0=None, args=(shapes, args), term=aug_term
+                    problem, dt0=dt0, args=(shapes, args), term=aug_term
                 )
                 aug_state = solution.ys[:, -1]
 
@@ -405,7 +408,7 @@ class BacksolveFunction(torch.autograd.Function):
                         t_eval=None,
                     )
                     solution = aug_solver.solve(
-                        problem, dt0=None, args=(shapes, args), term=aug_term
+                        problem, dt0=dt0, args=(shapes, args), term=aug_term
                     )
                     aug_state = solution.ys[:, -1]
                     aug_state[:, 1 : 1 + n_features] = ys[:, i - 1]
@@ -419,7 +422,20 @@ class BacksolveFunction(torch.autograd.Function):
         # Accumulate gradients over samples
         grad_params = [p.sum(dim=0) for p in grad_params]
 
-        return None, None, None, None, grad_y0, None, None, None, None, *grad_params
+        return (
+            None,
+            None,
+            None,
+            None,
+            grad_y0,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            *grad_params,
+        )
 
 
 class AugmentedDynamicsTerm(nn.Module):
@@ -432,22 +448,24 @@ class AugmentedDynamicsTerm(nn.Module):
         self.vmap_args_dims = vmap_args_dims
         self.vmap_randomness = vmap_randomness
 
-        self.func, self.params, self.buffers = functorch.make_functional_with_buffers(
-            self.term.f
-        )
+        f = self.term.f
+        orig_params = dict(f.named_parameters())
+        buffers = dict(f.named_buffers())
+        params_names = orig_params.keys()
 
         def vjp_single_sample(t_i, y_i, adj_y_i, arg_i):
             def wrapper(params, t_, y_):
+                args = (t_, y_)
                 if self.term.with_args:
-                    return self.func(params, self.buffers, t_, y_, arg_i)
-                else:
-                    return self.func(params, self.buffers, t_, y_)
+                    args = args + (arg_i,)
+                params_dict = {name: value for name, value in zip(params_names, params)}
+                return torch.func.functional_call(f, (params_dict, buffers), args)
 
-            dy, vjp = functorch.vjp(wrapper, self.params, t_i, y_i)
+            dy, vjp = torch.func.vjp(wrapper, tuple(orig_params.values()), t_i, y_i)
             vjp_params, vjp_t, vjp_y = vjp(-adj_y_i)
             return dy, vjp_t, vjp_y, vjp_params
 
-        self.vjp_vf = functorch.vmap(
+        self.vjp_vf = torch.func.vmap(
             vjp_single_sample,
             in_dims=(0, 0, 0, self.vmap_args_dims),
             randomness=self.vmap_randomness,
@@ -491,7 +509,10 @@ class BacksolveAdjoint(nn.Module):
         term: Optional[ODETerm] = None,
         dt0: Optional[TimeTensor] = None,
         args: Any = None,
+        backward_dt0: Optional[TimeTensor] = None,
     ) -> Solution:
+        if backward_dt0 is None and dt0 is not None:
+            backward_dt0 = -dt0
         ts, ys, stats, status = BacksolveFunction.apply(
             self.forward_adjoint,
             self.backward_adjoint,
@@ -501,6 +522,8 @@ class BacksolveAdjoint(nn.Module):
             problem.t_start,
             problem.t_end,
             problem.t_eval,
+            dt0,
+            backward_dt0,
             args,
             *list(self.term.parameters()),
         )
@@ -601,6 +624,7 @@ class JointBacksolveAdjoint(nn.Module):
         term: Optional[ODETerm] = None,
         dt0: Optional[TimeTensor] = None,
         args: Any = None,
+        backward_dt0: Optional[TimeTensor] = None,
     ) -> Solution:
         y0 = problem.y0.flatten()[None]
         t_start = problem.t_start[:1]
@@ -623,6 +647,8 @@ class JointBacksolveAdjoint(nn.Module):
         t_slope = all_slopes / all_slopes[0]
         t_intercept = problem.t_start - t_slope * problem.t_start[0]
 
+        if backward_dt0 is None and dt0 is not None:
+            backward_dt0 = -dt0
         _, ys, stats, status = BacksolveFunction.apply(
             self.forward_loop,
             self.backward_loop,
@@ -632,6 +658,8 @@ class JointBacksolveAdjoint(nn.Module):
             t_start,
             t_end,
             t_eval,
+            dt0,
+            backward_dt0,
             (problem.batch_size, t_intercept, t_slope, args),
             *list(self.term.parameters()),
         )