0.3.1 (#84)

* adjust nbook

* s to t

* Initial work for symplectic integrators

* odeint_symplectic added to API

* hypersolver base

* tutorial hypersolvers

* new tutorial, better readme

* better readme

Author: Zymrael

README.md

@@ -7,7 +7,7 @@
 
 ## A PyTorch based library for all things **neural differential equations**. Maintained by [DiffEqML](https://github.com/DiffEqML).
 
-![license](https://img.shields.io/badge/license-Apache%202.0-red.svg?)
+![license](https://img.shields.io/badge/license-Apache%202.0-blue.svg?)
 ![CI](https://github.com/DiffEqML/torchdyn/actions/workflows/os-coverage.yml/badge.svg)
 [![Slack](https://img.shields.io/badge/slack-chat-blue.svg?logo=slack)](https://join.slack.com/t/diffeqml/shared_invite/zt-gq3jjj5x-LuHSB4m4gc9MsnvoF1UB6A)
 [![codecov](https://codecov.io/gh/DiffEqML/torchdyn/branch/master/graph/badge.svg)](https://codecov.io/gh/DiffEqML/torchdyn)
@@ -40,17 +40,12 @@ Contribute to the library with your benchmark and model variants! No need to rei
 
 `pip install torchdyn`
 
-* NOTE: temporarily requires additional manual installation of `torchsde`: 
-
-`pip install git+https://github.com/google-research/torchsde.git`
 
 **Bleeding-edge** version:
 
-`git clone https://github.com/DiffEqML/torchdyn.git`
-
-`cd torchdyn`
+`git clone https://github.com/DiffEqML/torchdyn.git && cd torchdyn && python setup.py install`
 
-`python setup.py install`
+Don't forget to install in your environment of choice if necessary. We offer an automated method for setting up your `torchdyn` environment designed specifically for contributors or those planning to tinker with the internals. Check `Contributing` below for more details.
 
 ## Documentation
 Check our [wiki](https://torchdyn.readthedocs.io/) for a full description of available features.
@@ -64,7 +59,7 @@ Interest in the blend of differential equations, deep learning and dynamical sys
 </p>
 
 
-By providing a centralized, easy-to-access collection of model templates, tutorial and application notebooks, we hope to speed-up research in this area and ultimately establish neural differential equations and implicit models into an effective tool for control, system identification and general machine learning tasks. W
+By providing a centralized, easy-to-access collection of model templates, tutorial and application notebooks, we hope to speed-up research in this area and ultimately establish neural differential equations and implicit models into an effective tool for control, system identification and general machine learning tasks.
 
 #### Dependencies
 `torchdyn` leverages modern PyTorch best practices and handles training with `pytorch-lightning` [[6](https://github.com/PyTorchLightning/pytorch-lightning)]. We build Graph Neural ODEs utilizing the Graph Neural Networks (GNNs) API of `dgl` [[7](https://www.dgl.ai/)]. For a complete list of references, check `pyproject.toml`. We offer a complete suite of ODE solvers and sensitivity methods, extending the functionality offered by `torchdiffeq` [[1](https://arxiv.org/abs/1806.07366)]. We have light dependencies on `torchsde` [[7](https://arxiv.org/abs/2001.01328)] and `torchcde` [[8](https://arxiv.org/abs/2005.08926)].

benchmarks/numerics/bench_hypersolver.ipynb

@@ -2,21 +2,9 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
-   "outputs": [
-    {
-     "ename": "ModuleNotFoundError",
-     "evalue": "No module named 'torchdyn.numerics'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
-      "\u001b[0;32m<ipython-input-2-11c594d46672>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mmatplotlib\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpyplot\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mplt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mtorchdyn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumerics\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mEuler\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mRungeKutta4\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mTsitouras45\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mDormandPrince45\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtorchdyn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumerics\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0modeint\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
-      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'torchdyn.numerics'"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import torchdyn\n",
     "import torch\n",
@@ -471,9 +459,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "torchdyn",
    "language": "python",
-   "name": "python3"
+   "name": "torchdyn"
   },
   "language_info": {
    "codemirror_mode": {

benchmarks/numerics/bench_odeint.ipynb

@@ -30,7 +30,7 @@ papermill = "*"
 poethepoet = "^0.10.0"
 
 [tool.poe.tasks]
-force-cuda11 = "python -m pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html"
+force-cuda11 = "python -m pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio===0.9.0 -f https://download.pytorch.org/whl/torch_stable.html"
 
 [build-system]
 build-backend = "poetry.masonry.api"

benchmarks/numerics/bench_symplectic.ipynb

@@ -11,8 +11,24 @@
 
 
 class ODEProblem(nn.Module):
-    def __init__(self, vector_field, solver:Union[str, nn.Module], atol:float=1e-4, rtol:float=1e-4, sensitivity='autograd',
+    def __init__(self, vector_field, solver:Union[str, nn.Module], order:int=1, atol:float=1e-4, rtol:float=1e-4, sensitivity='autograd',
                  solver_adjoint:Union[str, nn.Module, None] = None, atol_adjoint:float=1e-6, rtol_adjoint:float=1e-6, seminorm:bool=False):
+        """An ODE Problem coupling a given vector field with solver and sensitivity algorithm to compute gradients w.r.t different quantities.
+
+        Args:
+            vector_field ([Callable]): the vector field, called with `vector_field(t, x)` for `vector_field(x)`. 
+                                       In the second case, the Callable is automatically wrapped for consistency
+            solver (Union[str, nn.Module]): [description]
+            order (int, optional): [description]. Defaults to 1.
+            atol (float, optional): [description]. Defaults to 1e-4.
+            rtol (float, optional): [description]. Defaults to 1e-4.
+            sensitivity (str, optional): [description]. Defaults to 'autograd'.
+            solver_adjoint (Union[str, nn.Module, None], optional): [description]. Defaults to None.
+            atol_adjoint (float, optional): [description]. Defaults to 1e-6.
+            rtol_adjoint (float, optional): [description]. Defaults to 1e-6.
+            seminorm (bool, optional): Indicates whether the a seminorm should be used for error estimation during adjoint backsolves. Defaults to False.
+        
+        """
         super().__init__()
         # instantiate solver at initialization
         if type(solver) == str: 
@@ -38,7 +54,7 @@ def __init__(self, vector_field, solver:Union[str, nn.Module], atol:float=1e-4,
                 vector_field = DEFuncBase(vector_field, has_time_arg=False)   
             else: vector_field = DEFuncBase(vector_field, has_time_arg=True) 
 
-        self.vf, self.sensalg = vector_field, sensitivity
+        self.vf, self.order, self.sensalg = vector_field, order, sensitivity
         if len(tuple(self.vf.parameters())) > 0:
             self.vf_params = torch.cat([p.contiguous().flatten() for p in self.vf.parameters()])
         else:
@@ -69,6 +85,16 @@ def forward(self, x:Tensor, t_span:Tensor):
 
 class MultipleShootingProblem(nn.Module):
     def __init__(self, solver:str, vector_field, sensalg='autograd'):
+        """[summary]
+
+        Args:
+            solver (str): [description]
+            vector_field ([type]): [description]
+            sensalg (str, optional): [description]. Defaults to 'autograd'.
+
+        Returns:
+            [type]: [description]
+        """
         super().__init__()
         #
         self.solver
@@ -98,3 +124,4 @@ def forward(self, x0, t_span, atol=1e-4, rtol=1e-4, t_eval=[]):
         x0, t_span = prep_input(x0, t_span)
         t_eval, sol = self.odefunc(self.vf_params, x0, t_span, t_eval)
         return t_eval, sol
+        

pyproject.toml

@@ -1,24 +1,22 @@
 ###  Goals of `torchdyn`
-Our aim with  `torchdyn` aims is to provide a unified, flexible API  to the most recent advances in continuous and implicit neural networks. Examples include neural differential equations variants, e.g.
+Our aim with  `torchdyn` aims is to provide a unified, flexible API to aid in the implementation of recent advances in continuous and implicit learning. Some models already implemented, either here under `torchdyn.models` or in the tutorials, are:
 
 * Neural Ordinary Differential Equations (Neural ODE) [[1](https://arxiv.org/abs/1806.07366)]
-* Neural Stochastic Differential Equations (Neural SDE) [[7](https://arxiv.org/abs/1905.09883),[8](https://arxiv.org/abs/1906.02355)]
-* Graph Neural ODEs [[9](https://arxiv.org/abs/1911.07532)]
-* Hamiltonian Neural Networks [[10](https://arxiv.org/abs/1906.01563)]
+* Galerkin Neural ODE [[2](https://arxiv.org/abs/2002.08071)]
+* Neural Stochastic Differential Equations (Neural SDE) [[3](https://arxiv.org/abs/1905.09883),[4](https://arxiv.org/abs/1906.02355)]
+* Graph Neural ODEs [[5](https://arxiv.org/abs/1911.07532)]
+* Hamiltonian Neural Networks [[6](https://arxiv.org/abs/1906.01563)]
 
-Depth-variant versions,
-* ANODEv2 [[11](https://arxiv.org/abs/1906.04596)]
-* Galerkin Neural ODE [[12](https://arxiv.org/abs/2002.08071)]
+Recurrent or "hybrid" versions for sequences
+* ODE-RNN [[7](https://arxiv.org/abs/1907.03907)]
 
-Recurrent or "hybrid" versions
-* ODE-RNN [[13](https://arxiv.org/abs/1907.03907)]
-* GRU-ODE-Bayes [[14](https://arxiv.org/abs/1905.12374)]
+Neural numerical methods
+* Hypersolvers [[12](https://arxiv.org/pdf/2007.09601.pdf)]
 
 Augmentation strategies to relieve neural differential equations of their expressivity limitations and reduce the computational burden of the numerical solver
-* ANODE (0-augmentation) [[15](https://arxiv.org/abs/1904.01681)]
-* Input-layer augmentation [[16](https://arxiv.org/abs/2002.08071)]
-* Higher-order augmentation [[17](https://arxiv.org/abs/2002.08071)]
+* ANODE (0-augmentation) [[8](https://arxiv.org/abs/1904.01681)]
+* Input-layer augmentation [[9](https://arxiv.org/abs/2002.08071)]
+* Higher-order augmentation [[10](https://arxiv.org/abs/2002.08071)]
 
-Alternative or modified adjoint training techniques
-* Integral loss adjoint [[18](https://arxiv.org/abs/2003.08063)]
-* Checkpointed adjoint [[19](https://arxiv.org/abs/1902.10298)]
+Various sensitivity algorithms / variants
+* Integral loss adjoint [[11](https://arxiv.org/abs/2003.08063)]

torchdyn/core/problems.py

@@ -1,5 +1,11 @@
+import math
 import torch
 import torch.nn as nn
+from torch.distributions import Normal, kl_divergence
+import pytorch_lightning as pl
+import torchsde
+
+from torchdyn.models import LSDEFunc
 
 
 class HybridNeuralDE(nn.Module):
@@ -42,16 +48,6 @@ def _jump_latent_cell(self, *args):
         return self.jump(x_t, (h, c))
 
 
-
-import math
-import torch
-import torch.nn as nn
-from torch.distributions import Normal, kl_divergence
-import pytorch_lightning as pl
-import torchsde
-
-from torchdyn.models import LSDEFunc
-
 class LatentNeuralSDE(NeuralSDE, pl.LightningModule): # pragma: no cover
     def __init__(self, post_drift, diffusion, prior_drift, sigma, theta, mu, options,
                  noise_type, order, sensitivity, s_span, solver, atol, rtol, intloss):
@@ -74,9 +70,14 @@ def __init__(self, post_drift, diffusion, prior_drift, sigma, theta, mu, options
         self.qy0_logvar = nn.Parameter(torch.tensor([[logvar]]), requires_grad=True)
 
     def forward(self, eps: torch.Tensor, s_span=None):
-        """
-        :param s_span: (optional) Series span -- can pass extended version for additional regularization
-        :param eps: Noise sample
+        """[summary]
+
+        Args:
+            eps (torch.Tensor): [description]
+            s_span ([type], optional): [description]. Defaults to None.
+
+        Returns:
+            [type]: [description]
         """
 
         eps = eps.to(self.qy0_std)
@@ -86,15 +87,12 @@ def forward(self, eps: torch.Tensor, s_span=None):
         py0 = Normal(loc=self.py0_mean, scale=self.py0_std)
         logqp0 = kl_divergence(qy0, py0).sum(1).mean(0)  # KL(time=0).
 
-        # Expand s_span to penalize out-of-datasets region and spread uncertainty -- moved
-        # s_span_ext = torch.cat((torch.tensor([0.0]), self.s_span.to('cpu'), torch.tensor([2.0])))
-
         if s_span is not None:
             s_span_ext = s_span
         else:
             s_span_ext = self.s_span.cpu()
 
-        zs, logqp = sdeint(sde=self.defunc, x0=x0, s_span=s_span_ext,
+        zs, logqp = torchsde.sdeint(sde=self.defunc, x0=x0, s_span=s_span_ext,
                            rtol=self.rtol, atol=self.atol, logqp=True, options=self.options,
                            adaptive=self.adaptive, method=self.solver)
 
@@ -104,24 +102,34 @@ def forward(self, eps: torch.Tensor, s_span=None):
         return zs, log_ratio
 
     def sample_p(self, vis_span, n_sim, eps=None, bm=None, dt=0.01):
-        """
-        :param vis_span:
-        :param n_sim:
-        :param eps:
-        :param bm:
-        :param dt:
+        """[summary]
+
+        Args:
+            vis_span ([type]): [description]
+            n_sim ([type]): [description]
+            eps ([type], optional): [description]. Defaults to None.
+            bm ([type], optional): [description]. Defaults to None.
+            dt (float, optional): [description]. Defaults to 0.01.
+
+        Returns:
+            [type]: [description]
         """
         eps = torch.randn(n_sim, 1).to(self.py0_mean).to(self.device) if eps is None else eps
         y0 = self.py0_mean + eps.to(self.device) * self.py0_std
         return torchsde.sdeint(self.defunc, y0, vis_span, bm=bm, method='srk', dt=dt, names={'drift': 'h'})
 
     def sample_q(self, vis_span, n_sim, eps=None, bm=None, dt=0.01):
-        """
-        :param vis_span:
-        :param n_sim:
-        :param eps:
-        :param bm:
-        :param dt:
+        """[summary]
+
+        Args:
+            vis_span ([type]): [description]
+            n_sim ([type]): [description]
+            eps ([type], optional): [description]. Defaults to None.
+            bm ([type], optional): [description]. Defaults to None.
+            dt (float, optional): [description]. Defaults to 0.01.
+
+        Returns:
+            [type]: [description]
         """
         eps = torch.randn(n_sim, 1).to(self.qy0_mean) if eps is None else eps
         y0 = self.qy0_mean + eps.to(self.device) * self.qy0_std

torchdyn/models/README.md

@@ -174,10 +174,10 @@ def __init__(self, bias=True, expfunc=Fourier(5), dilation=True, shift=True):
     def reset_parameters(self):
         torch.nn.init.zeros_(self.coeffs)
 
-    def calculate_weights(self, s):
-        "Expands `s` following the chosen eigenbasis"
+    def calculate_weights(self, t):
+        "Expands `t` following the chosen eigenbasis"
         n_range = torch.linspace(0, self.deg, self.deg).to(self.coeffs.device)
-        basis = self.expfunc(n_range, s*self.dilation.to(self.coeffs.device) + self.shift.to(self.coeffs.device))
+        basis = self.expfunc(n_range, t*self.dilation.to(self.coeffs.device) + self.shift.to(self.coeffs.device))
         B = []  
         for i in range(self.n_eig):
             Bin = torch.eye(self.deg).to(self.coeffs.device)

torchdyn/models/hybrid.py

@@ -10,8 +10,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from torchdyn.numerics.solvers import Euler, RungeKutta4, Tsitouras45, DormandPrince45
-from torchdyn.numerics.odeint import odeint
+from torchdyn.numerics.solvers import Euler, RungeKutta4, Tsitouras45, DormandPrince45, AsynchronousLeapfrog
+from torchdyn.numerics.hypersolvers import HyperEuler
+from torchdyn.numerics.odeint import odeint, odeint_symplectic
 
 
-__all__ =   ['odeint', 'Euler', 'RungeKutta4', 'DormandPrince45', 'Tsitouras45']
+__all__ =   ['odeint', 'odeint_symplectic', 'Euler', 'RungeKutta4', 'DormandPrince45', 'Tsitouras45',
+            'AsynchronousLeapfrog', 'HyperEuler']
+            

torchdyn/nn/galerkin.py

@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+from torchdyn.numerics.solvers import  Euler, RungeKutta4
+
+class HyperEuler(Euler):
+    def __init__(self, hypernet, dtype=torch.float32):
+        super().__init__(dtype)
+        self.hypernet = hypernet
+        self.stepping_class = 'fixed'
+
+    def step(self, f, x, t, dt, k1=None):
+        _, _, x_sol = super().step(f, x, t, dt, k1)
+        return None, None, x_sol + dt**2 * self.hypernet(t, x)
+
+class HyperRungeKutta4(RungeKutta4):
+    def __init__(self, hypernet, dtype=torch.float32):
+        super().__init__(dtype)
+        self.hypernet = hypernet
+
+    def step(self, f, x, t, dt, k1=None):
+        _, _, x_sol = super().step(f, x, t, dt, k1)
+        return None, None, x_sol + dt**5 * self.hypernet(t, x)