base.py

from crypt import methods
import torch
from einops import rearrange
from torch import nn
from torchdiffeq import odeint_adjoint

from basehelper import *


def bmul(vec, mat, axis=0):
    mat = mat.transpose(axis, -1)
    return (mat * vec.expand_as(mat)).transpose(axis, -1)


def bmul(vec, mat, axis=0):
    mat = mat.transpose(axis, -1)
    return (mat * vec.expand_as(mat)).transpose(axis, -1)

class mu_net(nn.Module):
	def __init__(self, input_dim, hid):
		super(mu_net, self).__init__()
		self.input_dim = input_dim
		self.fc1 = nn.Linear(input_dim, hid)
		self.fc2 = nn.Linear(hid, 1)
		self.actv = nn.ReLU()
	
	def forward(self, x):
		mu_out = self.fc1(x)
		mu_out = self.actv(mu_out)
		mu_out = self.fc2(mu_out)
		mu_out = torch.pow(mu_out, 2)
		return mu_out[0][0][0] # silverbox
  		# return mu_out[0][0][0][0] # human activity


class Tinvariant_NLayerNN(NLayerNN):
    def forward(self, t, x):
        return super(Tinvariant_NLayerNN, self).forward(x)


class dfwrapper(nn.Module):
    def __init__(self, df, shape, recf=None):
        super(dfwrapper, self).__init__()
        self.df = df
        self.shape = shape
        self.recf = recf

    def forward(self, t, x):
        bsize = x.shape[0]
        if self.recf:
            x = x[:, :-self.recf.osize].reshape(bsize, *self.shape)
            dx = self.df(t, x)
            dr = self.recf(t, x, dx).reshape(bsize, -1)
            dx = dx.reshape(bsize, -1)
            dx = torch.cat([dx, dr], dim=1)
        else:
            x = x.reshape(bsize, *self.shape)
            dx = self.df(t, x)
            dx = dx.reshape(bsize, -1)
        return dx


class NODEintegrate(nn.Module):

    def __init__(self, df, shape=None, tol=1e-5, adjoint=True, evaluation_times=None, recf=None, algebraic_from_differential=None, differential_from_algebraic=None, activation_h=None, activation_output=None, abs_second=None, time_requires_grad=True, verbose=True):
        """
        Create an OdeRnnBase model
            x' = df(x)
            x(t0) = x0
        :param df: a function that computes derivative. input & output shape [batch, channel, feature]
        :param x0: initial condition.
            - if x0 is set to be nn.parameter then it can be trained.
            - if x0 is set to be nn.Module then it can be computed through some network.
        """
        super().__init__()
        self.df = dfwrapper(df, shape, recf) if shape else df
        self.tol = tol
        if verbose:
            print(f"Tolerance: {self.tol}")
            print(f"Adjoint:", adjoint)
        self.odeint = torchdiffeq.odeint_adjoint if adjoint else torchdiffeq.odeint
        self.evaluation_times = evaluation_times if evaluation_times is not None else torch.Tensor([0.0, 1.0])
        if time_requires_grad == False:
            self.evaluation_times.requires_grad = False
        if verbose:
            print("Evaluation times:", self.evaluation_times)
            print("Learnable eval times:", self.evaluation_times.requires_grad)
        self.shape = shape
        self.recf = recf
        self.algebraic_from_differential = algebraic_from_differential
        self.activation_h = nn.Identity() if activation_h is None else activation_h
        self.activation_output = nn.Identity() if activation_output is None else activation_output
        self.abs_second = abs_second
        if verbose:
            print("self.algebraic_from_differential:", self.algebraic_from_differential)
            print("self.activation_h for NODEintegrate:", self.activation_h)
        if recf:
            assert shape is not None
        self.differential_from_algebraic = differential_from_algebraic

    def forward(self, x0):
        """
        Evaluate odefunc at given evaluation time
        :param x0: shape [batch, channel, feature]. Set to None while training.
        :param evaluation_times: time stamps where method evaluates, shape [time]
        :param x0stats: statistics to compute x0 when self.x0 is a nn.Module, shape required by self.x0
        :return: prediction by ode at evaluation_times, shape [time, batch, channel, feature]
        """

        if self.differential_from_algebraic:
            x, m = self.calc_differential_from_algebraic(x0, self.evaluation_times[0])
            x0 = torch.cat((x, m), dim=1)

        bsize = x0.shape[0]
        if self.shape:
            assert x0.shape[1:] == torch.Size(self.shape), \
                'Input shape {} does not match with model shape {}'.format(x0.shape[1:], self.shape)
            x0 = x0.reshape(bsize, -1)
            if self.abs_second:
                theta, m = torch.split(x0, 1, dim=1)
                m = torch.abs(m)
                x0 = torch.cat((theta, m), dim=1)
            if self.recf:
                reczeros = torch.zeros_like(x0[:, :1])
                reczeros = repeat(reczeros, 'b 1 -> b c', c=self.recf.osize)
                x0 = torch.cat([x0, reczeros], dim=1)
            out = odeint(self.df, x0, self.evaluation_times, rtol=self.tol, atol=self.tol)
            if self.algebraic_from_differential:
                out = self.calc_algebraic_from_differential(out)
            if self.recf:
                rec = out[-1, :, -self.recf.osize:]
                out = out[:, :, :-self.recf.osize]
                out = out.reshape(-1, bsize, *self.shape)
                return out, rec
            else:
                return out
        else:
            if self.abs_second:
                theta, m = torch.split(x0, 1, dim=1)
                m = torch.abs(m)
                x0 = torch.cat((theta, m), dim=1)
            out = odeint(self.df, x0, self.evaluation_times, rtol=self.tol, atol=self.tol)
            if self.algebraic_from_differential:
                out = self.calc_algebraic_from_differential(out)
            return out

    @property
    def nfe(self):
        return self.df.nfe

    def to(self, device, *args, **kwargs):
        super().to(device, *args, **kwargs)
        self.evaluation_times.to(device)

    def calc_algebraic_from_differential(self, z):
        # split the input into the starting time step and the other time steps
        z_0 = z[:1]
        z_T = z[1:] 
        # get the corresponding value of t for the other time steps
        if len(self.evaluation_times.shape) == 2:
            T = self.evaluation_times[:, 1:]
        else:
            T = self.evaluation_times[1:]
        x, m = torch.split(z_T, 1, dim=2)
        # T^(-3/2) * e^(T/2)
        k = torch.pow(T, -3/2) * torch.exp(T / 2)
        if z.is_cuda:
            k = k.to(z.get_device())
            T = T.to(z.get_device())
        # h(T) = [x(T) m(T)] * Transpose([T^(-3/2)*e^(T/2) I])
        # h = x * k
        # dh = k * (m - (3/2 * torch.pow(T, 1/2) * torch.exp(-T/2) - 1/2 * 1/k) * h)
        k = self.activation_h(k)
        h = self.activation_output(bmul(k, x))
        dh = self.activation_output(bmul(k, m - bmul(3/2 * torch.pow(T, 1/2) * torch.exp(-T/2) - 1/2 * 1/k, h)))
        z_t = torch.cat((h, dh), dim=2)
        out = torch.cat((z_0, z_t), dim=0)
        return out

    def calc_differential_from_algebraic(self, z, t):
        h, dh = torch.split(z, 1, dim=1)
        k_reciprocal = 1 / (torch.pow(t, -self.nesterov_factor/2) * torch.exp(t/2))
        if z.is_cuda:
            k_reciprocal = k_reciprocal.to(z.get_device())
        m = (self.nesterov_factor/2 * (1/t) * k_reciprocal - 1/2 * k_reciprocal) * h \
                + k_reciprocal * dh
        x = h * k_reciprocal
        return x, m
    
    def trajectory(self, x, timesteps):
        """Returns ODE trajectory.
        Parameters
        ----------
        x : torch.Tensor
            Shape (batch_size, self.odefunc.data_dim)
        timesteps : int
            Number of timesteps in trajectory.
        """
        integration_time = torch.linspace(0., 1., timesteps)
        return self.forward(x, eval_times=integration_time)


class NODElayer(NODEintegrate):
    def forward(self, x0):
        out = super(NODElayer, self).forward(x0)
        if isinstance(out, tuple):
            out, rec = out
            return out[-1], rec
        else:
            return out[-1]


'''
class ODERNN(nn.Module):
    def __init__(self, node, rnn, evaluation_times, nhidden):
        super(ODERNN, self).__init__()
        self.t = torch.as_tensor(evaluation_times).float()
        self.n_t = len(self.t)
        self.node = node
        self.rnn = rnn
        self.nhidden = (nhidden,) if isinstance(nhidden, int) else nhidden

    def forward(self, x):
        assert len(x) == self.n_t
        batchsize = x.shape[1]
        out = torch.zeros([self.n_t, batchsize, *self.nhidden]).to(x.device)
        for i in range(1, self.n_t):
            odesol = odeint(self.node, out[i - 1], self.t[i - 1:i + 1])
            h_ode = odesol[1]
            out[i] = self.rnn(h_ode, x[i])
        return out
'''


class NODE(nn.Module):
    def __init__(self, df=None, **kwargs):
        super(NODE, self).__init__()
        self.__dict__.update(kwargs)
        self.df = df
        self.nfe = 0
        self.elem_t = None

    def forward(self, t, x):
        self.nfe += 1
        if self.elem_t is None:
            return self.df(t, x)
        else:
            return self.elem_t * self.df(self.elem_t, x)

    def update(self, elem_t):
        self.elem_t = elem_t.view(*elem_t.shape, 1)


class SONODE(NODE):
    def forward(self, t, x):
        """
        Compute [y y']' = [y' y''] = [y' df(t, y, y')]
        :param t: time, shape [1]
        :param x: [y y'], shape [batch, 2, vec]
        :return: [y y']', shape [batch, 2, vec]
        """
        self.nfe += 1
        v = x[:, 1:, :]
        out = self.df(t, x)
        return torch.cat((v, out), dim=1)


class HeavyBallNODE(NODE):
    def __init__(self, df, actv_h=None, gamma_guess=-3.0, gamma_act='sigmoid', corr=-100, corrf=True, sign=1):
        super().__init__(df)
        # Momentum parameter gamma
        self.gamma = Parameter([gamma_guess], frozen=False)
        self.gammaact = nn.Sigmoid() if gamma_act == 'sigmoid' else gamma_act
        self.corr = Parameter([corr], frozen=corrf)
        self.sp = nn.Softplus()
        self.sign = sign # Sign of df
        self.actv_h = nn.Identity() if actv_h is None else actv_h # Activation for dh, GHBNODE only

    def forward(self, t, x):
        """
        Compute [theta' m' v'] with heavy ball parametrization in
        $$ h' = -m $$
        $$ m' = sign * df - gamma * m $$
        based on paper https://www.jmlr.org/papers/volume21/18-808/18-808.pdf
        :param t: time, shape [1]
        :param x: [theta m], shape [batch, 2, dim]
        :return: [theta' m'], shape [batch, 2, dim]
        """
        self.nfe += 1
        h, m = torch.split(x, 1, dim=1)
        dh = self.actv_h(- m)
        dm = self.df(t, h) * self.sign - self.gammaact(self.gamma()) * m
        dm = dm + self.sp(self.corr()) * h
        out = torch.cat((dh, dm), dim=1)
        if self.elem_t is None:
            return out
        else:
            return self.elem_t * out

    def update(self, elem_t):
        self.elem_t = elem_t.view(*elem_t.shape, 1, 1)


HBNODE = HeavyBallNODE # Alias


class NesterovNODE(NODE):
    def __init__(self, df, actv_h=None, corr=-100, corrf=True, use_h=False, full_details=False, algebraic_from_differential=True, actv_m=None, actv_dm=None, actv_df=None, sign=1):
        super().__init__(df)
        self.corr = Parameter([corr], frozen=corrf)
        self.sp = nn.Softplus()
        self.sign = sign # Sign of df
        self.actv_h = nn.Identity() if actv_h is None else actv_h # Activation for dh, GNNODE only
        self.actv_m = nn.Identity() if actv_m is None else actv_m # Activation for dh, GNNODE only
        self.actv_dm = nn.Identity() if actv_dm is None else actv_dm # Activation for dh, GNNODE only
        self.actv_df = nn.Identity() if actv_df is None else actv_df # Activation for df, GNNODE only
        self.use_h = use_h
        self.full_details = full_details
        self.algebraic_from_differential = algebraic_from_differential

    def forward(self, t, z):
        """
        Compute [x' m'] with diff-alg nesterov parametrization in
        $$ h' = -m $$
        $$ m' = sign * df(t, h) - m - xi * h $$
        :param t: time, shape [1]
        :param z: [h dh], shape [batch, 2, dim]
        :return: [x' m'], shape [batch, 2, dim]
        """
        self.nfe += 1
        x, m = torch.split(z, 1, dim=1)
        dx = self.actv_h(m)
        k = torch.pow(t, -3/2) * torch.exp(t/2)
        h = x * self.actv_h(k)
        dm = self.actv_df(self.df(t, h)) * self.sign - m
        dm = self.actv_dm(self.actv_m(dm) - self.sp(self.corr()) * h)
        out = torch.cat((dx, dm), dim=1)
        if self.elem_t is None:
            return out
        else:
            return self.elem_t * out

    def update(self, elem_t):
        self.elem_t = elem_t.view(*elem_t.shape, 1, 1)

NNODE = NesterovNODE # Alias


class RMSpropNODE(NODE):
    def __init__(self, df, alpha=1e-1, actv_dtheta=None, corr=-100, corrf=True):
        super().__init__(df)
        self.alpha = alpha
        self.corr = Parameter([corr], frozen=corrf)
        self.sp = nn.Softplus()
        self.actv_dtheta = nn.Identity() if actv_dtheta is None else actv_dtheta

    def forward(self, t, x):
        self.nfe += 1
        theta, m = torch.split(x, 1, dim=1)

        m = nn.Tanh()(torch.abs(m))
        gtheta = self.df(t, theta) * self.df(t, m)
        
        dtheta = -self.actv_dtheta(gtheta)
        dm = 1 / self.alpha * (self.actv_dtheta(gtheta ** 2) - 1) * m - self.sp(self.corr()) * theta
        
        out = torch.cat((dtheta, dm), dim=1)
        if self.elem_t is None:
            return out
        else:
            return self.elem_t * out

    def update(self, elem_t):
        self.elem_t = elem_t.view(*elem_t.shape, 1, 1)


class TOAES_NODE(NODE):
    def __init__(self, df, actv_df=None, thetaact=None, learnable_mu = False, mu=1.0, corr=-100):
        super().__init__(df)
        # Momentum parameter gamma
        self.actv_df = nn.Identity() if actv_df is None else actv_df 
        self.thetaact = nn.Identity() if thetaact is None else thetaact
        self.mu = mu
        self.learnable_mu = learnable_mu
        self.sp = nn.Softplus()
        self.corr = Parameter([corr], frozen=True)
        # self.mu_net = mu_net(28, 10) # human activity
        self.mu_net = mu_net(1, 10) # silverbox

    def forward(self, t, x):
        self.nfe += 1
        h, p, q = torch.split(x, 1, dim=1)
			
        if self.learnable_mu:
            self.mu = self.mu_net(x)
            f = self.df(t, h + 1/torch.sqrt(self.mu) * p)
            dh = self.thetaact(p)
            dp = self.thetaact(q)
            dq = -3 * torch.sqrt(self.mu) * q - 2 * self.mu * p - torch.sqrt(self.mu) * self.actv_df(f) - self.sp(self.corr()) * h
            out = torch.cat((dh, dp, dq), dim=1)
            
        else:
            f = self.df(t, h + 1/np.sqrt(self.mu) * p)
            dh = self.thetaact(p)
            dp = self.thetaact(q)
            dq = -3 * np.sqrt(self.mu) * q - 2 * self.mu * p - np.sqrt(self.mu) * self.actv_df(f) - self.sp(self.corr()) * h
            out = torch.cat((dh, dp, dq), dim=1)
        
        if self.elem_t is None:
            return out
        else:
            return self.elem_t * out
    

class ODE_RNN(nn.Module):
    def __init__(self, ode, rnn, nhid, ic, rnn_out=False, both=False, tol=1e-7):
        super().__init__()
        self.ode = ode
        self.t = torch.Tensor([0, 1])
        self.nhid = [nhid] if isinstance(nhid, int) else nhid
        self.rnn = rnn
        self.tol = tol
        self.rnn_out = rnn_out
        self.ic = ic
        self.both = both

    def forward(self, t, x, multiforecast=None):
        """
        --
        :param t: [time, batch]
        :param x: [time, batch, ...]
        :return: [time, batch, *nhid]
        """
        n_t, n_b = t.shape
        h_ode = torch.zeros(n_t + 1, n_b, *self.nhid, device=x.device)
        h_rnn = torch.zeros(n_t + 1, n_b, *self.nhid, device=x.device)
        if self.ic:
            h_ode[0] = h_rnn[0] = self.ic(rearrange(x, 't b c -> b (t c)')).view(h_ode[0].shape)
        if self.rnn_out:
            for i in range(n_t):
                self.ode.update(t[i])
                h_ode[i] = odeint(self.ode, h_rnn[i], self.t, atol=self.tol, rtol=self.tol)[-1]
                h_rnn[i + 1] = self.rnn(h_ode[i], x[i])
            out = (h_rnn,)
        else:
            for i in range(n_t):
                self.ode.update(t[i])
                h_rnn[i] = self.rnn(h_ode[i], x[i])
                h_ode[i + 1] = odeint(self.ode, h_rnn[i], self.t, atol=self.tol, rtol=self.tol)[-1]
            out = (h_ode,)

        if self.both:
            out = (h_rnn, h_ode)

        if multiforecast is not None:
            self.ode.update(torch.ones_like((t[0])))
            forecast = odeint(self.ode, out[-1][-1], multiforecast * 1.0, atol=self.tol, rtol=self.tol)
            out = (*out, forecast)

        return out


class ODE_RNN_with_Grad_Listener(nn.Module):
    def __init__(self, ode, rnn, nhid, ic, rnn_out=False, both=False, tol=1e-7, method="dopri5", evaluation_times=None, algebraic_from_differential=None, differential_from_algebraic=None, activation_h=None, time_requires_grad=True):
        super().__init__()
        self.ode = ode
        self.evaluation_times = evaluation_times if evaluation_times is not None else torch.Tensor([0.0, 1.0])
        if time_requires_grad == False:
            self.evaluation_times.requires_grad = False
        self.algebraic_from_differential = algebraic_from_differential
        self.differential_from_algebraic = differential_from_algebraic
        self.nesterov_factor = 3
        self.activation_h = nn.Identity() if activation_h is None else activation_h
        self.nhid = [nhid] if isinstance(nhid, int) else nhid
        self.rnn = rnn
        self.tol = tol
        self.rnn_out = rnn_out
        self.ic = ic
        self.both = both
        self.method = method

    def forward(self, t, x, multiforecast=None, retain_grad=False):
        """
        --
        :param t: [time, batch]
        :param x: [time, batch, ...]
        :return: [time, batch, *nhid]
        """
        n_t, n_b = t.shape
        h_ode = [None] * (n_t + 1)
        h_rnn = [None] * (n_t + 1)
        h_ode[-1] = h_rnn[-1] = torch.zeros(n_b, *self.nhid, device=x.device)

        if self.ic:
            h_ode[0] = h_rnn[0] = self.ic(rearrange(x, 't b c -> b (t c)')).view((n_b, *self.nhid))
        else:
            h_ode[0] = h_rnn[0] = torch.zeros(n_b, *self.nhid, device=x.device)
        
        if self.rnn_out:
            for i in range(n_t):
                self.ode.update(t[i])
                if self.differential_from_algebraic:
                    h_rnn[i] = self.calc_differential_from_algebraic(h_rnn[i], t[i])
                h_ode[i] = odeint(self.ode, h_rnn[i], self.evaluation_times, atol=self.tol, rtol=self.tol, method=self.method)[-1]
                if self.algebraic_from_differential:
                    h_ode[i] = self.calc_algebraic_from_differential(h_ode[i])
                h_rnn[i + 1] = self.rnn(h_ode[i], x[i])
            out = (h_rnn,)
        else:
            for i in range(n_t):
                self.ode.update(t[i])
                h_rnn[i] = self.rnn(h_ode[i], x[i])
                if self.differential_from_algebraic:
                    h_rnn[i] = self.calc_differential_from_algebraic(h_rnn[i], t[i])
                h_ode[i + 1] = odeint(self.ode, h_rnn[i], self.evaluation_times, atol=self.tol, rtol=self.tol, method=self.method)[-1]
                if self.algebraic_from_differential:
                    h_ode[i + 1] = self.calc_algebraic_from_differential(h_ode[i + 1])
            out = (h_ode,)

        if self.both:
            out = (h_rnn, h_ode)

        out = [torch.stack([k.to(x.device) for k in h], dim=0) for h in out]

        if multiforecast is not None:
            self.ode.update(torch.ones_like((t[0]), device=x.device))
            forecast = odeint(self.ode, out[-1][-1], multiforecast * 1.0, atol=self.tol, rtol=self.tol, method=self.method)
            if self.algebraic_from_differential:
                forecast = self.calc_algebraic_from_differential(forecast)
            out = (*out, forecast)

        if retain_grad:
            self.h_ode = h_ode
            self.h_rnn = h_rnn
            for i in range(n_t + 1):
                if self.h_ode[i].requires_grad:
                    self.h_ode[i].retain_grad()
                if self.h_rnn[i].requires_grad:
                    self.h_rnn[i].retain_grad()

        return out

    def calc_algebraic_from_differential(self, z):
        # split the input into the starting time step and the other time steps
        z_0 = z[:1]
        z_T = z[1:] 
        # get the corresponding value of t for the other time steps
        if len(self.evaluation_times.shape) == 2:
            T = self.evaluation_times[:, 1:]
        else:
            T = self.evaluation_times[1:]
        x, m = torch.split(z_T, 1, dim=-2)
        # T^(-3/2) * e^(T/2)
        k = torch.pow(T, -self.nesterov_factor/2) * torch.exp(T / 2)
        if z.is_cuda:
            k = k.to(z.get_device())
            T = T.to(z.get_device())
        # h(T) = [x(T) m(T)] * Transpose([T^(-3/2)*e^(T/2) I])
        # h = x * k
        # dh = k * (m - (3/2 * torch.pow(T, 1/2) * torch.exp(-T/2) - 1/2 * 1/k) * h)
        k = self.activation_h(k)
        h = bmul(k, x)
        dh = bmul(k, m - bmul(self.nesterov_factor/2 * torch.pow(T, 1/2) * torch.exp(-T/2) - 1/2 * 1/k, h))
        z_t = torch.cat((h, dh), dim=-2)
        out = torch.cat((z_0, z_t), dim=0)
        return out


    def calc_differential_from_algebraic(self, z, t):
        h, dh = torch.split(z, 1, dim=1)
        k_reciprocal = 1 / (torch.pow(t, -self.nesterov_factor/2) * torch.exp(t/2))
        if z.is_cuda:
            k_reciprocal = k_reciprocal.to(z.get_device())
        m = bmul((self.nesterov_factor/2 * (1/t) * k_reciprocal - 1/2 * k_reciprocal), h) \
                + bmul(k_reciprocal, dh)
        x = bmul(h, k_reciprocal)
        return torch.cat([x, m], dim=1)