vae.py

import math
import numpy as np
import numpy.random as npr
from sklearn.preprocessing import StandardScaler
from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau
from tqdm import tqdm_notebook as tqdm

import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns

from base import split_on_change
from utils import get_data

sns.color_palette("bright")
import matplotlib as mpl
import matplotlib.cm as cm

import torch
from torch import Tensor
from torch import nn
from torch.nn import functional as F
from torch.autograd import Variable


def ode_solve(z0, t0, t1, f):
    """
    Simplest Euler ODE initial value solver
    """
    h_max = 0.05
    n_steps = math.ceil((abs(t1 - t0) / h_max).max().item())

    h = (t1 - t0) / n_steps
    t = t0
    z = z0

    for i_step in range(n_steps):
        z = z + h * f(z, t)
        t = t + h
    return z


def runge_kutta_4(z0, t0, t1, f):
    """
    Runge Kutta 4th order solver
    """
    h_max = 0.05
    n_steps = math.ceil((abs(t1 - t0) / h_max).max().item())

    h = (t1 - t0) / n_steps
    t = t0
    z = z0

    for i in range(n_steps):
        k1 = f(z, t)
        k2 = f(z + k1 * h / 2.0, t + h / 2.0)
        k3 = f(z + k2 * h / 2.0, t + h / 2.0)
        k4 = f(z + k3 * h, t + h)
        z = z + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
        t = t + h
    return z


class ODEF(nn.Module):
    def forward_with_grad(self, z, t, grad_outputs):
        """Compute f and a df/dz, a df/dp, a df/dt"""
        batch_size = z.shape[0]

        out = self.forward(z, t)

        a = grad_outputs
        adfdz, adfdt, *adfdp = torch.autograd.grad(
            (out,),
            (z, t) + tuple(self.parameters()),
            grad_outputs=(a),
            allow_unused=True,
            retain_graph=True,
        )
        # grad method automatically sums gradients for batch items, we have to expand them back
        if adfdp is not None:
            adfdp = torch.cat([p_grad.flatten() for p_grad in adfdp]).unsqueeze(0)
            adfdp = adfdp.expand(batch_size, -1) / batch_size
        if adfdt is not None:
            adfdt = adfdt.expand(batch_size, 1) / batch_size
        return out, adfdz, adfdt, adfdp

    def flatten_parameters(self):
        p_shapes = []
        flat_parameters = []
        for p in self.parameters():
            p_shapes.append(p.size())
            flat_parameters.append(p.flatten())
        return torch.cat(flat_parameters)


class ODEAdjoint(torch.autograd.Function):
    @staticmethod
    def forward(ctx, z0, t, flat_parameters, func):
        assert isinstance(func, ODEF)
        bs, *z_shape = z0.size()
        time_len = t.size(0)

        with torch.no_grad():
            z = torch.zeros(time_len, bs, *z_shape).to(z0)
            z[0] = z0
            for i_t in range(time_len - 1):
                z0 = runge_kutta_4(z0, t[i_t], t[i_t + 1], func)
                z[i_t + 1] = z0

        ctx.func = func
        ctx.save_for_backward(t, z.clone(), flat_parameters)
        return z

    @staticmethod
    def backward(ctx, dLdz):
        """
        dLdz shape: time_len, batch_size, *z_shape
        """
        func = ctx.func
        t, z, flat_parameters = ctx.saved_tensors
        time_len, bs, *z_shape = z.size()
        n_dim = np.prod(z_shape)
        n_params = flat_parameters.size(0)

        # Dynamics of augmented system to be calculated backwards in time
        def augmented_dynamics(aug_z_i, t_i):
            """
            tensors here are temporal slices
            t_i - is tensor with size: bs, 1
            aug_z_i - is tensor with size: bs, n_dim*2 + n_params + 1
            """
            z_i, a = (
                aug_z_i[:, :n_dim],
                aug_z_i[:, n_dim : 2 * n_dim],
            )  # ignore parameters and time

            # Unflatten z and a
            z_i = z_i.view(bs, *z_shape)
            a = a.view(bs, *z_shape)
            with torch.set_grad_enabled(True):
                t_i = t_i.detach().requires_grad_(True)
                z_i = z_i.detach().requires_grad_(True)
                func_eval, adfdz, adfdt, adfdp = func.forward_with_grad(
                    z_i, t_i, grad_outputs=a
                )  # bs, *z_shape
                adfdz = (
                    adfdz.to(z_i)
                    if adfdz is not None
                    else torch.zeros(bs, *z_shape).to(z_i)
                )
                adfdp = (
                    adfdp.to(z_i)
                    if adfdp is not None
                    else torch.zeros(bs, n_params).to(z_i)
                )
                adfdt = (
                    adfdt.to(z_i) if adfdt is not None else torch.zeros(bs, 1).to(z_i)
                )

            # Flatten f and adfdz
            func_eval = func_eval.view(bs, n_dim)
            adfdz = adfdz.view(bs, n_dim)
            return torch.cat((func_eval, -adfdz, -adfdp, -adfdt), dim=1)

        dLdz = dLdz.view(time_len, bs, n_dim)  # flatten dLdz for convenience
        with torch.no_grad():
            ## Create placeholders for output gradients
            # Prev computed backwards adjoints to be adjusted by direct gradients
            adj_z = torch.zeros(bs, n_dim).to(dLdz)
            adj_p = torch.zeros(bs, n_params).to(dLdz)
            # In contrast to z and p we need to return gradients for all times
            adj_t = torch.zeros(time_len, bs, 1).to(dLdz)

            for i_t in range(time_len - 1, 0, -1):
                z_i = z[i_t]
                t_i = t[i_t]
                f_i = func(z_i, t_i).view(bs, n_dim)

                # Compute direct gradients
                dLdz_i = dLdz[i_t]
                dLdt_i = torch.bmm(
                    torch.transpose(dLdz_i.unsqueeze(-1), 1, 2), f_i.unsqueeze(-1)
                )[:, 0]

                # Adjusting adjoints with direct gradients
                adj_z += dLdz_i
                adj_t[i_t] = adj_t[i_t] - dLdt_i

                # Pack augmented variable
                aug_z = torch.cat(
                    (
                        z_i.view(bs, n_dim),
                        adj_z,
                        torch.zeros(bs, n_params).to(z),
                        adj_t[i_t],
                    ),
                    dim=-1,
                )

                # Solve augmented system backwards
                aug_ans = runge_kutta_4(aug_z, t_i, t[i_t - 1], augmented_dynamics)

                # Unpack solved backwards augmented system
                adj_z[:] = aug_ans[:, n_dim : 2 * n_dim]
                adj_p[:] += aug_ans[:, 2 * n_dim : 2 * n_dim + n_params]
                adj_t[i_t - 1] = aug_ans[:, 2 * n_dim + n_params :]

                del aug_z, aug_ans

            ## Adjust 0 time adjoint with direct gradients
            # Compute direct gradients
            dLdz_0 = dLdz[0]
            dLdt_0 = torch.bmm(
                torch.transpose(dLdz_0.unsqueeze(-1), 1, 2), f_i.unsqueeze(-1)
            )[:, 0]

            # Adjust adjoints
            adj_z += dLdz_0
            adj_t[0] = adj_t[0] - dLdt_0
        return adj_z.view(bs, *z_shape), adj_t, adj_p, None


class NeuralODE(nn.Module):
    def __init__(self, func):
        super(NeuralODE, self).__init__()
        assert isinstance(func, ODEF)
        self.func = func

    def forward(self, z0, t=Tensor([0.0, 1.0]), return_whole_sequence=False):
        t = t.to(z0)
        z = ODEAdjoint.apply(z0, t, self.func.flatten_parameters(), self.func)
        if return_whole_sequence:
            return z
        else:
            return z[-1]


class LinearODEF(ODEF):
    def __init__(self, W):
        super(LinearODEF, self).__init__()
        self.lin = nn.Linear(2, 2, bias=False)
        self.lin.weight = nn.Parameter(W)

    def forward(self, x, t):
        return self.lin(x)


def gen_batch(samp_trajs, samp_ts, batch_size, n_sample=100):
    n_batches = samp_trajs.shape[1] // batch_size
    time_len = samp_trajs.shape[0]
    n_sample = min(n_sample, time_len)
    batches = []
    for i in range(n_batches):
        if n_sample > 0:
            t0_idx = npr.multinomial(
                1, [1.0 / (time_len - n_sample)] * (time_len - n_sample)
            )
            t0_idx = np.argmax(t0_idx)
            tM_idx = t0_idx + n_sample
        else:
            t0_idx = 0
            tM_idx = time_len

        frm, to = batch_size * i, batch_size * (i + 1)

        yield samp_trajs[t0_idx:tM_idx, frm:to], samp_ts[t0_idx:tM_idx, frm:to]


class NNODEF(ODEF):
    def __init__(self, in_dim, hid_dim, time_invariant=False):
        super(NNODEF, self).__init__()
        self.time_invariant = time_invariant

        if time_invariant:
            self.lin1 = nn.Linear(in_dim, hid_dim)
        else:
            self.lin1 = nn.Linear(in_dim + 1, hid_dim)
        self.lin2 = nn.Linear(hid_dim, hid_dim)
        self.lin3 = nn.Linear(hid_dim, in_dim)
        self.elu = nn.ELU(inplace=True)

    def forward(self, x, t):
        if not self.time_invariant:
            x = torch.cat((x, t), dim=-1)

        h = self.elu(self.lin1(x))
        h = self.elu(self.lin2(h))
        out = self.lin3(h)
        return out


class RNNEncoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(RNNEncoder, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim

        self.rnn = nn.GRU(input_dim + 1, hidden_dim)
        self.hid2lat = nn.Linear(hidden_dim, 2 * latent_dim)

    def forward(self, x, t):

        t = t.clone()
        t[1:] = t[:-1] - t[1:]
        t[0] = 0.0
        xt = torch.cat((x, t.unsqueeze(1)), dim=-1)

        _, h0 = self.rnn(xt.flip((0,)))  # Reversed
        # Compute latent dimension
        z0 = self.hid2lat(h0[0])
        z0_mean = z0[:, : self.latent_dim]
        z0_log_var = z0[:, self.latent_dim :]
        return z0_mean, z0_log_var


class NeuralODEDecoder(nn.Module):
    def __init__(self, output_dim, hidden_dim, latent_dim):
        super(NeuralODEDecoder, self).__init__()
        self.output_dim = output_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim

        func = NNODEF(latent_dim, hidden_dim, time_invariant=True)
        self.ode = NeuralODE(func)
        self.l2h = nn.Linear(latent_dim, hidden_dim)
        self.h2o = nn.Linear(hidden_dim, output_dim)

    def forward(self, z0, t):
        zs = self.ode(z0, t, return_whole_sequence=True)

        hs = self.l2h(zs)
        xs = self.h2o(hs)
        return xs


class ODEVAE(nn.Module):
    def __init__(self, output_dim, hidden_dim, latent_dim):
        super(ODEVAE, self).__init__()
        self.output_dim = output_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim

        self.encoder = RNNEncoder(output_dim, hidden_dim, latent_dim)
        self.decoder = NeuralODEDecoder(output_dim, hidden_dim, latent_dim)

    def forward(self, x, t, MAP=False):
        z_mean, z_log_var = self.encoder(x, t)
        if MAP:
            z = z_mean
        else:
            z = z_mean + torch.randn_like(z_mean) * torch.exp(0.5 * z_log_var)
        x_p = self.decoder(z, t)
        return x_p, z, z_mean, z_log_var

    def generate_with_seed(self, seed_x, t):
        seed_t_len = seed_x.shape[0]
        seed_x = seed_x[:, 0].unsqueeze(1)
        z_mean, z_log_var = self.encoder(seed_x, t[:seed_t_len])
        x_p = self.decoder(z_mean, t)
        return x_p


def get_train_samples(train_data):
    f46 = train_data[:2500]
    t_46 = f46[:, 0]
    train_46 = f46[:, 1:]

    f49 = train_data[2500:]
    t_49 = f49[:, 0]
    train_49 = f49[:, 1:]

    return np.stack((train_46, train_49), axis=1), np.stack((t_46, t_49), axis=1)


def to_np(x):
    return x.detach().cpu().numpy()


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model_name = "ltc_2"
    architecture = "LTC"
    input_dim = 3
    output_dim = 3

    batch_size = 246
    hidden_dim = 16
    num_layers = 3
    lr = 0.01
    num_epochs = 5000

    pred_steps = 2500
    val_freq = 5000
    plot_range = 250

    model = ODEVAE(output_dim=3, hidden_dim=128, latent_dim=4)

    param_dict = {
        "Model Name": model_name,
        "Architecture": architecture,
        "Input Dim": input_dim,
        "Output Dim": output_dim,
        "Hidden Dim": hidden_dim,
        "Num Layers": num_layers,
        "Batch Size": batch_size,
        "Epochs": num_epochs,
        "Learning Rate": lr,
        "Loss": "RMSE",
        "Shuffle Train": "True",
        "Scaler": "StandardScaler",
    }

    torch.manual_seed(97)

    train_data = get_data("DORA_Train.csv")[:, 0:4]
    train_data, t_train = get_train_samples(train_data)

    train_data = torch.tensor(train_data, dtype=torch.float32)[:2500]
    t_train = torch.tensor(t_train, dtype=torch.float32)[:2500]

    test_data = get_data("DORA_Test.csv")[:, 0:5]

    splitted_test = split_on_change(test_data)
    t_test = splitted_test[:, :, 0]
    data_test = splitted_test[:, :, 1:]

    optim = torch.optim.Adam(model.parameters(), lr=0.0001)
    scheduler = ReduceLROnPlateau(optimizer=optim, factor=0.5)
    n_epochs = 10000
    noise_std = 0.02
    losses = []
    best_loss = 1e20
    for epoch_idx in range(n_epochs):
        epoch_loss = 0.0
        train_iter = gen_batch(
            samp_trajs=train_data, samp_ts=t_train, batch_size=1, n_sample=200
        )
        for x, t in train_iter:
            optim.zero_grad()

            max_len = np.random.choice([50])
            permutation = np.random.permutation(t.shape[0])
            np.random.shuffle(permutation)
            permutation = np.sort(permutation[:max_len])

            x, t = x[permutation], t[permutation]

            x_p, z, z_mean, z_log_var = model(x, t)
            kl_loss = -0.5 * torch.sum(
                1 + z_log_var - z_mean**2 - torch.exp(z_log_var), -1
            )
            loss = (0.5 * ((x - x_p) ** 2).sum(-1).sum(0) / noise_std**2) + kl_loss
            loss = torch.mean(loss)
            loss /= max_len
            epoch_loss += loss.item()
            loss.backward()
            optim.step()
            # scheduler.step(loss)
            # print(scheduler.get_last_lr())
        losses.append(epoch_loss)

        if epoch_loss < best_loss:
            best_loss = epoch_loss
            torch.save(model.state_dict(), "models/best_vae_rk4.pth")
            print(f"New best loss: {best_loss}")

        print(f"Epoch {epoch_idx}, Loss: {epoch_loss}")

    # if epoch_idx % 2499 == 0:
    model.load_state_dict(torch.load("models/best_vae_rk4.pth"))
    frm, to, to_seed = 0, 100, 10
    seed_trajs = train_data[frm:to_seed]
    ts = t_train[frm:to][:, 0].unsqueeze(-1)

    samp_trajs_p = to_np(model.generate_with_seed(seed_trajs, ts))

    plt.plot(ts.numpy(), samp_trajs_p[:, 0, 0], label="Pred")
    plt.plot(
        ts.numpy(),
        train_data[1 : samp_trajs_p.shape[0] + 1, 0, 0],
        label="True",
    )
    plt.legend()
    plt.grid(True)
    plt.show()

    plt.figure(figsize=(10, 5))
    plt.plot(range(1, n_epochs + 1), losses, label="Training Loss")
    plt.title("Training Loss Curve")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.yscale("log")
    plt.legend()
    plt.grid(True)
    plt.show()
    frm, to, to_seed = 0, 200, 50
    seed_trajs = train_data[frm:to_seed]
    ts = t_train[frm:to]

    # samp_trajs_p = to_np(model.generate_with_seed(seed_trajs, ts))

    # fig, axes = plt.subplots(nrows=3, ncols=3, figsize=(15, 9))
    # axes = axes.flatten()
    # for i, ax in enumerate(axes):
    #     ax.scatter(
    #         to_np(seed_trajs[:, i, 0]),
    #         to_np(seed_trajs[:, i, 1]),
    #         c=to_np(ts[frm:to_seed, i, 0]),
    #         cmap=cm.plasma,
    #     )
    #     ax.plot(to_np(orig_trajs[frm:to, i, 0]), to_np(orig_trajs[frm:to, i, 1]))
    #     ax.plot(samp_trajs_p[:, i, 0], samp_trajs_p[:, i, 1])
    # plt.show()
    #
    # print(np.mean(losses), np.median(losses))
    #
    # scaler = StandardScaler()
    #
    # scaler.fit(train_data)
    # train_scaled = scaler.transform(train_data)
    #
    # for i in range(len(data_test)):
    #     data_test[i] = scaler.transform(data_test[i])


if __name__ == "__main__":
    main()