`PotentialODE.reference.ipynb`#

[1]:

%load_ext nb_black
# pip install neural-diffeqs

import neural_diffeqs

print(f"Version: {neural_diffeqs.__version__}")
import torch

Version: 0.3.2

Default `PotentialODE`#

As in the NeuralODE and NeuralSDE, the only required parameter is:

state_size

[2]:

ODE = neural_diffeqs.PotentialODE(state_size=20, mu_hidden=[64, 64])
print(ODE)

PotentialODE(
  (mu): TorchNet(
    (hidden_1): Sequential(
      (linear): Linear(in_features=20, out_features=64, bias=True)
      (activation): LeakyReLU(negative_slope=0.01)
    )
    (hidden_2): Sequential(
      (linear): Linear(in_features=64, out_features=64, bias=True)
      (activation): LeakyReLU(negative_slope=0.01)
    )
    (output): Sequential(
      (linear): Linear(in_features=64, out_features=1, bias=False)
    )
  )
)

Notice that, the output layer of the mu function (drift network) contains only a single feature, without bias (by default):

(output): Sequential(
(linear): Linear(in_features=512, out_features=1, bias=False)
)

[3]:

# 5 samples x 20 dim
y = torch.randn([5, 20])
print(y)

tensor([[-0.3228,  0.1208, -0.4852, -1.5084, -1.5721, -0.3639,  0.7440, -0.3686,
          0.0311,  0.0771, -0.2299,  0.6744,  1.0946,  1.8330, -0.0296,  0.0648,
          1.1006, -0.2764,  1.1437,  2.3773],
        [-0.3958,  0.5138,  0.0432, -0.0932,  2.2951, -1.3561,  0.4870, -1.4136,
          1.0158, -1.5997, -2.5283,  0.5416, -0.0307, -1.3237, -0.1767,  1.9054,
          1.5694, -0.5196,  0.3701,  1.1737],
        [ 0.4174,  0.5138, -1.1092, -0.0051, -1.0940, -0.1167,  1.0324, -1.6913,
         -0.0920, -0.4754,  0.2900, -0.9673, -2.0961, -0.1045,  1.0242, -0.9995,
          1.1757, -0.4290,  1.8097,  0.9739],
        [ 1.1494,  0.0153,  0.2669, -0.2239, -0.3777,  0.6096,  0.1383, -0.6071,
         -0.9103,  0.7419, -0.2132, -0.0099, -1.1543, -0.0944,  1.1588,  0.7124,
         -2.2216, -0.0356,  0.0506,  0.1203],
        [ 0.9027, -0.2452,  0.3295,  1.3359,  0.5843,  0.5813,  0.2054,  0.1454,
          1.4566,  0.4892, -0.5222, -0.6388,  0.4159,  0.8533, -0.8192, -1.3142,
          1.6547,  0.9800,  0.4647, -0.4676]])

PotentialODE.f and PotentialODE.drift are identical. Under the hood, the following occurs:

def drift(self, y: torch.Tensor) -> torch.Tensor:

        y = y.requires_grad_()
        ψ = self._potential(y)
        return self._gradient(ψ, y) * self._coef_drift

Wherein PotentialODE._potential(y) computes ψ = PotentialODE.mu(y). ψ is of shape: [y.shape[0], 1] or [n_samples x 1]. While the regular NeuralODE computes y(t) = net(y), directly, the next step for a PotentialODE occurs in PotentialODE._gradient(ψ, y):

y_hat = torch.autograd.grad(ψ, y, torch.ones_like(ψ), create_graph=True)[0]

We’ll skip self._coef_drift for now.

[4]:

y_f_hat = ODE.f(t=None, y=y)
print(y_f_hat)

tensor([[-0.0191, -0.0182,  0.0078,  0.0020,  0.0070, -0.0190, -0.0366, -0.0121,
         -0.0267, -0.0194,  0.0034, -0.0013, -0.0185, -0.0411, -0.0280,  0.0304,
         -0.0010, -0.0459, -0.0191, -0.0101],
        [-0.0118, -0.0113,  0.0267, -0.0601, -0.0217, -0.0136, -0.0339,  0.0420,
          0.0257,  0.0173, -0.0147,  0.0158, -0.0204,  0.0089, -0.0070,  0.0195,
         -0.0142, -0.0115,  0.0208, -0.0027],
        [ 0.0078, -0.0325,  0.0557, -0.0108,  0.0017, -0.0035, -0.0211, -0.0244,
         -0.0189, -0.0022, -0.0195, -0.0373,  0.0397, -0.0341, -0.0472,  0.0457,
         -0.0276, -0.0150,  0.0251,  0.0113],
        [ 0.0117, -0.0229,  0.0043,  0.0080, -0.0022, -0.0083, -0.0038, -0.0119,
         -0.0275, -0.0035,  0.0126,  0.0075,  0.0523,  0.0035, -0.0165,  0.0104,
          0.0084,  0.0045,  0.0319, -0.0330],
        [-0.0044, -0.0056,  0.0090,  0.0182,  0.0087,  0.0157, -0.0245,  0.0270,
         -0.0402, -0.0136, -0.0249,  0.0182, -0.0473,  0.0016,  0.0184,  0.0178,
         -0.0443, -0.0500, -0.0312,  0.0028]], grad_fn=<MmBackward0>)

[5]:

y_f_hat = ODE.drift(y=y)
print(y_f_hat)

tensor([[-0.0191, -0.0182,  0.0078,  0.0020,  0.0070, -0.0190, -0.0366, -0.0121,
         -0.0267, -0.0194,  0.0034, -0.0013, -0.0185, -0.0411, -0.0280,  0.0304,
         -0.0010, -0.0459, -0.0191, -0.0101],
        [-0.0118, -0.0113,  0.0267, -0.0601, -0.0217, -0.0136, -0.0339,  0.0420,
          0.0257,  0.0173, -0.0147,  0.0158, -0.0204,  0.0089, -0.0070,  0.0195,
         -0.0142, -0.0115,  0.0208, -0.0027],
        [ 0.0078, -0.0325,  0.0557, -0.0108,  0.0017, -0.0035, -0.0211, -0.0244,
         -0.0189, -0.0022, -0.0195, -0.0373,  0.0397, -0.0341, -0.0472,  0.0457,
         -0.0276, -0.0150,  0.0251,  0.0113],
        [ 0.0117, -0.0229,  0.0043,  0.0080, -0.0022, -0.0083, -0.0038, -0.0119,
         -0.0275, -0.0035,  0.0126,  0.0075,  0.0523,  0.0035, -0.0165,  0.0104,
          0.0084,  0.0045,  0.0319, -0.0330],
        [-0.0044, -0.0056,  0.0090,  0.0182,  0.0087,  0.0157, -0.0245,  0.0270,
         -0.0402, -0.0136, -0.0249,  0.0182, -0.0473,  0.0016,  0.0184,  0.0178,
         -0.0443, -0.0500, -0.0312,  0.0028]], grad_fn=<MmBackward0>)

Compute the potential, ψ of the observed state, y - i.e., (ψ(y)):

[6]:

y_hat_potential = ODE._potential(y=y)
print(y_hat_potential)

tensor([[-0.2161],
        [-0.1079],
        [-0.2329],
        [-0.0970],
        [-0.1882]], grad_fn=<MmBackward0>)

PotentialODE.reference.ipynb#

Default PotentialODE#

`PotentialODE.reference.ipynb`#

Default `PotentialODE`#