Commit 1d964d61 authored by hazrmard's avatar hazrmard
Browse files

added Recurrent integrator models in python

parent 94f08eb8
%% Cell type:code id: tags:
``` python
%reload_ext autoreload
%autoreload 2
import numpy as np
from numpy.random import randn
import torch
import torch.nn as nn
import torch.nn.functional as F
from models import TorchEstimator, Accumulator, AffineIntegrator, System
from state_space import SS
```
%% Cell type:code id: tags:
``` python
xrand = torch.rand(100, 20, 1) - 0.5
xneg = torch.rand(100, 20, 1) - 1
xpos = torch.rand(100, 20, 1)
yrand = np.cumsum(xrand, axis=0)
yneg = np.cumsum(xneg, axis=0)
ypos = np.cumsum(xpos, axis=0)
```
%% Cell type:markdown id: tags:
## Anatomy of a dynamic model
![anatomy](../docs/img/dynamic_system_anatomy.png)
A dynamic system consistes of 3 parts:
1. **State dynamics**: The system determines the rate of change of its internal state in response to its current state and outside stimulus. This can be an instantaneous/feed-forward mapping from inputs to rate of change.
2. **State**: The rates of change of state are accumulated to detemine the current state. The system has *memory*.
3. **Output**: The system determines the output from its current state. It can be a feed-forward mapping from state to output.
%% Cell type:code id: tags:
``` python
m = Accumulator()
reg = TorchEstimator(module=m, optimizer=None, epochs=500,
loss=None, batch_size=None, verbose=False, cuda=False)
# reg.fit(xrand, yrand) # <= Accumulator has not training parameters
print(reg.score(xrand, yrand))
print(reg.predict(torch.ones(10,1,1)).squeeze())
```
%% Output
tensor(0.) tensor(9644.0918)
1.0
tensor([ 1., 2., 3., 4., 5., 6., 7., 8., 9., 10.])
%% Cell type:code id: tags:
``` python
m = AffineIntegrator(fixed=False)
reg = TorchEstimator(module=m, optimizer=torch.optim.Adam(m.parameters(), lr=0.01), epochs=100,
loss=nn.MSELoss(), batch_size=10, verbose=False, cuda=False)
reg.fit(xrand, yrand)
print(reg.score(xrand, yrand))
print(reg.predict(torch.ones(10,1,1)).squeeze())
```
%% Output
-368458.125
tensor([ 0.6547, 1.6504, 2.7213, 3.8730, 5.1116, 6.4437, 7.8763, 9.4171,
11.0741, 12.8562])
%% Cell type:code id: tags:
``` python
m = System(fixed=False)
reg = TorchEstimator(module=m, optimizer=torch.optim.Adam(m.parameters(), lr=0.01), epochs=100,
loss=nn.MSELoss(), batch_size=10, verbose=False, cuda=False)
reg.fit(xrand, yrand)
print(reg.score(xrand, yrand))
print(reg.predict(torch.ones(10,1,1)).squeeze())
```
%% Output
-182146056192.0
tensor([ 0.7766, 1.6182, 2.5302, 3.5186, 4.5896, 5.7502, 7.0079, 8.3708,
9.8477, 11.4483])
%% Cell type:code id: tags:
``` python
k = 1 # spring constant
a = 0.5 # drag coefficient
m = 1 # mass
A = np.asarray([[0, 1], [-k/m, -a/m]])
B = np.asarray([[0],[1/m]])
C = np.asarray([[1, 0]])
D = np.asarray([[0]])
model = SS(A,B,C,D,tstep=1e-3)
U = randn(1000, 1)
Y, _ = model.run(U, 0.5)
print(f'Input shape:{U.shape}, output shape: {Y.shape}')
```
%% Output
Input shape:(1000, 1), output shape: (1000, 1)
%% Cell type:code id: tags:
``` python
seqY = torch.from_numpy(Y).view(-1, 1, 1).float()
seqU = torch.from_numpy(U).view(-1, 1, 1).float()
m = AffineIntegrator(input_size=1, fixed=False)
reg = TorchEstimator(module=m, optimizer=torch.optim.Adam(m.parameters(), lr=0.01), epochs=20,
loss=nn.MSELoss(), batch_size=1, verbose=False, cuda=False)
reg.fit(seqU, seqY)
reg.score(seqU, seqY)
```
%% Output
-3.7041192054748535
%% Cell type:code id: tags:
``` python
seqY = torch.from_numpy(Y).view(-1, 1, 1).float()
seqU = torch.from_numpy(U).view(-1, 1, 1).float()
AB = nn.Linear(3, 2)
CD = nn.Linear(3, 1)
m = System(dx=AB, y=CD, input_size=1, state_size=2, output_size=1)
reg = TorchEstimator(module=m, optimizer=torch.optim.Adam(m.parameters(), lr=0.01), epochs=20,
loss=nn.MSELoss(), batch_size=1, verbose=False, cuda=False)
reg.fit(seqU, seqY)
reg.score(seqU, seqY)
```
%% Output
-3.7753424644470215
%% Cell type:code id: tags:
``` python
mfc = nn.Sequential(nn.Linear(1,3), nn.Tanh(), # 6
nn.Linear(3,3), nn.Tanh(), # 12
nn.Linear(3,4), nn.Tanh(), # 16
nn.Linear(4,1)) # 5
reg = TorchEstimator(module=mfc, optimizer=torch.optim.Adam(mfc.parameters(), lr=0.01), epochs=20,
loss=nn.MSELoss(), batch_size=10, verbose=False, cuda=False)
reg.fit(torch.from_numpy(U).float(), torch.from_numpy(Y).float())
reg.score(torch.from_numpy(U).float(), torch.from_numpy(Y).float())
```
%% Output
-3.7677178382873535
"""
Utility definitions for Models.ipynb
"""
from multiprocessing import Pool
from os import cpu_count
from typing import Iterable, Tuple, Iterator
import numpy as np
from sklearn import clone
from sklearn.neural_network import MLPRegressor
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm.auto import trange
def _fit(*args):
"""
Multi-processing payload function used by `fit_composite_model()`.
"""
est, (x, y) = args
return est.fit(x, y)
def fit_composite_model(estimator: MLPRegressor,
data: Iterable[Tuple[np.ndarray, np.ndarray]]) -> Iterable[MLPRegressor]:
"""
Fits copies of an estimator to different datasets in parallel.
Args:
* `estimator`: An estimator instance with a `fit(X, y)` method which returns
the instance.
* `data`: An iterable of tuples of arrays: [(train1, train2,..), (test1, test2,..)].
Returns:
* A list of fitted estimators.
"""
data = list(data)
estimators = [clone(estimator) for _ in data]
with Pool(min(len(data), cpu_count())) as pool:
return pool.starmap(_fit, zip(estimators, data))
class TorchEstimator:
"""
Wraps a `torch.nn.Module` instance with a scikit-learn `Estimator` API.
Args:
* `module`: A `nn.Module` describing the neural network,
* `optimizer`: An `Optimizer` instance which iteratively modifies weights,
* `loss`: a `_Loss` instance which calculates the loss metric,
* `epochs`: The number of times to iterate over the training data,
* `verbose`: Whether to log training progress or not,
* `batch_size`: Chunk size of data for each training step,
* `cuda`: Whether to use GPU acceleration if available.
"""
def __init__(self, module: nn.Module, optimizer: optim.Optimizer,
loss: nn.modules.loss._Loss, epochs: int=10, verbose=True,
batch_size: int=8, cuda=True):
self.module = module
self.optimizer = optimizer
self.loss = loss
self.epochs = epochs
self.verbose = verbose
self.batch_size = batch_size
self.batch_first = self._is_batch_first()
# pylint: disable=no-member
self._device = torch.device('cuda') if torch.cuda.is_available() and cuda\
else torch.device('cpu')
def fit(self, X: torch.Tensor, y: torch.Tensor) -> 'TorchEstimator':
"""
Fit target to features.
Args:
* `X`: `Tensor` of shape (SeqLen, N, Features) or (N, SeqLen, Features)
for recurrent modules or (N, Features) for other modules.
* `y`: `Tensor` of shape ([SeqLen,] N, OutputFeatures) for recurrent
modules of (N, OutputFeatures).
Returns:
* self
"""
if self.verbose:
print()
ranger = trange(self.epochs)
self.module.to(self._device)
for e in ranger:
total_loss = 0.
for instance, target in zip(self._to_batches(X), self._to_batches(y)):
instance, target = instance.to(self._device), target.to(self._device)
self.module.zero_grad()
output = self.module(instance)
loss = self.loss(output, target)
loss.backward()
self.optimizer.step()
total_loss += loss.item()
if self.verbose:
ranger.write(f'Epoch {e+1:3d}\tLoss: {total_loss:10.2f}')
return self
def predict(self, X: torch.Tensor) -> torch.Tensor:
"""
Predict output from inputs.
Args:
* `X`: `Tensor` of shape (SeqLen, N, Features) or (N, SeqLen, Features)
for recurrent modules or (N, Features) for other modules.
Returns:
* `Tensor` of shape ([SeqLen,] N, OutputFeatures) for recurrent
modules of (N, OutputFeatures).
"""
with torch.no_grad():
result = self.module(X)
return result
def score(self, X, y_true) -> float:
"""
Measure how well the estimator learned through the coefficient of
determination.
Args:
* `X`: `Tensor` of shape (SeqLen, N, Features) or (N, SeqLen, Features)
for recurrent modules or (N, Features) for other modules.
* `y_true`: True target values for each of N instances
Returns:
* float: Coefficient of determination.
"""
y_pred = self.predict(X)
residual_squares_sum = ((y_true - y_pred) ** 2).sum()
total_squares_sum = ((y_true - y_true.mean()) ** 2).sum()
return (1 - residual_squares_sum / total_squares_sum).item()
def _to_batches(self, X: torch.Tensor) -> Iterator[torch.Tensor]:
"""
Convert ([SeqLen,] N, Features) to a generator of ([SeqLen,] n, Features)
mini-batches. So for recurrent layers, training can be done in batches.
"""
if not self.batch_first:
# Recurrent layers take inputs of the shape (SeqLen, N, Features...)
# So if there is any recurrent layer in the module, assume that this