PSL Systems to Neuromancer Blocks

[cjb873]

The PSL system library has a lot of systems that are not implemented in neuromancer.dynamics, such as the Boid's system, RC Networks, etc. I see that the base classes have a forward function defined so that these dynamical systems can be wrapped with a Node and used with the System class, yet these forward functions as currently implemented do not support minibatching. I believe this is due to the fact that integrators are called within the forward function, as well as the actual equations of these systems not being able to act on two dimensional tensors.

I think it could be useful to pursue an implementation of forward for these systems that can be wrapped and placed in an integrator, in an example like:

sys = RC_Network(nx=5)


class RC_Block(ode.ODESystem):

    def __init__(self, net):
        super().__init__(net.nx, net.nu)
        self.net = net

    def ode_equations(self, x, u):
        return self.net(x, u)

integrator = integrators.Euler(RC_Block(sys), h=ts)
integrator_node = Node(integrator, ['x', 'u'], ['x'])

Or, it may be best to just implement a version of these dynamical systems inside of neuromancer.dynamics. Either way, I personally would like it to be easier to use white-box representations of these coupled systems to do DPC.

[drgona]

@cjb873 you can actually do that with the current PSL models. We have implemented pytorch backend for all PSL models.
All PSL models inherit from this base class.
All you need to do is to set the argument backend='torch' when instantiating the model.

[cjb873]

@drgona Then this might be better to discuss as a bug, because the following code produces an error:

import torch
import numpy as np
from neuromancer.psl.coupled_systems import RC_Network

torch.manual_seed(0)

adj = np.array([[0,1],[0,2],[0,3],[1,0],[1,3],[1,4],[2,0],[2,3],[3,0],[3,1],[3,2],[3,4],[4,1],[4,3]]).T
sys = RC_Network(nx=5, backend='torch')
nx = sys.nx
nu = sys.nx+1
nref = nx
ts = sys.ts

x = torch.rand(2, nx)
u = torch.ones(2, nu)

sys(x, u)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 29
     26 x = torch.rand(2, nx)
     27 u = torch.ones(2, nu)
---> 29 sys(x, u)

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/neuromancer/psl/base.py:101, in cast_backend.<locals>._impl(self, *method_args, **method_kwargs)
     99 @functools.wraps(method)
    100 def _impl(self, *method_args, **method_kwargs):
--> 101     method_output = method(self, *method_args, **method_kwargs)
    102     if type(method_output) is dict:
    103         return {k: self.B.cast(v, dtype=self.B.core.float32) if k != 'ts' else v
    104                 for k, v in method_output.items()}

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/neuromancer/psl/base.py:346, in ODE_NonAutonomous.forward(self, x, u)
    344 U = self.B.cat([u, u], axis=0)
    345 Time = self.B.core.arange(0, 2) * self.ts
--> 346 equation = EquationWrapper(Time, U, self.equations, self.B)
    347 if self.B.core is torch:
    348     xdot = self.B.odeint(equation, x.flatten(), Time, options={"grid_points": Time, "eps": 1e-6})

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/neuromancer/psl/base.py:84, in EquationWrapper.__init__(self, Time, U, equations, backend)
     82     self.ufunc = scipy.interpolate.interp1d(Time, U, kind='previous', axis=0, fill_value='extrapolate')
     83 elif backend.backend == 'torch':
---> 84     self.ufunc = scipy.interpolate.interp1d(Time, numpy.array(U.numpy(force=True)), kind='previous', axis=0, fill_value='extrapolate')
     85 self.equations = equations
     86 self.B = backend

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/scipy/interpolate/_interpolate.py:267, in interp1d.__init__(self, x, y, kind, axis, copy, bounds_error, fill_value, assume_sorted)
    263 def __init__(self, x, y, kind='linear', axis=-1,
    264              copy=True, bounds_error=None, fill_value=np.nan,
    265              assume_sorted=False):
    266     """ Initialize a 1-D linear interpolation class."""
--> 267     _Interpolator1D.__init__(self, x, y, axis=axis)
    269     self.bounds_error = bounds_error  # used by fill_value setter
    271     # `copy` keyword semantics changed in NumPy 2.0, once that is
    272     # the minimum version this can use `copy=None`.

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/scipy/interpolate/_polyint.py:57, in _Interpolator1D.__init__(self, xi, yi, axis)
     55 self.dtype = None
     56 if yi is not None:
---> 57     self._set_yi(yi, xi=xi, axis=axis)

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/scipy/interpolate/_polyint.py:127, in _Interpolator1D._set_yi(self, yi, xi, axis)
    125     shape = (1,)
    126 if xi is not None and shape[axis] != len(xi):
--> 127     raise ValueError("x and y arrays must be equal in length along "
    128                      "interpolation axis.")
    130 self._y_axis = (axis % yi.ndim)
    131 self._y_extra_shape = yi.shape[:self._y_axis] + yi.shape[self._y_axis+1:]

ValueError: x and y arrays must be equal in length along interpolation axis.

Even if the backend is specified as torch it seems that the forward function implemented in ODE_Nonautonomous which inherits EmulatorBase is causing the error.

[cjb873]

Additionally, when the batch size is equal to 1, we have something like:

x = torch.rand(1, nx)
u = torch.rand(1, nu)

sys(x, u)

Then the error becomes:

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/neuromancer/psl/base.py:89, in EquationWrapper.__call__(self, t, x)
     88 def __call__(self, t, x):
---> 89     return self.equations(t, x, self.B.cast(self.ufunc(t)))

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/neuromancer/psl/coupled_systems.py:229, in RC_Network.equations(self, t, x, u)
    226 dx = np.zeros_like(x)
    228 #Internal heat transfer
--> 229 messages = self.message_passing(x[self.adj_list[0]], x[self.adj_list[1]], t, u)
    230 np.add.at(dx, self.adj_list[0], messages)
    232 #Outside heat transfer

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/neuromancer/psl/coupled_systems.py:219, in RC_Network.message_passing(self, receivers, senders, t, u)
    217 else:
    218     C = self.C
--> 219 messages = (1.0 / (R*C)) * (senders - receivers)
    220 return messages

TypeError: unsupported operand type(s) for *: 'numpy.ndarray' and 'Tensor'

It seems like the actual RC_Network system is hardcoded for numpy.

[drgona]

@cjb873 I agree this is a bug for coupled systems.
Did you try other non-coupled systems with torch backend?

[cjb873]

Others work when the batch size is one. However, when you increase the batch size, you get the same error.
Example:

from neuromancer.psl.nonautonomous import LorenzControl
import torch
torch.manual_seed(0)

modelSystem = LorenzControl(backend='torch')
nx = modelSystem.nx
nu = modelSystem.nu
x = torch.rand(2, nx)
u = torch.rand(2, nu)
modelSystem(x, u)

File ~/miniconda3/envs/neuromancer/lib/python3.10/site-packages/scipy/interpolate/_polyint.py:127, in _Interpolator1D._set_yi(self, yi, xi, axis)
    125     shape = (1,)
    126 if xi is not None and shape[axis] != len(xi):
--> 127     raise ValueError("x and y arrays must be equal in length along "
    128                      "interpolation axis.")
    130 self._y_axis = (axis % yi.ndim)
    131 self._y_extra_shape = yi.shape[:self._y_axis] + yi.shape[self._y_axis+1:]

ValueError: x and y arrays must be equal in length along interpolation axis.