Code for reproducing results from the paper "Unified Long Horizon Time Series Benchmark"
import model
import trainers
import utils
from model import LTSF, FullModel, LSTMModel, SpaceTime
from model.components.mlp import MLP
from model.components.encoders import (
    RecognitionWaveNet,
    MultIdEnc,
    LSTMEnc,
)
from model.components.decoders import IdDec
import problem_specs.instances as problem_spec_instances  # noqa
from model.solvers import ODEINTWEvents


def build_wavenet_enc(enc):
    return RecognitionWaveNet(
        latent_dim=enc["latent_dim"],
        obs_dim=enc["obs_dim"],
        nfilters=enc["filter_dim"],
        depth=enc["depth"],
        activation=enc["activation"],
        use_batchnorm=enc["use_batchnorm"],
    )


def build_mult_wavenet_enc(enc):
    return RecognitionWaveNet(
        latent_dim=enc["latent_dim"],
        obs_dim=enc["obs_dim"],
        nfilters=enc["filter_dim"],
        depth=enc["depth"],
        activation=enc["activation"],
        additional_convs_per_layer=enc["additional_convs_per_layer"],
        encoded_len=enc["encoded_len"],
    )


def build_mult_id_enc(enc):
    return MultIdEnc(enc["encoded_len"])


def build_lstm_enc(enc):
    return LSTMEnc(
        obs_dim=enc["obs_dim"],
        latent_dim=enc["latent_dim"],
        hidden_dim=enc["hidden_dim"],
        depth=enc["depth"],
        encoded_len=enc["encoded_len"],
    )


def build_enc(enc):
    if enc["name"] == "WaveNetEnc":
        return build_wavenet_enc(enc)
    elif enc["name"] == "MultWaveNetEnc":
        return build_mult_wavenet_enc(enc)
    elif enc["name"] == "MultIdEnc":
        return build_mult_id_enc(enc)
    elif enc["name"] == "LSTMEnc":
        return build_lstm_enc(enc)
    else:
        raise NotImplementedError(enc)


def build_mlp(latent_ode):
    return MLP(
        layers=latent_ode["layers"],
        activation=latent_ode["activation"],
        use_batchnorm=latent_ode["use_batchnorm"],
        add_last=latent_ode.get("add_last", False),
    )


def build_lstm(latent_ode):
    return LSTMModel(
        latent_dim=latent_ode["latent_dim"],
        hidden_dim=latent_ode["width"],
        depth=latent_ode["depth"],
    )


def build_latent_ode(latent_ode):
    if latent_ode["type"] == "mlp":
        return build_mlp(latent_ode)
    elif latent_ode["type"] == "lstm":
        return build_lstm(latent_ode)
    else:
        raise NotImplementedError(latent_ode)


def build_id_dec(dec):
    return IdDec()


def build_dec(dec):
    if dec["type"] == "mlp":
        return build_mlp(dec)
    elif dec["type"] == "IdDec":
        return build_id_dec(dec)
    else:
        raise NotImplementedError(dec)


def build_standard_node_model(model):
    return {
        "type": "standard_node",
        "enc": build_enc(model["enc"]).to(utils.DEFAULT_DEVICE),
        "latent_ode": build_latent_ode(model["latent_ode"]).to(
            utils.DEFAULT_DEVICE
        ),
        "dec": build_dec(model["dec"]).to(utils.DEFAULT_DEVICE),
    }


def build_latent_ltsf_linear_model(model):
    return {
        "type": "latent_ltsf_linear",
        "enc": build_enc(model["enc"]).to(utils.DEFAULT_DEVICE),
        "dec": build_dec(model["dec"]).to(utils.DEFAULT_DEVICE),
        "latent_dim": model["latent_dim"],
    }


def build_latent_ltsf_nlinear_model(model):
    return {
        "type": "latent_ltsf_nlinear",
        "enc": build_enc(model["enc"]).to(utils.DEFAULT_DEVICE),
        "dec": build_dec(model["dec"]).to(utils.DEFAULT_DEVICE),
        "latent_dim": model["latent_dim"],
        "N": model["N"],
    }


def build_lstm_model(model):
    return model


def build_spacetime_model(model):
    return model


def build_nhits_model(model):
    return model


def build_model(model):
    if model["type"] == "standard_node":
        return build_standard_node_model(model)
    elif model["type"] == "latent_ltsf_linear":
        return build_latent_ltsf_linear_model(model)
    elif model["type"] == "latent_ltsf_nlinear":
        return build_latent_ltsf_nlinear_model(model)
    elif model["type"] == "lstm":
        return build_lstm_model(model)
    elif model["type"] == "spacetime":
        return build_spacetime_model(model)
    else:
        raise ValueError(f"unknown model type: {model}")


def build_from_configuration(configuration, checkpoint_name):
    model = build_model(
        configuration["model"],
    )
    problem_spec = eval(
        "problem_spec_instances.{}".format(configuration["problem"])
    )()
    trainer_class = trainers.ODEINTWEventsTrainer

    if model["type"] in ["standard_node", "standard_su_node"]:
        solver = ODEINTWEvents(
            default_step_size=configuration["trainer_args"]["step_size"]
        )
        full_model = FullModel(
            model_type=model["type"],
            obs_dim=configuration["obs_dim"],
            latent_dim=configuration["latent_dim"],
            rec=model["enc"],
            dec=model["dec"],
            latent_ode=model["latent_ode"],
            T=configuration.get("T"),
            solver=solver,
        )
    elif model["type"] == "latent_ltsf_linear":
        full_model = LTSF(
            rec=model["enc"],
            dec=model["dec"],
            latent_dim=model["latent_dim"],
            out_values=configuration["trainer_args"]["train_seq_len"] - model["N"],
            ltsf_type="linear",
            N=model["N"],
        )
    elif model["type"] == "latent_ltsf_nlinear":
        full_model = LTSF(
            rec=model["enc"],
            dec=model["dec"],
            latent_dim=model["latent_dim"],
            out_values=configuration["trainer_args"]["train_seq_len"] - model["N"],
            ltsf_type="nlinear",
            N=model["N"],
        )
    elif model["type"] == "lstm":
        full_model = LSTMModel(
            obs_dim=model["obs_dim"],
            hidden_dim=model["hidden_dim"],
            depth=model["depth"],
            separate_lstm_encoder=model["separate_lstm_encoder"],
            encoded_len=model["encoded_len"],
        )
    elif model["type"] == "spacetime":
        full_model = SpaceTime(
            obs_dim=model["obs_dim"],
            embedding_dim=model["embedding_dim"],
            encoder_dim=model["encoder_dim"],
            encoded_len=model["encoded_len"],
            pred_len=configuration["trainer_args"]["train_seq_len"] - model["encoded_len"],
            hidden_dim=model["hidden_dim"],
            encoder_layers=model["encoder_layers"],
            decoder_layers=model["decoder_layers"],
        )
    else:
        raise ValueError(f"unknown model type: {model['type']}")

    full_model = full_model.to(utils.DEFAULT_DEVICE)

    count_parameters = 0

    for param in full_model.parameters():
        count_parameters += param.contiguous().view(-1).shape[0]

    print(f"\nModel has {count_parameters} parameters\n")

    trainer = trainer_class(
        problem_spec=problem_spec,
        model=full_model,
        problem_spec_name=configuration["problem"],
        adjoint=configuration["trainer_args"].get("adjoint"),
        train_seq_len=configuration["trainer_args"]["train_seq_len"],
        batch_size=configuration["trainer_args"]["batch_size"],
        max_lr=configuration["trainer_args"]["max_lr"],
        trainer_name="su_ode",
        train_dir="./models_cpth/{}/{}".format(
            configuration["problem"],
            checkpoint_name,
        ),
        configuration=configuration,
        solver_method=configuration["trainer_args"]["solver_method"],
        pretraining_epochs=configuration["trainer_args"]["pretraining_epochs"],
        increasing_seq_lens=configuration["trainer_args"][
            "increasing_seq_lens"
        ],
        trajectories_per_epoch=configuration["trainer_args"][
            "trajectories_per_epoch"
        ],
        betas=configuration["trainer_args"]["betas"],
        increasing_seq_lens_epochs=configuration["trainer_args"][
            "increasing_seq_lens_epochs"
        ],
        test_train_data_len=configuration["trainer_args"][
            "test_train_data_len"
        ],
        test_data_len=configuration["trainer_args"]["test_data_len"],
        entropy_loss=configuration["trainer_args"]["entropy_loss"],
        atol=configuration["trainer_args"]["atol"],
        rtol=configuration["trainer_args"]["rtol"],
        min_dt=configuration["trainer_args"]["min_dt"],
        entropy_factor=configuration["trainer_args"]["entropy_factor"],
        seq_len_step=configuration["trainer_args"]["seq_len_step"],
        pretraining_len=configuration["trainer_args"]["pretraining_len"],
        train_enc_dec_after_pretraining=configuration["trainer_args"][
            "train_enc_dec_after_pretraining"
        ],
        pretraining_ode_lr=configuration["trainer_args"]["pretraining_ode_lr"],
        checkpoint_name=checkpoint_name,
        aug_step=configuration["trainer_args"]["aug_step"],
        optimizer=configuration["trainer_args"]["optimizer"],
        sample_vae=configuration["trainer_args"]["sample_vae"],
        add_kl=configuration["trainer_args"]["add_kl"],
    )

    return trainer

src/build_from_configuration.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
# flake8: noqa
from .configurations_ode import *

src/configurations/__init__.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
def mlp_layers(in_dim, out_dim, width, depth):
    if depth < 0:
        raise ValueError("")
    elif depth == 0:
        return [
            (
                in_dim,
                out_dim,
            )
        ]
    else:
        layers = [
            (
                in_dim,
                width,
            )
        ]
        layers += [(width, width) for _ in range(depth - 1)]
        layers.append(
            (
                width,
                out_dim,
            )
        )

        return layers


def mlp(
    in_dim, out_dim, width, depth, activation, resnet_style, use_batchnorm
):
    return {
        "type": "mlp",
        "layers": mlp_layers(in_dim, out_dim, width, depth),
        "activation": activation,
        "use_batchnorm": use_batchnorm,
    }


def mlp_linear(latent_dim):
    return {
        "type": "linear",
        "latent_size": latent_dim,
    }


def mlp_const(latent_dim):
    return {
        "type": "const",
        "latent_size": latent_dim,
    }


def mlp_zero(latent_dim):
    return {
        "type": "zero",
        "latent_size": latent_dim,
    }


def rnn(latent_dim, width, depth):
    return {
        "type": "rnn",
        "latent_dim": latent_dim,
        "width": width,
        "depth": depth,
    }


def lstm(latent_dim, width, depth):
    return {
        "type": "lstm",
        "latent_dim": latent_dim,
        "width": width,
        "depth": depth,
    }


def ltsf_linear(latent_dim, lookback):
    return {
        "type": "ltsf_linear",
        "latent_dim": latent_dim,
        "lookback": lookback,
    }


def ltsf_nlinear(latent_dim, lookback):
    return {
        "type": "ltsf_nlinear",
        "latent_dim": latent_dim,
        "lookback": lookback,
    }


def ltsf_dlinear(latent_dim, lookback):
    return {
        "type": "ltsf_dlinear",
        "latent_dim": latent_dim,
        "lookback": lookback,
    }


def classifier(
    latent_dim, width, depth, activation, resnet_style, use_batchnorm
):
    return {
        "layers": mlp_layers(latent_dim, 1, width, depth),
        "activation": activation,
        "use_batchnorm": use_batchnorm,
    }


def wavenet_enc(filter_dim, latent_dim, obs_dim, depth, activation):
    return {
        "name": "WaveNetEnc",
        "filter_dim": filter_dim,
        "depth": depth,
        "activation": activation,
        "latent_dim": latent_dim,
        "obs_dim": obs_dim,
    }


def mult_wavenet_enc(
    filter_dim,
    latent_dim,
    obs_dim,
    depth,
    additional_convs_per_layer,
    activation,
    n_out,
    encoded_len,
):
    return {
        "name": "MultWaveNetEnc",
        "filter_dim": filter_dim,
        "depth": depth,
        "activation": activation,
        "latent_dim": latent_dim,
        "obs_dim": obs_dim,
        "n_out": n_out,
        "additional_convs_per_layer": additional_convs_per_layer,
        "encoded_len": encoded_len,
    }
    

def lstm_enc(
    hidden_dim,
    latent_dim,
    obs_dim,
    depth,
    n_out,
    encoded_len,
):
    return {
        "name": "LSTMEnc",
        "hidden_dim": hidden_dim,
        "latent_dim": latent_dim,
        "obs_dim": obs_dim,
        "depth": depth,
        "n_out": n_out,
        "encoded_len": encoded_len,
    }


def mult_id_enc(encoded_len):
    return {
        "name": "MultIdEnc",
        "encoded_len": encoded_len,
    }


def id_dec():
    return {
        "type": "IdDec",
    }

src/configurations/common.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
from problem_specs.instances import get_obs_dim

from . import model_instances  # noqa

def clconfig_1000(**kwargs):
    latent_dim = kwargs["latent_dim"]
    problem_spec_name = kwargs["problem_spec_name"]
    model_instance_name = kwargs["model_instance"]
    obs_dim = get_obs_dim(problem_spec_name)
    enc_dim = kwargs["enc_dim"]
    hidden_dim = kwargs["hidden_dim"]
    pretraining_len = kwargs["pretraining_len"]
    dec_width = kwargs["dec_width"]
    dec_depth = kwargs["dec_depth"]
    model = eval(
        f"model_instances.{model_instance_name}(\
        latent_dim={latent_dim},\
        obs_dim={obs_dim},\
        enc_dim={enc_dim},\
        hidden_dim={hidden_dim},\
        dec_depth={dec_depth},\
        dec_width={dec_width},\
    )"
    )
    batch_size = kwargs["batch_size"]
    max_lr = 0.001
    return {
        "problem": problem_spec_name,
        "trainer": "su_ode",
        "trainer_args": {
            "batch_size": batch_size,
            "max_lr": max_lr,
            "pretraining_ode_lr": max_lr,
            "epochs": 100,
            "step_size": 0.01 / 4,
            "min_dt": 0.01 / 32,
            "train_seq_len": 1000,
            "solver_method": "dopri5",
            "pretraining_epochs": 10,
            "increasing_seq_lens": True,
            "trajectories_per_epoch": 18000,
            "increasing_seq_lens_epochs": 1,
            "betas": (0.9, 0.999),
            "test_train_data_len": 2000,
            "test_data_len": 100000000,
            "entropy_loss": False,
            "atol": 1e-3,
            "rtol": 1e-2,
            "entropy_factor": 0.2,
            "seq_len_step": 50,
            "pretraining_len": pretraining_len,
            "train_enc_dec_after_pretraining": True,
            "aug_step": 100,
            "optimizer": "Adam",
            "add_kl": False,
            "sample_vae": False,
        },
        "latent_dim": latent_dim,
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "model": model,
        "T": 0.2,
    }

def clconfig_1000_lesstraj(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["trajectories_per_epoch"] = 3000
    cf["trainer_args"]["pretraining_epochs"] = 5
    return cf

def fullconfig_1000(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["pretraining_epochs"] = 0
    cf["trainer_args"]["increasing_seq_lens"] = False
    cf["trainer_args"]["train_enc_dec_after_pretraining"] = True
    return cf

def clconfig_2000(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 2000
    cf["trainer_args"]["seq_len_step"] = 100
    cf["trainer_args"]["pretraining_epochs"] = 5
    return cf

def clconfig_500(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 500
    return cf

def clconfig_288(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 288
    return cf

def fullconfig_288(**kwargs):
    cf = fullconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 288
    return cf

def clconfig_200(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 200
    cf["trainer_args"]["max_lr"] = 0.001
    return cf

def fullconfig_500(**kwargs):
    cf = fullconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 500
    return cf

def fullconfig_1440(**kwargs):
    cf = fullconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 1440
    return cf

def clconfig_1440(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 1440
    return cf

def fullconfig_2000(**kwargs):
    cf = fullconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 2000
    return cf

def fullconfig_200(**kwargs):
    cf = fullconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 200
    return cf

def clconfig_192(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 192
    return cf

def fullconfig_192(**kwargs):
    cf = fullconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 192
    return cf

def clconfig_247(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 247
    return cf

def fullconfig_247(**kwargs):
    cf = fullconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 247
    return cf

def clconfig_228(**kwargs):
    cf = clconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 228
    return cf

def fullconfig_228(**kwargs):
    cf = fullconfig_1000(**kwargs)
    cf["trainer_args"]["train_seq_len"] = 228
    return cf
src/configurations/configurations_ode.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
from .common import id_dec, mult_id_enc, mult_wavenet_enc, mlp, classifier, lstm_enc


def standard_node_1(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    assert hidden_dim is not None
    return {
        "type": "standard_node",
        # "enc": mult_wavenet_enc(enc_dim, latent_dim, obs_dim, 5, 0, "ELU", 1),
        "enc": lstm_enc(enc_dim, latent_dim, obs_dim, 2, 1, 32),
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            max(latent_dim, obs_dim) * 4,
            4,
            "ELU",
            True,
            False,
        ),
    }

# ===ablation definitions===

SMALLDEC2 = 200
ENCLAYERS = 4

#===lstm encoders latent ode===

def latentode_1_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            1,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_2_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            2,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_8_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            8,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_96_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            96,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_125_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            125,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_168_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            168,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_250_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            250,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_325_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            325,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }
    
def latentode_336_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            336,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_425_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            425,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_500_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            500,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_720_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            720,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def latentode_1000_lstm(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "standard_node",
        "latent_ode": mlp(
            latent_dim, latent_dim, hidden_dim, 3, "ELU", True, False
        ),
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            1000,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
    }

def obs_ltsf_nlinear_144(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(144),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 144,
    }
    

def obs_ltsf_nlinear_168(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(168),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 168,
    }

def obs_ltsf_nlinear_500(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(500),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 500,
    }

def obs_ltsf_nlinear_1(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(1),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 1,
    }

def obs_ltsf_nlinear_2(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(2),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 2,
    }

def obs_ltsf_nlinear_4(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(4),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 4,
    }


def obs_ltsf_nlinear_8(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(8),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 8,
    }

def obs_ltsf_nlinear_16(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(16),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 16,
    }

def obs_ltsf_nlinear_32(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(32),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 32,
    }

def obs_ltsf_nlinear_96(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(96),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 96,
    }

def obs_ltsf_nlinear_250(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(250),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 250,
    }

def obs_ltsf_nlinear_336(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(336),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 336,
    }

def obs_ltsf_nlinear_720(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(720),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 720,
    }

def obs_ltsf_nlinear_1000(obs_dim, enc_dim, **kwargs):
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": mult_id_enc(1000),
        "dec": id_dec(),
        "latent_dim": obs_dim,
        "N": 1000,
    }


def lstm_1(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 1,
        "separate_lstm_encoder": False,
    }


def lstm_2(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 2,
        "separate_lstm_encoder": True,
    }

def lstm_3_1(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 1,
    }

def lstm_3_2(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 2,
    }

def lstm_3_8(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 8,
    }

def lstm_3_96(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 96,
    }

def lstm_3_144(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 144,
    }

def lstm_3_168(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 168,
    }

def lstm_3_250(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 250,
    }

def lstm_3_336(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 336,
    }

def lstm_3_500(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 500,
    }

def lstm_3_720(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 720,
    }

def lstm_3_1000(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": False,
        "encoded_len": 1000,
    }

def lstm_4(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 4,
        "separate_lstm_encoder": False,
    }


def lstm_5(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 5,
        "separate_lstm_encoder": False,
    }


def lstm_6(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 6,
        "separate_lstm_encoder": False,
    }


def lstm_7(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 7,
        "separate_lstm_encoder": False,
    }


def dlstm_1(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 1,
        "separate_lstm_encoder": True,
        "encoded_len": 75,
    }


def dlstm_2(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 2,
        "separate_lstm_encoder": True,
    }


def dlstm_3(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 75,
    }

def dlstm_3_1(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 1,
    }

def dlstm_3_2(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 2,
    }

def dlstm_3_8(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 8,
    }

def dlstm_3_96(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 96,
    }

def dlstm_3_144(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 144,
    }

def dlstm_3_168(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 168,
    }

def dlstm_3_250(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 250,
    }

def dlstm_3_336(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 336,
    }

def dlstm_3_500(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 500,
    }

def dlstm_3_720(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 720,
    }

def dlstm_3_1000(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 3,
        "separate_lstm_encoder": True,
        "encoded_len": 1000,
    }

def dlstm_4(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 4,
        "separate_lstm_encoder": True,
    }


def dlstm_5(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 5,
        "separate_lstm_encoder": True,
    }


def dlstm_6(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 6,
        "separate_lstm_encoder": True,
    }


def dlstm_7(obs_dim, hidden_dim, **kwargs):
    assert obs_dim is not None
    assert hidden_dim is not None
    return {
        "type": "lstm",
        "obs_dim": obs_dim,
        "hidden_dim": hidden_dim,
        "depth": 7,
        "separate_lstm_encoder": True,
    }

### === latent LTSF begin ===

def latent_ltsf_nlinear_1(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            1,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 1,
    }

def latent_ltsf_nlinear_2(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            2,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 2,
    }

def latent_ltsf_nlinear_8(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            8,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 8,
    }

def latent_ltsf_nlinear_96(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            96,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 96,
    }

def latent_ltsf_nlinear_144(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            144,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 144,
    }

def latent_ltsf_nlinear_168(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            168,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 168,
    }

def latent_ltsf_nlinear_250(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            250,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 250,
    }

def latent_ltsf_nlinear_336(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            336,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 336,
    }

def latent_ltsf_nlinear_500(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            500,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 500,
    }

def latent_ltsf_nlinear_720(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            720,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 720,
    }

def latent_ltsf_nlinear_1000(latent_dim, obs_dim, enc_dim, **kwargs):
    assert latent_dim is not None
    assert obs_dim is not None
    assert enc_dim is not None
    return {
        "type": "latent_ltsf_nlinear",
        "enc": lstm_enc(
            enc_dim,
            latent_dim,
            obs_dim,
            2,
            1,
            1000,
        ),
        "dec": mlp(
            latent_dim,
            obs_dim,
            SMALLDEC2,
            ENCLAYERS,
            "ELU",
            True,
            False,
        ),
        "latent_dim": latent_dim,
        "N": 1000,
    }

### === latent LTSF end ===

def spacetime_test(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 75,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_1(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 1,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_2(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 2,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_8(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 8,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_96(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 96,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_144(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 144,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_168(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 168,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_250(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 250,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_336(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 336,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_500(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 500,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_720(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 720,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_1000(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 1000,
        "hidden_dim": hidden_dim,
        "encoder_layers": 2,
        "decoder_layers": 2,
    }

def spacetime_500_3layer(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 500,
        "hidden_dim": hidden_dim,
        "encoder_layers": 3,
        "decoder_layers": 3,
    }

def spacetime_1000_3layer(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "spacetime",
        "obs_dim": obs_dim,
        "embedding_dim": hidden_dim,
        "encoder_dim": hidden_dim,
        "encoded_len": 1000,
        "hidden_dim": hidden_dim,
        "encoder_layers": 3,
        "decoder_layers": 3,
    }

def nhits_test(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 96,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 256,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_1(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 1,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 256,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_2(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 2,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 256,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_8(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 8,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 256,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_96(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 96,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 128,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_144(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 144,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 128,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_168(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 168,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 128,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_250(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 250,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 128,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_336(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 336,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 128,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_500(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 500,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 128,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_720(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 720,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 128,
        "n_freq_downsample": [168, 24, 1],
    }

def nhits_1000(latent_dim, obs_dim, enc_dim, hidden_dim, **kwargs):
    return {
        "type": "nhits",
        "obs_dim": obs_dim,
        "encoded_len": 1000,
        "n_stacks": 3,
        "n_blocks": 3,
        "n_hidden_layers_in_block": 2,
        "hidden_dim": 128,
        "n_freq_downsample": [168, 24, 1],
    }
src/configurations/model_instances.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import os
import gdown


def main():
    gdown.download_folder(
        "https://drive.google.com/drive/folders/1IfAHpka3hu2kM4j6ebzAPnGpLpSMnTlf?usp=share_link",  # noqa
        quiet=False,
        use_cookies=False,
    )
    os.system("mv Final_hdf5 data")


if __name__ == "__main__":
    main()

src/download_data.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import os

import gdown


def main():
    gdown.download_folder(
        "https://drive.google.com/drive/folders/172QgaXtOIPfPiI9b_anLQuuDRR-ZMBWj?usp=share_link",  # noqa
        quiet=False,
        use_cookies=False,
    )
    os.system("mkdir data/cheetah_medium_deter/")
    os.system("mv cheetah_medium_deter/* data/cheetah_medium_deter/")

    gdown.download_folder(
        "https://drive.google.com/drive/folders/1n6D3krCd6rDvQlOvb-z317LGl9qAZsB7?usp=share_link",  # noqa
        quiet=False,
        use_cookies=False,
    )
    os.system("mkdir data/hopper_medium_deter/")
    os.system("mv hopper_medium_deter/* data/hopper_medium_deter/")

    gdown.download_folder(
        "https://drive.google.com/drive/folders/16idSzUl35L33mFzzrx3MvQ2Yxkn9nnln?usp=share_link",  # noqa
        quiet=False,
        use_cookies=False,
    )
    os.system("mkdir data/walker_medium_deter/")
    os.system("mv walker_medium_deter/* data/walker_medium_deter/")


if __name__ == "__main__":
    main()

src/download_data_subset.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import os

import gdown


def main():
    gdown.download_folder(
        "https://drive.google.com/drive/folders/1u3HN75OCSOl-V7VFdkDCkkQTqMGB5j7K?usp=share_link",  # noqa
        quiet=False,
        use_cookies=False,
    )
    os.system("mkdir data/mackey_glass/")
    os.system("mv mackey_glass/* data/mackey_glass/")

    gdown.download_folder(
        "https://drive.google.com/drive/folders/1xrgVk9dN0IfFUPH7pGnYo1vbuXcpoDt5?usp=share_link",  # noqa
        quiet=False,
        use_cookies=False,
    )
    os.system("mkdir data/lorenz/")
    os.system("mv lorenz/* data/lorenz/")

    gdown.download_folder(
        "https://drive.google.com/drive/folders/1EvrTrh2uvZ4RrmLluzFM1-sQk_9y-hfq?usp=share_link",  # noqa
        quiet=False,
        use_cookies=False,
    )
    os.system("mkdir data/mso/")
    os.system("mv mso/* data/mso/")

    gdown.download_folder(
        "https://drive.google.com/drive/folders/1fkWPXvizNylEPbTApOUx_hk2Eczuruy-?usp=share_link",  # noqa
        quiet=False,
        use_cookies=False,
    )
    os.system("mkdir data/sspiral/")
    os.system("mv sspiral/* data/sspiral/")



if __name__ == "__main__":
    main()

src/download_data_subset2.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
# flake8: noqa
from .ltsf import LTSF
from .full_model import FullModel
from .lstm import LSTMModel
from .spacetime import SpaceTime

src/model/__init__.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import torch
from torch import nn


class LinearDecoder(nn.Module):
    def __init__(self, latent_dim, obs_dim):
        super(LinearDecoder, self).__init__()
        self.bn_in = nn.BatchNorm1d(latent_dim)
        self.fc = nn.Linear(latent_dim, obs_dim)

    def forward(self, z):
        return self.fc(z)


class LinearLSTMDecoder(nn.Module):
    def __init__(self, latent_dim, obs_dim):
        super(LinearLSTMDecoder, self).__init__()
        self.bn_in = nn.BatchNorm1d(latent_dim)
        self.fc = nn.Linear(latent_dim, obs_dim)

    def forward(self, z):
        return self.fc(z)


class VideoDecoder(nn.Module):
    def __init__(self, latent_dim, obs_dim, nhidden, d):
        super(VideoDecoder, self).__init__()
        self.tanh = torch.tanh
        self.bn0 = nn.BatchNorm1d(latent_dim)
        self.fc1 = nn.Linear(latent_dim, nhidden)
        self.bn1 = nn.BatchNorm1d(nhidden)
        self.fcs = nn.ModuleList(
            [nn.Linear(nhidden, nhidden) for _ in range(d)]
        )
        self.bns = nn.ModuleList([nn.BatchNorm1d(nhidden) for _ in self.fcs])

        self.ct1 = nn.ConvTranspose2d(
            in_channels=nhidden,
            out_channels=obs_dim[2],
            kernel_size=(
                obs_dim[0],
                obs_dim[1],
            ),
        )

    def forward(self, z):
        out = self.bn0(z)
        out = self.fc1(out)
        out = self.bn1(out)
        for fc, bn in zip(self.fcs, self.bns):
            t = self.tanh(fc(out))
            out = t + out
            out = bn(out)

        out = torch.reshape(
            out,
            (
                out.shape[0],
                out.shape[1],
                1,
                1,
            ),
        )
        out = self.ct1(out)

        out = torch.reshape(
            out,
            (
                out.shape[0],
                out.shape[1],
                out.shape[2],
                out.shape[3],
            ),
        )
        out = out.permute(
            0,
            2,
            3,
            1,
        )

        return out


class TimeDistributedDec(nn.Module):
    def __init__(self, module):
        super(TimeDistributedDec, self).__init__()
        self.module = module

    def forward(self, x):
        x_reshape = x.contiguous().view(-1, x.size(-1))
        y = self.module(x_reshape)
        y = y.contiguous().view(x.size(0), x.size(1), y.size(-1))

        return y


class IdDec(nn.Module):
    def forward(self, x):
        return x
    
    def set_freeze(self, freeze):
        pass

src/model/components/decoders.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import torch
from torch import nn

import problem_specs
import utils
from model.solvers import ODEINTWEvents


class RecognitionRNN(nn.Module):
    def __init__(self, latent_dim, obs_dim, nhidden, d):
        super(RecognitionRNN, self).__init__()
        self.nhidden = nhidden
        self.i2h = nn.Linear(obs_dim + nhidden, nhidden)
        self.i2hs = nn.ModuleList(
            [nn.Linear(nhidden, nhidden) for _ in range(d)]
        )
        self.h2o = nn.Linear(nhidden, latent_dim)

    def forward(self, x, h):
        h = torch.cat((x, h), dim=1)
        h = torch.tanh(self.i2h(h))
        for i2 in self.i2hs:
            t = torch.tanh(i2(h))
            h = t + h
        out = self.h2o(h)
        return out, h

    def initHidden(self, nbatch):
        return torch.zeros(nbatch, self.nhidden)

    def recognize(self, device, dataset, rec_num):
        h = self.initHidden(dataset.size(0)).to(device)

        for t in range(rec_num):
            obs = dataset[:, rec_num - t - 1, :]
            out, h = self.forward(obs, h)

        return out


class MultIdEnc(nn.Module):
    def __init__(self, encoded_len):
        super(MultIdEnc, self).__init__()
        self.encoded_len = encoded_len
        self._rec_len = encoded_len

    def recognize(self, tsdata, needed_last):
        means = tsdata.dataset[
            :, self.encoded_len - needed_last : self.encoded_len
        ]
        return means, None, None, None

    def set_freeze(self, freeze):
        pass


class ResNetBlock1d(nn.Module):
    def __init__(self, nfilters, activation):
        super(ResNetBlock1d, self).__init__()


        self.conv1 = nn.Conv1d(nfilters, nfilters, 1)
        self.conv2 = nn.Conv1d(nfilters, nfilters, 1)
        self.activation = activation

    def forward(self, x):
        old_x = x
        x = self.conv1(x)
        x = self.activation(x)
        x = self.conv2(x)

        return self.activation(old_x + x)


class RecognitionWaveNet(nn.Module):
    def __init__(
        self,
        latent_dim,
        obs_dim,
        nfilters,
        depth,
        additional_convs_per_layer,
        encoded_len,
        activation,
    ):
        super(RecognitionWaveNet, self).__init__()

        assert (
            encoded_len <= 2**depth
        ), "TCN encoder must have its receptive field >= encoded_len"

        self.latent_dim = latent_dim
        self.depth = depth
        self._rec_len = encoded_len

        if activation == "ELU":
            self.activation = nn.ELU()
        else:
            raise NotImplementedError(
                "unknown activation: {}".format(activation)
            )

        self.conv1d_first = nn.Sequential(
            # nn.BatchNorm1d(obs_dim),
            nn.Conv1d(
                in_channels=obs_dim,
                out_channels=nfilters,
                kernel_size=2,
                dilation=1,
            ),
            # nn.BatchNorm1d(nfilters),
            self.activation,
            *[
                nn.Sequential(ResNetBlock1d(nfilters, self.activation))
                for _ in range(additional_convs_per_layer)
            ]
        )
        # self.bn_first = nn.BatchNorm1d(nfilters)
        self.conv1ds_middle = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Sequential(
                        nn.Conv1d(
                            in_channels=nfilters,
                            out_channels=nfilters,
                            kernel_size=2,
                            dilation=2**x,
                        ),
                        self.activation,
                    ),
                    *[
                        nn.Sequential(ResNetBlock1d(nfilters, self.activation))
                        for _ in range(additional_convs_per_layer)
                    ]
                )
                for x in range(1, depth - 1)
            ]
        )
        self.conv1d_last = nn.Conv1d(
            in_channels=nfilters,
            out_channels=latent_dim
            * 2,  # + latent_dim * latent_dim + latent_dim,
            kernel_size=2,
            dilation=2 ** (depth - 1),
        )

        self.out_param = nn.Parameter(
            torch.Tensor(encoded_len, latent_dim * latent_dim + latent_dim, 1)
        )
        torch.nn.init.xavier_normal_(self.out_param)

    def set_freeze(self, freeze):
        for param in self.conv1d_first.parameters():
            param.requires_grad = not freeze
        for param in self.conv1ds_middle.parameters():
            param.requires_grad = not freeze
        for param in self.conv1d_last.parameters():
            param.requires_grad = not freeze

    def forward(self, x, needed_last):
        if needed_last > 1:
            additional_pad = torch.zeros_like(x[:, :, : needed_last - 1]) + utils.DEFAULT_MISSING_VALUE
            x = torch.cat([additional_pad, x], dim=2)
        x = self.conv1d_first(x)

        for conv1d in self.conv1ds_middle:
            x = conv1d(x)

        seq_state = self.conv1d_last(x)

        return seq_state, self.out_param[-needed_last:]

    def _check_input(self, tsdata):
        assert isinstance(tsdata, problem_specs.TSData)

        timestamps = tsdata.timestamps[:, : self.rec_len]
        timestamp_firsts = timestamps[:, 0].view(-1, 1).repeat(1, self.rec_len)
        timestamps_norm = timestamps - timestamp_firsts
        timestamps_example = timestamps_norm[0]

        assert torch.all(timestamps_example > timestamps_norm - 0.00001)
        assert torch.all(timestamps_example < timestamps_norm + 0.00001)

        diffs = timestamps_norm[:, 1:] - timestamps_norm[:, :-1]
        diffs = diffs.view(-1)

        assert torch.all(diffs > 0.01 - 0.00001)
        assert torch.all(diffs < 0.01 + 0.00001)

    def recognize(self, tsdata, needed_last):
        recognizable_dataset = tsdata.dataset[:, : self.rec_len].permute(
            0, 2, 1
        )
        if self._rec_len < 2**self.depth:
            pad = torch.zeros_like(recognizable_dataset)[
                :, :, : 2**self.depth - self._rec_len
            ]
            pad += utils.DEFAULT_MISSING_VALUE
            recognizable_dataset = torch.cat(
                [pad, recognizable_dataset], dim=2
            )

        z0, out_param = self.forward(recognizable_dataset, needed_last)
        latent_dim = self.latent_dim
        # return z0[:, :latent_dim],
        # torch.nn.functional.elu((z0[:, latent_dim:])) + 1.0001
        mean = z0[:, :latent_dim]
        std = z0[:, latent_dim : 2 * latent_dim] + 0.0001
        weight = out_param[:, : latent_dim * latent_dim].view(
            needed_last, latent_dim, latent_dim
        )
        bias = out_param[:, latent_dim * latent_dim :]
        return mean, std, weight, bias

    @property
    def rec_len(self):
        return self._rec_len


class LSTMEnc(nn.Module):
    def __init__(self, obs_dim, latent_dim, hidden_dim, depth, encoded_len):
        super(LSTMEnc, self).__init__()
        self.lstm = nn.LSTM(obs_dim, hidden_dim, depth, batch_first=True)
        self.lstm_proj = nn.Linear(hidden_dim, latent_dim)
        self._rec_len = encoded_len
        self.encoded_len = encoded_len
        self.latent_dim = latent_dim

    def set_freeze(self, freeze):
        for param in self.lstm.parameters():
            param.requires_grad = not freeze

    def forward(self, x):
        bs = x.shape[0]
        x = self.lstm(x)[0]
        x = x.reshape(-1, x.shape[2])
        x = self.lstm_proj(x)
        x = x.view(bs, -1, x.shape[1])

        return x

    def recognize(self, tsdata, needed_last):
        dataset = tsdata.dataset[:, :self._rec_len]
        z0 = self(dataset)
        z0 = z0[:, -needed_last:]
        latent_dim = self.latent_dim
        mean = z0[:, :, :latent_dim]
        return mean, None, None, None
        
src/model/components/encoders.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import torch
from torch import nn


class LTSFLinear(nn.Module):
    def __init__(self, latent_dim, lookback):
        super(LTSFLinear, self).__init__()
        self.latent_dim = latent_dim
        self.lookback = lookback

        self.ll = nn.Linear(latent_dim * lookback, latent_dim)

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        return self.ll(x[:, -self.lookback * self.latent_dim :])


class LTSFNLinear(nn.Module):
    def __init__(self, latent_dim, lookback):
        super(LTSFNLinear, self).__init__()
        self.latent_dim = latent_dim
        self.lookback = lookback

        self.ll = nn.Linear(latent_dim * lookback, latent_dim)

    def forward(self, x):
        last_l = x[:, 0].unsqueeze(1).repeat(1, x.shape[1], 1)
        x = x - last_l
        x = x.view(x.shape[0], -1)
        return self.ll(x[:, -self.lookback * self.latent_dim :]) + last_l[:, 0]


class moving_avg(nn.Module):
    """
    Moving average block to highlight the trend of time series
    """

    def __init__(self, kernel_size, stride):
        super(moving_avg, self).__init__()
        self.kernel_size = kernel_size
        self.avg = nn.AvgPool1d(
            kernel_size=kernel_size, stride=stride, padding=0
        )

    def forward(self, x):
        # padding on the both ends of time series
        front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)
        end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)
        x = torch.cat([front, x, end], dim=1)
        x = self.avg(x.permute(0, 2, 1))
        x = x.permute(0, 2, 1)
        return x


class series_decomp(nn.Module):
    """
    Series decomposition block
    """

    def __init__(self, kernel_size):
        super(series_decomp, self).__init__()
        self.moving_avg = moving_avg(kernel_size, stride=1)

    def forward(self, x):
        moving_mean = self.moving_avg(x)
        res = x - moving_mean
        return res, moving_mean


class LTSFDLinear(nn.Module):
    """
    Decomposition-Linear
    """

    def __init__(self, latent_dim, lookback):
        super(LTSFDLinear, self).__init__()

        kernel_size = 25
        self.decompsition = series_decomp(kernel_size)

        self.latent_dim = latent_dim
        self.lookback = lookback

        self.llt = nn.Linear(latent_dim * lookback, latent_dim)
        self.lls = nn.Linear(latent_dim * lookback, latent_dim)

    def forward(self, x):
        # x: [Batch, Input length, Channel]
        seasonal_init, trend_init = self.decompsition(x)
        trend_init = trend_init.view(trend_init.shape[0], -1)
        seasonal_init = seasonal_init.view(seasonal_init.shape[0], -1)
        trend_output = self.llt(
            trend_init[:, -self.latent_dim * self.lookback :]
        )
        seasonal_output = self.lls(
            seasonal_init[:, -self.latent_dim * self.lookback :]
        )

        x = seasonal_output + trend_output
        return x


class MLP(nn.Module):
    def __init__(self, layers, activation, use_batchnorm, add_last=False):
        super(MLP, self).__init__()

        if use_batchnorm:
            raise NotImplementedError

        if activation == "ELU":
            self.activation = nn.ELU()
        elif activation == "Tanh":
            self.activation = nn.Tanh()
        elif activation == "sigmoid":
            self.activation = nn.Sigmoid()
        else:
            raise NotImplementedError(
                "unknown activation: {}".format(activation)
            )

        self.activation = nn.Tanh()

        self.fcs = nn.ModuleList(
            [
                nn.Linear(
                    in_dim,
                    out_dim,
                )
                for (
                    in_dim,
                    out_dim,
                ) in layers
            ]
        )
        self.fcs_bn = nn.ModuleList(
            [
                nn.BatchNorm1d(in_dim)
                for (
                    in_dim,
                    _,
                ) in layers
            ]
        )

        self.add_last = add_last
        print(add_last)

        self.bias = nn.Parameter(torch.Tensor(layers[-1][1]))

    def set_freeze(self, freeze):
        for param in self.parameters():
            param.requires_grad = not freeze

    def forward(self, x):
        orig_dim = len(x.shape)
        orig_shape = x.shape

        dx = x

        if orig_dim == 3:
            dx = dx.view(-1, orig_shape[2])

        # dx = self.fcs_bn[0](dx)
        for fc, bn in zip(self.fcs[:-1], self.fcs_bn[1:]):
            dx = fc(dx)
            # dx = bn(dx)
            if dx.size(-1) == fc.out_features:
                ddx = self.activation(dx)
                dx = dx + ddx
            else:
                dx = self.activation(dx)

        dx = self.fcs[-1](dx)

        if orig_dim == 3:
            dx = dx.view(orig_shape[0], orig_shape[1], -1)

        if self.add_last:
            return x + dx / 100
        else:
            return dx

src/model/components/mlp.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import torch
from torch import nn

import problem_specs
import utils


class FullModel(nn.Module):
    def __init__(
        self,
        model_type,
        obs_dim,
        latent_dim,
        rec,
        dec,
        latent_ode,
        T,
        solver,
    ):
        super(FullModel, self).__init__()
        self.obs_dim = obs_dim
        self.latent_dim = latent_dim

        if T is not None:
            self.T = T

        self.latent_ode = latent_ode
        self.rec = rec
        self.dec = dec

        self.model_type = model_type
        self.solver = solver

    @property
    def rec_len(self):
        return self.rec.rec_len

    def forward(
        self,
        tsdata,
        method,
        atol,
        rtol,
        min_dt,
        sample_vae,
        discrete_updates,
    ):
        assert isinstance(tsdata, problem_specs.TSData)
        z0_mean, z0_std, _, _ = self.rec.recognize(tsdata, needed_last=1)
        encoded_observations = tsdata.dataset[:, :self.rec.encoded_len]
        pred_z, su_probs, steps = self.forward_solver(
            samp_ts=tsdata.timestamps,
            z0_mean=z0_mean,
            z0_std=z0_std,
            method=method,
            atol=atol,
            rtol=rtol,
            min_dt=min_dt,
            sample_vae=sample_vae,
            discrete_updates=discrete_updates,
        )

        pred_x = torch.cat([encoded_observations, self.dec(pred_z)], dim=1)
        obs_flag = tsdata.obs_flag
        pred_x[obs_flag == 0.0] = utils.DEFAULT_MISSING_VALUE

        return {
            "pred_x": pred_x,
            "su_probs": su_probs,
            "z0_mean": z0_mean,
            "z0_std": z0_std,
            "steps": steps,
        }

    def predict(
        self,
        tsdata,
        discrete_updates,
        method,
        atol,
        rtol,
        min_dt,
        sample_vae,
    ):
        self.eval()
        with torch.no_grad():
            ret_val = self.forward(
                tsdata=tsdata,
                discrete_updates=discrete_updates,
                method=method,
                atol=atol,
                rtol=rtol,
                min_dt=min_dt,
                sample_vae=sample_vae,
            )
            return ret_val

    def forward_solver(
        self,
        samp_ts,
        z0_mean,
        z0_std,
        method,
        atol,
        rtol,
        min_dt,
        sample_vae,
        discrete_updates,
        adjoint=False,
    ):
        zeros = torch.zeros_like(z0_mean)
        ones = torch.ones_like(z0_mean)
        if sample_vae:
            z0 = z0_mean + torch.normal(zeros, ones) * z0_std.exp()
        else:
            z0 = z0_mean
        prev_samp = samp_ts[:, self.rec.encoded_len - 1]

        next_step_size = self.solver.default_step_size

        su_probs = []
        steps = []

        results = []

        for samp_id in range(self.rec.encoded_len, samp_ts.shape[1]):
            samp = samp_ts[:, samp_id]

            pred, su_probs_, _, steps_ = self.solver(
                    latent_ode=self.latent_ode,
                    init_state=z0.squeeze(1),
                    first_t=prev_samp,
                    last_t=samp,
                    constant_step=next_step_size,
                    adjoint=adjoint,
                    method=method,
                    atol=atol,
                    rtol=rtol,
                    min_dt=min_dt,
                    discrete_updates=discrete_updates,
                )

            su_probs.append(su_probs_)
            steps.append(steps_)
            z0 = pred
            if torch.is_tensor(z0):
                results.append(z0.view(z0.shape[0], 1, -1))
            else:
                results.append(z0[0].view(z0[0].shape[0], 1, -1))

            prev_samp = samp

        results = torch.cat(results, axis=1)

        su_probs = torch.cat(su_probs, dim=1)
        steps = torch.cat(steps, dim=1)

        return results, su_probs, steps

src/model/full_model.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import torch

import utils


def mse(y_pred, y_act, obs_flag=None, **kwargs):
    res = y_pred - y_act

    if obs_flag is not None:
        res = res[obs_flag == 1.0]

    res = res**2

    return res.mean()


def mse_far(y_pred, y_act, obs_flag, first):
    if obs_flag is None:
        return mse(y_pred[:, first:], y_act[:, first:], try_omit_first=None)
    else:
        return mse(
            y_pred[:, first:],
            y_act[:, first:],
            try_omit_first=None,
            obs_flag=obs_flag[:, first:],
        )


def rmse(y_pred, y_act):
    abs_res = torch.abs(y_pred - y_act)
    return ((abs_res**2).mean()) ** (1 / 2)


def far_mape_loss(y_pred, y_act):
    far_seq_len_start = y_act.size(1) // 4
    y_act = y_act[:, far_seq_len_start:, :]
    y_pred = y_pred[:, far_seq_len_start:, :]
    diff = torch.abs(y_act - y_pred)
    loss = (diff / (torch.abs(y_act) + 0.0001)).mean()
    return loss


def mape_loss(y_pred, y_act):
    diff = torch.abs(y_act - y_pred)
    loss = (diff / (torch.abs(y_act) + 0.0001)).mean()
    return loss


def smape_loss(y_pred, y_act):
    diff = torch.abs(y_act - y_pred)
    loss = (diff / (torch.abs(y_act) + torch.abs(y_pred) + 0.0001)).mean()
    return loss * 2.0


def avg_far_mape_loss(y_pred, y_act):
    return (far_mape_loss(y_pred, y_act) + mape_loss(y_pred, y_act)) / 2.0


def wmape_loss(y_pred, y_act, obs_flag, **kwargs):
    diff = torch.abs(y_act - y_pred)[obs_flag == 1.0]
    loss = diff.sum() / (torch.abs(y_act[obs_flag == 1.0]) + 0.0001).sum()
    return loss


z0_prior = torch.distributions.Normal(
    torch.Tensor([0.0]).to(utils.DEFAULT_DEVICE),
    torch.Tensor([1.0]).to(utils.DEFAULT_DEVICE),
)


def kl_divergence(mean, std):
    std = std.exp()
    fp_distr = torch.distributions.Normal(mean, std)
    kl = torch.distributions.kl_divergence(fp_distr, z0_prior)
    return kl.mean()


def elbo_loss(y_pred, y_test, z0_stats, try_omit_first, obs_flag=None):
    z0_mean, z0_std = z0_stats
    if obs_flag:
        raise NotImplementedError
    recon_loss = mse(
        y_pred, y_test, try_omit_first=try_omit_first
    )  # for N(0, 1) we can use MSE
    kl = kl_divergence(z0_mean, z0_std)
    # TODO: maybe torch.losumexp like in the original repo
    loss = recon_loss + 0.01 * kl
    return loss


def mae(y_pred, y_test, obs_flag, **kwargs):
    res = y_pred - y_test

    if obs_flag is not None:
        res = res[obs_flag == 1.0]

    res = torch.abs(res)

    return res.mean()


def mae_far(y_pred, y_act, obs_flag, first):
    if obs_flag is None:
        return mae(y_pred[:, first:], y_act[:, first:], try_omit_first=None)
    else:
        return mae(
            y_pred[:, first:],
            y_act[:, first:],
            try_omit_first=None,
            obs_flag=obs_flag[:, first:],
        )

src/model/losses.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import torch
from torch import nn
from model.components.encoders import MultIdEnc


class LSTMModel(nn.Module):
    def __init__(
        self, obs_dim, hidden_dim, depth, separate_lstm_encoder, encoded_len
    ):
        super(LSTMModel, self).__init__()

        self.separate_lstm_encoder = separate_lstm_encoder

        if separate_lstm_encoder:
            self.enc_lstm = nn.LSTM(
                input_size=obs_dim,
                hidden_size=hidden_dim,
                num_layers=depth,
                batch_first=True,
            )

        self.main_lstm = nn.LSTM(
            input_size=obs_dim,
            hidden_size=hidden_dim,
            num_layers=depth,
            batch_first=True,
        )

        self.ll_out = nn.Linear(hidden_dim, obs_dim)
        self.model_type = "lstm"
        self.obs_dim = obs_dim

        self.rec = MultIdEnc(encoded_len)

    def forward(self, tsdata, **kwargs):
        dataset = tsdata.dataset
        first_x = dataset[:, : self.rec.encoded_len]

        if self.separate_lstm_encoder:
            h = self.enc_lstm(first_x)[1]

            next_in = dataset[:, self.rec.encoded_len :, :1] * 0
            next_in = next_in.repeat(1, 1, self.obs_dim)

            next_x = self.main_lstm(next_in, h)[0]
            next_x = next_x.reshape(-1, next_x.shape[2])
            next_x = self.ll_out(next_x)
            next_x = next_x.view(first_x.shape[0], -1, next_x.shape[-1])

            pred_x = torch.cat([first_x, next_x], dim=1)

        else:

            if self.rec.encoded_len > 1:
                _, h = self.main_lstm(first_x[:, : self.rec.encoded_len - 1])
            else:
                h = None

            z = first_x[:, self.rec.encoded_len - 1 : self.rec.encoded_len]
            pred_x = [first_x]

            for _ in range(self.rec.encoded_len, dataset.shape[1]):
                z, h = self.main_lstm(z, h)
                z = self.ll_out(z.reshape(-1, z.shape[2])).unsqueeze(1)
                pred_x.append(z)

            pred_x = torch.cat(pred_x, dim=1)

        return {
            "pred_x": pred_x,
            "z0_mean": None,
            "z0_std": None,
        }

    def predict(self, *args, **kwargs):
        self.eval()
        with torch.no_grad():
            return self.forward(*args, **kwargs)

src/model/lstm.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import torch
from torch import nn

import problem_specs
import utils


class LTSF(nn.Module):
    def __init__(self, rec, dec, latent_dim, out_values, ltsf_type, N):
        super(LTSF, self).__init__()

        self.ltsf_type = ltsf_type

        if ltsf_type == "linear":
            self.model_type = "latent_ltsf_linear"
        elif ltsf_type == "nlinear":
            self.model_type = "latent_ltsf_nlinear"
        else:
            raise NotImplementedError

        self.rec = rec
        self.dec = dec

        self.N = N

        self.latent_dim = latent_dim
        self.out_values = out_values

        self.ll = nn.Linear(N * latent_dim, out_values * latent_dim)

    def forward(self, tsdata, sample_vae, analytical=True, **kwargs):
        assert isinstance(tsdata, problem_specs.TSData)
        z0_mean, z0_std, _, _ = self.rec.recognize(tsdata, needed_last=self.N)
        if sample_vae:
            zeros = torch.zeros_like(z0_mean)
            ones = torch.ones_like(z0_mean)
            enc_zs = z0_mean + torch.normal(zeros, ones) * z0_std.exp()
        else:
            enc_zs = z0_mean

        if self.ltsf_type == "linear":
            z = enc_zs[:, -self.N :].reshape(-1, self.latent_dim * self.N)
            pred_z = self.ll(z).view(-1, self.out_values, self.latent_dim)
        elif self.ltsf_type == "nlinear":
            z = enc_zs[:, -self.N :]
            z = z - enc_zs[:, -1:]
            z = z.reshape(-1, self.latent_dim * self.N)
            pred_z = enc_zs[:, -1:] + self.ll(z).view(
                -1, self.out_values, self.latent_dim
            )
        else:
            raise NotImplementedError
        pred_x = self.dec(pred_z)

        pred_x = torch.cat(
            [tsdata.dataset[:, : self.rec.encoded_len], pred_x], dim=1
        )
        obs_flag = tsdata.obs_flag
        pred_x[obs_flag == 0.0] = utils.DEFAULT_MISSING_VALUE

        return {
            "pred_x": pred_x,
            "z0_mean": z0_mean,
            "z0_std": z0_std,
        }

    def predict(self, *args, **kwargs):
        self.eval()
        with torch.no_grad():
            return self.forward(*args, **kwargs)

src/model/ltsf.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
# flake8: noqa
from .odeint_w_events import ODEINTWEvents

src/model/solvers/__init__.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import torch
from torch import nn

import utils


# main solver interface
class ODEINTWEvents(nn.Module):
    def __init__(self, default_step_size):
        super(ODEINTWEvents, self).__init__()
        self.device = utils.DEFAULT_DEVICE
        self.default_step_size = default_step_size

    def forward(
        self,
        latent_ode,
        init_state,
        first_t,
        last_t,
        constant_step,
        discrete_updates,
        method,
        atol,
        rtol,
        min_dt,
        adjoint=False,
    ):
        return self._my_odeint(
            latent_ode=latent_ode,
            init_state=init_state,
            first_t=first_t,
            last_t=last_t,
            step_size=constant_step,
            discrete_updates=discrete_updates,
            adjoint=adjoint,
            method=method,
            atol=atol,
            rtol=rtol,
            min_dt=min_dt,
        )

    def _my_odeint(
        self,
        latent_ode,
        init_state,
        first_t,
        last_t,
        step_size,
        min_dt,
        discrete_updates,
        adjoint,
        method,
        atol,
        rtol,
    ):
        state = init_state
        su_probs = []

        min_dt = torch.zeros(state.shape[0]).to(utils.DEFAULT_DEVICE) + min_dt

        if not torch.is_tensor(step_size):
            step_size = (
                torch.zeros((state.shape[0],)).to(utils.DEFAULT_DEVICE)
                + step_size
            )

        next_step_size = step_size

        steps = []

        if adjoint:
            raise NotImplementedError("adjoint with variable ts")
        while torch.any(first_t + 0.00001 < last_t):
            should_change_state = (first_t + 0.00001 < last_t).float()
            should_change_state = should_change_state.detach().view(-1, 1)
            steps.append(should_change_state)
            if method == "rk4":
                after_step_state, step_size = rk4_step_func(
                    latent_ode, state, first_t, step_size
                )
                next_step_size = step_size
            elif method == "euler":
                after_step_state, step_size = euler_step_func(
                    latent_ode, state, first_t, step_size
                )
                next_step_size = step_size
            elif method == "dopri5":
                max_dt = last_t - first_t
                after_step_state, step_size, next_step_size = dopri5_step_func(
                    func=latent_ode,
                    y0=state,
                    t0=first_t,
                    dt=next_step_size,
                    max_dt=max_dt,
                    min_dt=min_dt,
                    atol=atol,
                    rtol=rtol,
                )
            else:
                raise NotImplementedError

            possible_next_state = after_step_state
            su_probs.append(state.detach()[:, :1] * 0)

            state = (
                should_change_state * possible_next_state
                + (1 - should_change_state) * state
            )
            assert len(state.shape) == 2, "state: {}".format(state.shape)

            first_t = first_t + step_size

        su_probs = torch.cat(su_probs, dim=1)
        steps = torch.cat(steps, dim=1)

        return state, su_probs, next_step_size, steps

"""
based on https://github.com/rtqichen/torchdiffeq/blob/master/torchdiffeq/_impl/rk_common.py
"""

def rk4_step_func(func, y0, t0, dt):
    dt = dt.detach()
    y0 = y0.detach()
    k1 = func(y0.detach())
    k2 = func((y0 + dt.view(-1, 1) * k1 / 3))
    k3 = func((y0 + dt.view(-1, 1) * (k2 - k1 / 3)))
    k4 = func((y0 + dt.view(-1, 1) * (k1 - k2 + k3)))
    return y0 + (k1 + 3 * (k2 + k3) + k4) * dt.view(-1, 1) / 8, dt


def euler_step_func(func, y0, history, t0, dt):
    dt = dt.detach()
    return y0 + dt.view(-1, 1) * func(y0), dt


c21 = 1 / 5
c31 = 3 / 40
c32 = 9 / 40
c41 = 44 / 45
c42 = -56 / 15
c43 = 32 / 9
c51 = 19372 / 6561
c52 = -25360 / 2187
c53 = 64448 / 6561
c54 = -212 / 729
c61 = 9017 / 3168
c62 = -355 / 33
c63 = -46732 / 5247
c64 = 49 / 176
c65 = -5103 / 18656
c71 = 35 / 384
c73 = 500 / 1113
c74 = 125 / 192
c75 = -2187 / 6784
c76 = 11 / 84

y1_1 = 35 / 384
y1_3 = 500 / 1113
y1_4 = 125 / 192
y1_5 = -2187 / 6784
y1_6 = 11 / 84

z1_1 = 5179 / 57600
z1_3 = 7571 / 16695
z1_4 = 393 / 640
z1_5 = -92097 / 339200
z1_6 = 187 / 2100
z1_7 = 1 / 40


def dopri5_step_func(func, y0, t0, dt, max_dt, min_dt, atol, rtol, attempts=1):
    i = 0
    d_stop = 0.0
    next_dt = dt
    while i < attempts:
        dt = next_dt
        dt = torch.maximum(dt, min_dt)
        dt = torch.minimum(dt, max_dt)
        dt = dt.detach()
        k1 = dt.view(-1, 1) * func(y0)
        k2 = dt.view(-1, 1) * func((y0 + c21 * k1))
        k3 = dt.view(-1, 1) * func((y0 + c31 * k1 + c32 * k2))
        k4 = dt.view(-1, 1) * func((y0 + c41 * k1 + c42 * k2 + c43 * k3))
        k5 = dt.view(-1, 1) * func(
            (y0 + c51 * k1 + c52 * k2 + c53 * k3 + c54 * k4)
        )
        k6 = dt.view(-1, 1) * func(
            (y0 + c61 * k1 + c63 * k3 + c64 * k4 + c65 * k5)
        )
        k7 = dt.view(-1, 1) * func(
            (y0 + c71 * k1 + c73 * k3 + c74 * k4 + c75 * k5 + c76 * k6)
        )

        y1 = y0 + (y1_1 * k1 + y1_3 * k3 + y1_4 * k4 + y1_5 * k5 + y1_6 * k6)
        z1 = y0 + (
            z1_1 * k1
            + z1_3 * k3
            + z1_4 * k4
            + z1_5 * k5
            + z1_6 * k6
            + z1_7 * k7
        )
        z1 = z1.detach()

        d = torch.max(torch.abs(y1 - z1), dim=-1).values

        v_rtol = (
            rtol
            * torch.max(
                torch.max(y0.detach().abs(), y1.detach().abs()), dim=-1
            ).values
        )
        error_ratio = atol + v_rtol

        i += 1

        s = (error_ratio / (d + 0.0000001)) ** (1 / 5)

        next_dt = d_stop * dt + (1 - d_stop) * dt * s
        next_dt = torch.maximum(next_dt, min_dt)
        next_dt = torch.minimum(next_dt, max_dt)
        
        d_stop = d < error_ratio.float()

        if d_stop.min() > 0.9:
            break

    return y1, dt.detach(), next_dt.detach()

src/model/solvers/odeint_w_events.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
from torch import nn

import utils


class RNNSolver(nn.Module):
    def __init__(
        self,
    ):
        super(RNNSolver, self).__init__()
        self.device = utils.DEFAULT_DEVICE

    def forward(
        self,
        latent_ode,
        init_state,
        *args,
        **kwargs,
    ):
        new_x = latent_ode(init_state)

        su_probs = new_x[:, :1].detach() * 0
        steps = su_probs + 1

        return new_x, su_probs, None, steps

src/model/solvers/rnn_solver.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
"""
modified, based on https://github.com/HazyResearch/spacetime
"""
import torch
import torch.nn as nn
from einops import repeat

from .ssm import (
    init_embedding, Encoder, Decoder, init_mlp
)
from model.components.encoders import MultIdEnc


class SpaceTime(nn.Module):
    def __init__(
        self,
        obs_dim,
        embedding_dim,
        encoder_dim,
        encoded_len,
        pred_len,
        hidden_dim,
        encoder_layers,
        decoder_layers,
        inference_only: bool=False,
    ):
        super().__init__()

        self.inference_only = inference_only
        self.lag     = encoded_len
        self.horizon = encoded_len + pred_len
        self.rec = MultIdEnc(encoded_len)
        self.model_type = "spacetime"
        
        self.embedding  = init_embedding({
            "method": "linear",
            "kwargs": {
                "input_dim": obs_dim,
                "embedding_dim": embedding_dim,
            }
        })
        n_kernels=8
        self.encoder    = self.init_encoder(
            {
                'blocks': [
                    {
                        'input_dim': embedding_dim, 
                        'pre_config': {
                            'method': 'residual', 
                            'kwargs': {
                                'max_diff_order': 4, 
                                'min_avg_window': 4, 
                                'max_avg_window': encoder_dim // 2, 
                                'model_dim': encoder_dim, 
                                'n_kernels': n_kernels, 
                                'kernel_dim': 2, 
                                'kernel_repeat': encoder_dim // n_kernels, 
                                'n_heads': 1, 
                                'head_dim': 1, 
                                'kernel_weights': None, 
                                'kernel_init': None, 
                                'kernel_train': False, 
                                'skip_connection': False, 
                                'seed': 0
                            }
                        }, 
                        'ssm_config': {
                            'method': 'companion', 
                            'kwargs': {
                                'model_dim': encoder_dim, 
                                'n_kernels': encoder_dim//n_kernels, 
                                'kernel_dim': encoder_dim // 2, 
                                'kernel_repeat': 1, 
                                'n_heads': 8, 
                                'head_dim': 1, 
                                'kernel_weights': None, 
                                'kernel_init': 'normal', 'kernel_train': True, 'skip_connection': True, 'norm_order': 1}}, 'mlp_config': {
                                    'method': 'mlp', 'kwargs': {
                                        'input_dim': hidden_dim, 
                                        'output_dim': hidden_dim, 
                                        'activation': 'gelu', 
                                        'dropout': 0.25, 
                                        'layernorm': False, 
                                        'n_layers': 1, 
                                        'n_activations': 1, 
                                        'pre_activation': True, 
                                        'input_shape': 'bld', 
                                        'skip_connection': True, 
                                        'average_pool': None
                                    }
                                }, 'skip_connection': True, 'skip_preprocess': False
                    } for _ in range(encoder_layers)
                ]
            }
        )
        self.decoder    = self.init_decoder(
            {
                'blocks': [
                    {
                        'input_dim': encoder_dim, 
                        'pre_config': {
                            'method': 'identity', 
                            'kwargs': None
                        }, 
                        'ssm_config': {
                            'method': 'closed_loop_companion', 
                            'kwargs': {
                                'lag': encoded_len, 
                                'horizon': pred_len + encoded_len, 
                                'model_dim': hidden_dim, 
                                'n_kernels': hidden_dim, 
                                'kernel_dim': hidden_dim // 2, 
                                'kernel_repeat': 1, 
                                'n_heads': 1, 
                                'head_dim': 1, 
                                'kernel_weights': None, 
                                'kernel_init': 'normal', 'kernel_train': True, 'skip_connection': False, 'norm_order': 1, 'use_initial': False}}, 'mlp_config': {'method': 'identity', 'kwargs': None}, 'skip_connection': False, 'skip_preprocess': False}
                for _ in range(decoder_layers) ]
            }
        )
        self.output     = self.init_output(
            {
                'input_dim': hidden_dim, 
                'output_dim': obs_dim, 
                'method': 'mlp', 
                'kwargs': {
                    'input_dim': hidden_dim, 
                    'output_dim': obs_dim,
                    'activation': 'gelu', 
                    'dropout': 0.25, 
                    'layernorm': False, 
                    'n_layers': 1,
                    'n_activations': 1, 
                    'pre_activation': True, 
                    'input_shape': 'bld', 
                    'skip_connection': 
                    False, 
                    'average_pool': None,
                }
            }
        )
        
    def init_encoder(self, config):
        self.encoder = Encoder(config)
        # Allow access to first encoder SSM kernel_dim
        self.kernel_dim = self.encoder.blocks[0].ssm.kernel_dim
        return self.encoder
    
    def init_decoder(self, config):
        self.decoder = Decoder(config)
        self.decoder.blocks.ssm.lag = self.lag
        self.decoder.blocks.ssm.horizon = self.horizon
        return self.decoder
    
    def init_output(self, config):
        return init_mlp(config)
    
    # -------------
    # Toggle things
    # -------------
    def set_inference_only(self, mode=False):
        self.inference_only = mode
        self.decoder.blocks.ssm.inference_only = mode
        
    def set_closed_loop(self, mode=True):
        self.decoder.blocks.ssm.closed_loop = mode
        
    def set_train(self):
        self.train()
        
    def set_eval(self):
        self.eval()
        self.set_inference_only(mode=True)
        
    def set_lag(self, lag: int):
        self.decoder.blocks.ssm.lag = lag
        
    def set_horizon(self, horizon: int):
        self.decoder.blocks.ssm.horizon = horizon
        
    # ------------
    # Forward pass
    # ------------
    def forward(self, tsdata, **kwargs):
        u = tsdata.dataset[:, :self.rec.encoded_len]
        self.set_closed_loop(True)
        # Assume u.shape is (batch x len x dim), 
        # where len = lag + horizon
        z = self.embedding(u)
        z = self.encoder(z)
        y_c, _ = self.decoder(z)  
        y_c = self.output(y_c)  # y_c is closed-loop output

        if not self.inference_only:  
            # Also compute outputs via open-loop
            self.set_closed_loop(False)
            y_o, z_u = self.decoder(z)
            y_o = self.output(y_o)    # y_o is "open-loop" output
            # Prediction and "ground-truth" for next-time-step 
            # layer input (i.e., last-layer output)
            z_u_pred, z_u_true = z_u  
        else:
            y_o = None
            z_u_pred, z_u_true = None, None
        # Return (model outputs), (model last-layer next-step inputs)
        # return (y_c, y_o), (z_u_pred, z_u_true)
        # raise Exception(u.shape, y_c.shape, y_o.shape, z_u_pred.shape, z_u_true.shape)
        y_c = y_c[:, self.rec.encoded_len:]
        y_c = torch.cat([u, y_c], dim=1)
        return {
            "pred_x": y_c,
            "z0_mean": None,
            "z0_std": None,
        }

    def predict(self, *args, **kwargs):
        with torch.no_grad():
            return self(*args, **kwargs)

src/model/spacetime/__init__.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
"""
modified, based on https://github.com/HazyResearch/spacetime
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import repeat, rearrange
import opt_einsum as oe
import math

class OurModule(nn.Module):
    def __init__(self): 
        super().__init__()

    def register(self, name, tensor, trainable=False, lr=None, wd=None):
        """Utility method: register a tensor as a buffer or trainable parameter"""
        if trainable:
            try:
                self.register_parameter(name, nn.Parameter(tensor))
            except KeyError:
                delattr(self, name)
                self.register_parameter(name, nn.Parameter(tensor))
        else:
            
            try:
                self.register_buffer(name, tensor)
            except KeyError:
                delattr(self, name)
                self.register_buffer(name, tensor)

        optim = {}
        if trainable and lr is not None: optim["lr"] = lr
        if trainable and wd is not None: optim["weight_decay"] = wd
        if len(optim) > 0: setattr(getattr(self, name), "_optim", optim)

class SSM(OurModule):
    def __init__(self, 
                 model_dim: int, 
                 n_kernels: int,     # Number of kernels / scales
                 kernel_dim: int,
                 kernel_repeat: int,
                 n_heads: int=None,  # Number of heads per kernel
                 head_dim: int=1,    # Dimension of each head
                 kernel_weights: torch.float=None,
                 kernel_init: str='normal',
                 kernel_train: bool=True,
                 skip_connection: bool=False,
                 seed: int=42):
        super().__init__()
        # At least one of these should be int
        assert not (n_heads is None and head_dim is None)
                 
        self.model_dim = model_dim
        self.n_kernels = n_kernels
        self.kernel_dim = kernel_dim
        self.kernel_repeat = kernel_repeat
        self.head_dim, self.n_heads = self.init_heads(n_heads, head_dim)
        self.kernel_weights  = kernel_weights
        self.kernel_init     = kernel_init
        self.kernel_train    = kernel_train
        self.skip_connection = skip_connection
        self.seed = seed
        self.generator = torch.Generator()
        self.generator.manual_seed(seed)
        
        self.init_weights()
        
    def init_heads(self, n_heads: int, head_dim: int):
        if head_dim is None:
            self.head_dim = self.model_dim // (self.kernel_repeat * 
                                               self.n_kernels * n_heads)
            self.n_heads  = n_heads
        else:
            self.head_dim = head_dim
            self.n_heads  = self.model_dim // (self.kernel_repeat * 
                                               self.n_kernels * head_dim)
        if self.n_heads == 5:
            print(self.n_heads, self.head_dim, n_heads, head_dim, self.kernel_repeat, self.n_kernels, self.model_dim)
            raise Exception
        return self.head_dim, self.n_heads
        
    def fft_conv(self, u_input: torch.tensor, v_kernel: torch.tensor):
        # Convolve u with v in O(n log n) time with FFT (n = len(u))
        L   = u_input.shape[-1]  # Assume u is input
        u_f = torch.fft.rfft(u_input, n=2*L) # (B H L)
        v_f = torch.fft.rfft(v_kernel[:, :L], n=2*L) # (H L)

        y_f = oe.contract('b h l, h l -> b h l', u_f, v_f) 
        y   = torch.fft.irfft(y_f, n=2*L)[..., :L]  # (B H L)
        return y
    
    def init_weights(self):
        if self.kernel_weights is not None:  
            # lr and wd as None sets them to be same as model lr and weight_decay
            register('k', self.kernel_weights, trainable=True, lr=None, wd=None)
        
        skip = torch.randn(self.model_dim)
        self.register('skip', skip, trainable=True, lr=None, wd=None)
    
    def get_kernel(self):
        raise NotImplementedError
        
    def forward(self, u):
        u = rearrange(u, 'b l d -> b d l')  # Assume u is B x L x D
        # Repeat kernels across heads
        if self.kernel_weights is None:
            k = self.get_kernel(u)
            k = repeat(k, 'nk kd -> (kr nk nh hd) kd', 
                   kr=self.kernel_repeat, nh=self.n_heads, hd=self.head_dim)
        else:
            k = self.k
        y = self.fft_conv(u, k)
        if self.skip_connection:
            y = y + oe.contract('b d l, d -> b d l', u, self.skip)
        y = rearrange(y, 'b d l -> b l d')
        return y

class CompanionSSM(SSM):
    """
    Open-loop implementation of Companion SSM:
    -> y_t = C \sum_{i = 0}^{k - 1 - i} A^k B u_i
       where A is companion matrix
    """
    def __init__(self, norm_order, **kwargs):
        self.norm_order = norm_order
        kwargs['kernel_repeat'] = 1
        kwargs['kernel_weights'] = None
        kwargs['kernel_train'] = True
        # Set kwargs['n_heads'] as n_kernels for preprocessing kernels
        # Set kwargs['head_dim'] to be original sample input dim
        super().__init__(**kwargs)
        
    def init_kernel_weights(self, kernel_init):
        if kernel_init == 'normal':
            kernel = torch.randn(self.n_kernels, self.kernel_dim)
        elif kernel_init == 'xavier':
            # Xavier-ish initialization
            stdv = 1. / math.sqrt(self.kernel_dim)
            kernel = torch.FloatTensor(self.n_kernels, 
                                       self.kernel_dim).uniform_(-stdv, stdv)
        else:
            raise NotImplementedError
        return kernel
        
    def init_weights(self):
        super().init_weights()  # Initializes skip connection
        self._fp = (self.n_kernels, self.kernel_dim)
        
        # Shift matrix initialization
        self.shift_matrix = torch.zeros(self.n_kernels, 
                                        self.kernel_dim, 
                                        self.kernel_dim)
        self.shift_matrix[:, 1:, :-1] = torch.eye(self.kernel_dim - 1)
        self.p_padding = torch.zeros(*self._fp)
        self.p_padding[:, -1] = 1.
        
        # A matrix
        a = self.init_kernel_weights(self.kernel_init)
        self.register("a", a, trainable=True, lr=None, wd=None)
        
        # B matrix
        b = self.init_kernel_weights(self.kernel_init) 
        self.register("b", b, trainable=True, lr=None, wd=None)
        
        # C matrix
        c = self.init_kernel_weights(self.kernel_init)
        self.register("c", c, trainable=True, lr=None, wd=None)
    
    def norm(self, x, ord=1):
        # x.shape is either (H x D) or (H x D x D)
        x_norm = torch.linalg.norm(x, ord=ord, dim=-1, keepdim=True)
        # If norm(x) in batch close to 0, don't normalize 
        # (heuristicky, but we norm for stability)
        try:
            x = x / x_norm if torch.abs(x_norm).mean().item() > 1e-4 else x  
        except Exception as e:
            print(e)
            breakpoint()
        # x = F.normalize(x, dim=1, p=ord, eps=1)
        return x
    
    def matrix_power(self, l, c, b, p):
        # Construct companion matrix
        A = self.shift_matrix.to(p.device) + (
            oe.contract('h i, h j -> h j i', 
                        self.p_padding.to(p.device), p)
        )
        # Use repeated squares to power A
        g = krylov(l, A, b, c)
        return g
    
    def get_kernel(self, u, c=None, l=None):
        l = u.shape[-1] if l is None else l
        c = self.c if c is None else c
        a = (self.norm(self.a, ord=self.norm_order) 
             if self.norm_order > 0 else self.a)
        f = self.matrix_power(l, c, self.b, a).to(u.device)
        return f
    
    def forward(self, u):
        return super().forward(u)
        

class ClosedLoopCompanionSSM(CompanionSSM):
    """
    Closed-loop implementation of Companion SSM:
    - Instantiate A, B, C; so we can compute both:
    - Open-loop inference:   
      -> y_{n + h} = \sum_{i = 0}^{n + h - 1} CA^{n + h - 1 - i} B u_i
    - Closed-loop inference: 
      -> y_{n + h} = C(A + BK)^{h} x_n
                   = C(A + BK)^{h} \sum_{j = 0}^{n - 1} A^{n - 1 - j} B u_j
                   = C(A + BK)^{n + h - 1} x_1
                   = C(A + BK)^{n + h - 1} B u_0
                   = \sum_{i = 0}^{n + h - 1} C(A + BK)^{n + h - 1 - i} B u_i, u_j = 0 for j > 0
    """
    def __init__(self, 
                 lag: int=1,
                 horizon: int=1,
                 use_initial: bool=False,
                 **kwargs):
        self.lag     = lag
        self.horizon = horizon
        self.use_initial = use_initial  # When False, assumes initial hidden_state x_0 = 0. True not implemented
        self.closed_loop = True         # Toggle closed or open-loop forward pass, see self.forward
        self.inference_only = False     # Toggle different behavior during training and test
        kwargs['kernel_repeat'] = 1
        kwargs['kernel_weights'] = None
        kwargs['kernel_train'] = True
        kwargs['skip_connection'] = False
        super().__init__(**kwargs)
        
    def init_kernel_weights(self, kernel_init):
        if kernel_init == 'normal':
            kernel = torch.randn(self.n_kernels, self.kernel_dim)
        elif kernel_init == 'xavier':
            # Xavier-ish initialization
            stdv = 1. / math.sqrt(self.kernel_dim)
            kernel = torch.FloatTensor(
                self.n_kernels, self.kernel_dim).uniform_(-stdv, stdv)
        else:
            raise NotImplementedError
        return kernel
        
    def init_weights(self):
        super().init_weights()  # Initializes skip connection, A, B, C
        # K matrix
        k = self.init_kernel_weights(self.kernel_init)
        self.register("k", k, trainable=True, lr=None, wd=None)
    
    def get_companion_matrix(self, p):
        # Construct companion matrix
        return self.shift_matrix.to(p.device) + (
            oe.contract('h i, h j -> h j i', 
                        self.p_padding.to(p.device), p)
        )
    
    def fft_conv_d(self, u, v):
        L   = u.shape[-1]
        u_f = torch.fft.rfft(u, n=2*L, dim=2) # (B H L)
        v_f = torch.fft.rfft(v, n=2*L, dim=2) # (H D L)

        y_f = oe.contract('b h l, h d l -> b h l d', u_f, v_f) 
        y   = torch.fft.irfft(y_f, n=2*L, dim=2)[:, :, :L, :] # (B H L D)
        return y
    
    def forward(self, u):
        """
        During training, call this function twice to compute closed-loop and open-loop
        -> minimize the closed-loop?
        """
        u = rearrange(u, 'b l d -> b d l')
        b, d, l = u.shape
        l_horizon = self.horizon
        
        # Normalize just the non-shift column, 
        # alternatively could normalize A + BK below 
        a = (self.norm(self.a, ord=self.norm_order) 
             if self.norm_order > 0 else self.a)
        A = self.get_companion_matrix(a)
        
        if self.closed_loop:  # Compute closed-loop forecast
            # Compute hidden state 
            # -> x_lag = \sum_{i = 0}^{lag - 1} A^{lag - 1 - i}B u_i
            k_x = krylov(l, A, self.b, c=None).to(u.device)
            x = self.fft_conv_d(u, k_x)  # shape: B x H x L x D
            
            # Compute A + BK matrix
            b = (self.norm(self.b, ord=self.norm_order) 
                 if self.norm_order > 0 else self.b)
            k = (self.norm(self.k, ord=self.norm_order) 
                 if self.norm_order > 0 else self.b)
            A_BK = A + oe.contract('h i, h j -> h i j', b, k)
            
            # Rollout: Compute C(A + BK)^{h} * x_lag and K(A + BK)^{h} * x_lag
            # First compute hidden state
            x = krylov(l_horizon, A_BK, x[:, :, -1, :], c=None)
            
            # Compute predictions for layer output
            c = self.norm(self.c, ord=self.norm_order) if self.norm_order > 0 else self.c
            y = torch.einsum('...nl, ...n -> ...l', x, c).contiguous()
            y = rearrange(y, 'b d l -> b l d')
            
            # Compute predictions for layer next-time-step input (prior layer next-time-step output)
            if not self.inference_only and self.closed_loop:
                u = torch.einsum('...nl, ...n -> ...l', x, self.k).contiguous()
                u = rearrange(u, 'b d l -> b l d')
            else:
                u = None
            # Layer outputs, and next-time-step layer inputs
            return y, u
        
        else:  # Compute open-loop forecast up to L
            # A = self.norm(A, ord=self.norm_order)
            # Return CA^{n}B where A = a is computed companion matrix from self.a
            b = (self.norm(self.b, ord=self.norm_order) 
                 if self.norm_order > 0 else self.b)
            c = self.norm(self.c, ord=self.norm_order) if self.norm_order > 0 else self.c
            k = krylov(l, A, b, c).to(u.device)
            k = repeat(k, 'nk kd -> (kr nk nh hd) kd', 
                       kr=self.kernel_repeat, nh=self.n_heads, hd=self.head_dim)
            y = rearrange(self.fft_conv(u, k), 'b d l -> b l d')
            
            if not self.inference_only:
                _k  = self.norm(self.k, ord=self.norm_order)
                k_u = krylov(l, A, b, _k).to(u.device)
                k_u = repeat(k_u, 'nk kd -> (kr nk nh hd) kd', 
                             kr=self.kernel_repeat, nh=self.n_heads, hd=self.head_dim)
                y_u = rearrange(self.fft_conv(u, k_u), 'b d l -> b l d')
            else:
                y_u = None
            return y, y_u

class ClosedLoopShiftSSM(ClosedLoopCompanionSSM):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
    def init_weights(self):
        super().init_weights()  # Initializes skip connection, A, B, C
        # A Matrix
        a = torch.zeros(self.n_kernels, self.kernel_dim)
        self.register("a", a, trainable=False, lr=None, wd=None)
        
        # B Matrix - make it not learnable
        b = torch.zeros(self.n_kernels, self.kernel_dim)
        b[:, 0] = 1
        self.register("b", b, trainable=False, lr=None, wd=None)
        
        # C matrix
        c = self.init_kernel_weights(self.kernel_init)
        self.register("c", c, trainable=True, lr=None, wd=None)
        
        # K matrix
        k = self.init_kernel_weights(self.kernel_init)
        self.register("k", k, trainable=True, lr=None, wd=None)
        raise Exception(k.shape)
        
    def get_companion_matrix(self, p):
        # Construct "companion" matrix
        return self.shift_matrix.to(p.device)


class ShiftSSM(CompanionSSM):
    """
    Open-loop implementation of Shift SSM:
    -> y_t = C \sum_{i = 0}^{k - 1 - i} S^k B u_i
       where S is shift matrix
    """
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
    def init_weights(self):
        super().init_weights()  # Initializes skip connection, B, C matrices
        # A column initialized in super().init_weights(), but now we zero-out
        a = torch.zeros(self.n_kernels, self.kernel_dim)
        self.register("a", a, trainable=False, lr=None, wd=None)
        
        # B Matrix - make it not learnable by default
        b = torch.zeros(self.n_kernels, self.kernel_dim)
        b[:, 0] = 1.
        self.register("b", b, trainable=False, lr=None, wd=None)
        
        # C matrix
        c = self.init_kernel_weights(self.kernel_init)
        self.register("c", c, trainable=True, lr=None, wd=None)
    
    def forward(self, u):
        return super().forward(u)


def krylov(L, A, b, c=None, return_power=False):
    """
    Compute the Krylov matrix (b, Ab, A^2b, ...) using the squaring trick.

    If return_power=True, return A^{L-1} as well
    """
    # TODO There is an edge case if L=1 where output doesn't get broadcasted, which might be an issue if caller is expecting broadcasting semantics... can deal with it if it arises

    x = b.unsqueeze(-1) # (..., N, 1)
    A_ = A

    AL = None
    if return_power:
        AL = torch.eye(A.shape[-1], dtype=A.dtype, device=A.device)
        _L = L-1

    done = L == 1
    # loop invariant: _L represents how many indices left to compute
    while not done:
        if return_power:
            if _L % 2 == 1: AL = A_ @ AL
            _L //= 2

        # Save memory on last iteration
        l = x.shape[-1]
        if L - l <= l:
            done = True
            _x = x[..., :L-l]
        else: _x = x

        try:
            _x = A_ @ _x
        except Exception as e:
            print(e)
            breakpoint()
        x = torch.cat([x, _x], dim=-1) # there might be a more efficient way of ordering axes
        if not done: A_ = A_ @ A_

    try:
        assert x.shape[-1] == L
    except:
        print('x.shape', x.shape)
        print('L', L)
        breakpoint()

    if c is not None:
        x = torch.einsum('...nl, ...n -> ...l', x, c)
    x = x.contiguous() # WOW!!
    if return_power:
        return x, AL
    else:
        return x


def krylov_sequential(L, A, b, c=None):
    """ Constant matrix A

    A : (..., N, N)
    b : (..., N)
    c : (..., N)

    Returns
    if c:
    x : (..., L)
    x[i, l] = c[i] @ A^l @ b[i]

    else:
    x : (..., N, L)
    x[i, l] = A^l @ b[i]
    """

    # Check which of dim b and c is smaller to save memory
    if c is not None and c.numel() < b.numel():
        return krylov_sequential(L, A.transpose(-1, -2), c, b)

    b_ = b
    x = []
    for _ in range(L):
        if c is not None:
            x_ = torch.sum(c*b_, dim=-1) # (...) # could be faster with matmul or einsum?
        else:
            x_ = b_
        x.append(x_)
        b_ = (A @ b_.unsqueeze(-1)).squeeze(-1)

    x = torch.stack(x, dim=-1)
    return x


class ResidualSSM(SSM):
    """
    Computes both order-N differencing and moving average residuals over input sequence
    """
    def __init__(self, 
                 max_diff_order: int=4, 
                 min_avg_window: int=4, 
                 max_avg_window: int=720,
                 n_kernels: int=8,
                 kernel_repeat: int=16,
                 **kwargs):
        self.max_diff_order = max_diff_order
        self.min_avg_window = min_avg_window
        self.max_avg_window = max_avg_window
        self.n_ma_kernels = (n_kernels - self.max_diff_order) * kernel_repeat
        kwargs['n_heads'] = 1
        kwargs['kernel_weights'] = None
        kwargs['kernel_train'] = False
        kwargs['skip_connection'] = False
        # Set kwargs['kernel_repeat'] to number of model n_kernels
        super().__init__(n_kernels=n_kernels, kernel_repeat=kernel_repeat, **kwargs)
        
    def init_weights(self):
        diff_kernel    = repeat(self.init_differencing_weights(), 'nk kd -> (kr nk) kd',
                                kr=self.kernel_repeat)
        ma_r_kernel = self.init_moving_average_weights()  # Shape: (kr x nk) x hd
        self.register('diff_kernel', diff_kernel, trainable=False, lr=None, wd=None)
        self.register('ma_r_kernel', ma_r_kernel, trainable=False, lr=None, wd=None)
        
    def init_differencing_weights(self):
        kernel = torch.zeros(self.max_diff_order, self.max_diff_order).float()
        diff_coeffs = get_pascal(self.max_diff_order, self.max_diff_order).float()
        kernel[:, :self.max_diff_order] += diff_coeffs
        return kernel
    
    def init_moving_average_weights(self):
        ma_window = torch.randint(low=self.min_avg_window,
                                  high=self.max_avg_window,
                                  size=(1, self.n_ma_kernels))
        # Compute moving average kernel 
        max_window = self.max_avg_window
        kernel = torch.zeros(self.n_ma_kernels, max_window)
        kernel[:, 0] = 1.
        
        moving_avg = (1. / torch.clamp(ma_window, min=self.min_avg_window, max=max_window))
        for ix, window in enumerate(ma_window[0]):
            kernel[ix, :window] -= moving_avg[:1, ix]
        return kernel

    def get_kernel(self, u):
        """
        Initialize weights for differencing kernel
        - Assume u is shape B x D x L
        """
        b, d, l = u.shape
        l = max(l, self.diff_kernel.shape[1])
        # Pad kernels to input length
        diff_kernel = F.pad(self.diff_kernel, (0, l - self.diff_kernel.shape[1]), 'constant', 0)
        ma_r_kernel = F.pad(self.ma_r_kernel, (0, l - self.ma_r_kernel.shape[1]), 'constant', 0)
        
        # Combine kernels
        diff_kernel = rearrange(diff_kernel, '(kr nk) kd -> kr nk kd', 
                                kr=self.kernel_repeat)
        ma_r_kernel = rearrange(ma_r_kernel, '(kr nk) kd -> kr nk kd', 
                                kr=self.kernel_repeat)
        
        kernel = torch.cat([diff_kernel, ma_r_kernel], dim=1)
        kernel = repeat(kernel, 'kr nk kd -> (kr nk hd) kd', hd=self.head_dim)
        return kernel
    
    def forward(self, u):
        # Same as base SSM forward, but kernel repeating already taken care of
        u = rearrange(u, 'b l d -> b d l')
        k = self.get_kernel(u)
        y = self.fft_conv(u, k)
        return rearrange(y, 'b d l -> b l d')


class MovingAvgResidualSSM(SSM):
    """
    Computes moving average residuals over input sequence
    """
    def __init__(self, min_avg_window=4, max_avg_window=720, **kwargs):
        self.min_avg_window = min_avg_window
        self.max_avg_window = max_avg_window
        kwargs['n_heads'] = 1
        kwargs['kernel_weights'] = None
        kwargs['kernel_train'] = False
        kwargs['skip_connection'] = False
        # Set kwargs['kernel_repeat'] to number of model n_kernels
        super().__init__(**kwargs)
        
    def init_weights(self):
        # Moving average window kernels
        kernel = torch.zeros(self.n_kernels, self.kernel_dim).float()
        kernel[:, 0] = 1.
        
        # Low is a heuristic for now
        ma_window = torch.randint(low=self.min_avg_window, 
                                  high=self.max_avg_window,  # self.kernel_dim
                                  size=(1, self.n_kernels)).float()
        
        self.register('ma_window', ma_window, trainable=True, lr=None, wd=None)
        self.register('kernel', kernel, trainable=False, lr=None, wd=None)
        
    def get_kernel(self, u):
        """
        Initialize weights for differencing kernel
        - Assume u is shape B x D x L
        """
        b, d, l = u.shape
        # Set kernel values s.t. convolution computes residuals
        # from moving average, i.e., y[t] - mean(y[t:t - m])
        max_window = min(self.max_avg_window, l)
        kernel = self.kernel - (1. / torch.clamp(torch.round(self.ma_window), 
                                                 min=self.min_avg_window, 
                                                 max=max_window).T)
        return F.pad(self.kernel, (0, l-self.kernel_dim, 0, 0), 'constant', 0)


def get_pascal(n, total_rows=None):
    total_rows = n if total_rows is None else total_rows
    # Compute binomial coeffs for all rows up to n
    line = torch.zeros(total_rows, n).float()
    line[:, 0] = 1.
    for j in range(1, n):      # For all rows,
        for k in range(0, j):  # Compute C(j, k)
            # Coefficients are binomial coeffs, 
            # C(n, k + 1) = C(n, k) * (n - k) / (k + 1)
            negate = 2 * k % 2 - 1  # Negate even elements
            line[j][k+1] += (line[j][k] * (j - k) / (k + 1)) * negate
    return line


class DifferencingSSM(SSM):
    """
    Computes order-N differencing over input sequence
    """
    def __init__(self, max_diff_order=4, **kwargs):
        self.max_diff_order = max_diff_order
        kwargs['n_heads'] = 1
        kwargs['kernel_weights'] = None
        kwargs['kernel_train'] = False
        kwargs['skip_connection'] = False
        # Set kwargs['kernel_repeat'] to number of model n_kernels
        super().__init__(**kwargs)
        
    def init_weights(self):
        kernel = torch.zeros(self.n_kernels, self.kernel_dim).float()
        # Hard-coded up to 4 orders, but just the binomial coeffs / Pascal's triangle (with negatives)
        diff_coeffs = get_pascal(self.max_diff_order)
        # Could be slow, but just done once at initialization
        for ix in range(self.n_kernels):
            try:
                kernel[ix, :self.max_diff_order] += diff_coeffs[ix % len(diff_coeffs)].float()
            except:
                breakpoint()
        self.register('kernel', kernel, trainable=False, lr=None, wd=None)
    
    def get_kernel(self, u):
        """
        Initialize weights for differencing kernel
        - Assume u is shape B x D x L
        """
        b, d, l = u.shape
        return F.pad(self.kernel, (0, l-self.kernel_dim, 0, 0), 'constant', 0)


def construct_toeplitz(v, f=0.0):
    """Explicit construction of Krylov matrix [v  A @ v  A^2 @ v  ...  A^{n-1} @ v]
    where A = Z_f. This uses vectorized indexing and cumprod so it's much
    faster than using the Krylov function.
    Parameters:
        v: the starting vector of size n or (rank, n).
        f: real number
    Returns:
        K: Krylov matrix of size (n, n) or (rank, n, n).
    """
    n  = v.shape[-1]
    a = torch.arange(n, device=v.device)
    b = -a
    indices = a[:, None] + b[None]
    K = v[..., indices]
    K[..., indices < 0] *= f
    return K

def triangular_toeplitz_multiply_(u, v, sum=None):
    n = u.shape[-1]
    u_expand = F.pad(u, (0, n))
    v_expand = F.pad(v, (0, n))
    u_f = torch.fft.rfft(u_expand, n=2*n, dim=-1)
    v_f = torch.fft.rfft(v_expand, n=2*n, dim=-1)
    uv_f = u_f * v_f
    if sum is not None:
        uv_f = uv_f.sum(dim=sum)
    output = torch.fft.irfft(uv_f, n=2*n, dim=-1)[..., :n]
    return output

def triangular_toeplitz_multiply_padded_(u, v):
    """ Same as triangular_toeplitz_multiply but inputs and output assume to be 0-padded already. """
    n = u.shape[-1]
    assert n % 2 == 0
    u_f = torch.fft.rfft(u, n=n, dim=-1)
    v_f = torch.fft.rfft(v, n=n, dim=-1)
    uv_f = u_f * v_f
    output = torch.fft.irfft(uv_f, n=n, dim=-1)
    output[..., n:] = 0
    return output

class TriangularToeplitzMult(torch.autograd.Function):
    @staticmethod
    def forward(ctx, u, v):
        ctx.save_for_backward(u, v)
        return triangular_toeplitz_multiply_(u, v)

    @staticmethod
    def backward(ctx, grad):
        u, v = ctx.saved_tensors
        d_u = triangular_toeplitz_multiply_(grad.flip(-1), v).flip(-1)
        d_v = triangular_toeplitz_multiply_(grad.flip(-1), u).flip(-1)
        return d_u, d_v

class TriangularToeplitzMultFast(torch.autograd.Function):
    @staticmethod
    def forward(ctx, u, v):
        n = u.shape[-1]
        u_expand = F.pad(u, (0, n))
        v_expand = F.pad(v, (0, n))
        u_f = torch.fft.rfft(u_expand, n=2*n, dim=-1)
        v_f = torch.fft.rfft(v_expand, n=2*n, dim=-1)

        ctx.save_for_backward(u_f, v_f)

        uv_f = u_f * v_f
        output = torch.fft.irfft(uv_f, n=2*n, dim=-1)[..., :n]
        return output

    @staticmethod
    def backward(ctx, grad):
        u_f, v_f = ctx.saved_tensors
        n = grad.shape[-1]
        g_expand = F.pad(grad.flip(-1), (0, n))
        g_f = torch.fft.rfft(g_expand, n=2*n, dim=-1)
        gu_f = g_f * u_f
        gv_f = g_f * v_f
        d_u = torch.fft.irfft(gv_f, n=2*n, dim=-1)[..., :n]
        d_v = torch.fft.irfft(gu_f, n=2*n, dim=-1)[..., :n]
        d_u = d_u.flip(-1)
        d_v = d_v.flip(-1)
        return d_u, d_v

class TriangularToeplitzMultPadded(torch.autograd.Function):
    @staticmethod
    def forward(ctx, u, v):
        ctx.save_for_backward(u, v)
        output = triangular_toeplitz_multiply_(u, v)
        return output

    @staticmethod
    def backward(ctx, grad):
        u, v = ctx.saved_tensors
        d_u = triangular_toeplitz_multiply_padded_(grad.flip(-1), v).flip(-1)
        d_v = triangular_toeplitz_multiply_padded_(grad.flip(-1), u).flip(-1)
        return d_u, d_v

class TriangularToeplitzMultPaddedFast(torch.autograd.Function):
    """ Trade off speed (20-25% faster) for more memory (20-25%) """

    @staticmethod
    def forward(ctx, u, v):
        n = u.shape[-1]
        u_f = torch.fft.rfft(u, n=n, dim=-1)
        v_f = torch.fft.rfft(v, n=n, dim=-1)

        ctx.save_for_backward(u_f, v_f)

        uv_f = u_f * v_f
        output = torch.fft.irfft(uv_f, n=n, dim=-1)
        output[..., n//2:].zero_()
        return output

    @staticmethod
    def backward(ctx, grad):
        u_f, v_f = ctx.saved_tensors
        n = grad.shape[-1]
        g_expand = F.pad(grad[..., :n//2].flip(-1), (0, n//2))
        g_f = torch.fft.rfft(g_expand, n=n, dim=-1)
        gu_f = g_f * u_f
        gv_f = g_f * v_f
        d_u = torch.fft.irfft(gv_f, n=n, dim=-1)
        d_v = torch.fft.irfft(gu_f, n=n, dim=-1)
        d_u[..., n//2:].zero_()
        d_v[..., n//2:].zero_()
        d_u[..., :n//2] = d_u[..., :n//2].flip(-1) # TODO
        d_v[..., :n//2] = d_v[..., :n//2].flip(-1) # TODO
        return d_u, d_v

# triangular_toeplitz_multiply = triangular_toeplitz_multiply_
triangular_toeplitz_multiply = TriangularToeplitzMult.apply
triangular_toeplitz_multiply_fast = TriangularToeplitzMultFast.apply
triangular_toeplitz_multiply_padded = TriangularToeplitzMultPadded.apply
triangular_toeplitz_multiply_padded_fast = TriangularToeplitzMultPaddedFast.apply


def causal_convolution(u, v, fast=True, pad=False):
    if not pad and not fast:
        return triangular_toeplitz_multiply(u, v)
    if not pad and fast:
        return triangular_toeplitz_multiply_fast(u, v)
    if pad and not fast:
        return triangular_toeplitz_multiply_padded(u, v)
    if pad and fast:
        return triangular_toeplitz_multiply_padded_fast(u, v)
        

def _fft(x, N): return torch.fft.rfft(F.pad(x, (0, 2*N-x.shape[-1])), n=2*N, dim=-1)
def _ifft(x, N): return torch.fft.irfft(x, n=2*N, dim=-1)[..., :N]

def causal_convolution_inverse(u):
    """ Invert the causal convolution/polynomial/triangular Toeplitz matrix represented by u.

    This is easiest in the polynomial view:
    https://www.csa.iisc.ac.in/~chandan/courses/CNT/notes/lec5.pdf
    The idea is that
    h = g^{-1} (mod x^m) => 2h - gh^2 = g^{-1} (mod x^{2m})

    # TODO this can be numerically unstable if input is "poorly conditioned",
    # for example if u[0] is magnitudes different from the rest of u
    """
    N = u.shape[-1]
    v = u[..., :1].reciprocal()
    while v.shape[-1] < N:
        M = v.shape[-1]
        v_f = _fft(v, 2*M)
        u_f = _fft(u[..., :2*M], 2*M)
        _v = -_ifft(u_f * v_f**2, 2*M)
        _v[..., :M] = _v[..., :M] + 2*v
        v = _v
    # TODO contiguous?
    v = v[..., :N]
    return v

""" Below are experimental functions for improving the stability of LSSL/S3 algorithm. Currently not used anywhere. """

def causal_convolution_inverse_wrong(u, v):
    """ Solve u * x = v. Initial attempt by inverting the multiplication algorithm, which I think doesn't work. """
    n = u.shape[-1]
    u_expand = F.pad(u, (0, n))
    v_expand = F.pad(v, (0, n))
    u_f = torch.fft.rfft(u_expand, n=2*n, dim=-1)
    v_f = torch.fft.rfft(v_expand, n=2*n, dim=-1)
    uv_f = v_f / u_f
    x = torch.fft.irfft(uv_f, n=2*n, dim=-1)[..., :n]
    return x

def construct_toeplitz_log(v):
    n  = v.shape[-1]
    a = torch.arange(n, device=v.device)
    b = -a
    indices = a[:, None] + b[None]
    K = v[..., indices]
    K[..., indices < 0] = -100.0
    return K

def _logsumexp(x, dim=-1):
    """ logsumexp for complex """
    m = torch.max(torch.real(x), dim=dim, keepdim=True)[0]
    x = x - m
    x = torch.log(torch.sum(torch.exp(x), dim=dim))
    x = x + m.squeeze(dim)
    return x

def causal_convolution_inverse_log(u, N=-1):
    """ Invert the causal convolution/polynomial/triangular Toeplitz matrix represented by u.

    This is easiest in the polynomial view:
    https://www.csa.iisc.ac.in/~chandan/courses/CNT/notes/lec5.pdf
    The idea is that
    h = g^{-1} (mod x^m) => 2h - gh^2 = g^{-1} (mod x^{2m})

    # TODO this can be numerically unstable if input is "poorly conditioned",
    # for example if u[0] is magnitudes different from the rest of u
    """
    if N < 0:
        N = u.shape[-1]
    v = - u[..., :1]
    while v.shape[-1] < N:
        M = v.shape[-1]
        _v = F.pad(v, (0, M), value=-100.0)
        _v_ = construct_toeplitz_log(_v)
        u_ = u[..., :2*M] if u.shape[-1] >= 2*M else F.pad(u, (0, 2*M-u.shape[-1]), value=-100.0)
        _u = _logsumexp(_v_ + u_, dim=-1)
        _u = _logsumexp(_v_ + _u, dim=-1)
        _u = _u + torch.log(-torch.ones_like(_u))
        _v = _v + torch.log(2.0 * torch.ones_like(_u))
        v = _logsumexp(torch.stack([_v, _u], dim=-1), dim=-1)
    # TODO contiguous?
    v = v[..., :N]

    check = _logsumexp(construct_toeplitz_log(v) + F.pad(u, (0, N-u.shape[-1]), value=-100.0))
    print("check", check, torch.exp(check))
    return v



def power(L, A, v=None):
    """ Compute A^L and the scan sum_i A^i v_i

    A: (..., N, N)
    v: (..., N, L)
    """

    I = torch.eye(A.shape[-1]).to(A) # , dtype=A.dtype, device=A.device)

    powers = [A]
    l = 1
    while True:
        if L % 2 == 1: I = powers[-1] @ I
        L //= 2
        if L == 0: break
        l *= 2
        powers.append(powers[-1] @ powers[-1])

    if v is None: return I

    # Invariants:
    # powers[-1] := A^l
    # l := largest po2 at most L

    # Note that an alternative divide and conquer to compute the reduction is possible and can be embedded into the above loop without caching intermediate powers of A
    # We do this reverse divide-and-conquer for efficiency reasons:
    # 1) it involves fewer padding steps for non-po2 L
    # 2) it involves more contiguous arrays

    # Take care of edge case for non-po2 arrays
    # Note that this initial step is a no-op for the case of power of 2 (l == L)
    k = v.size(-1) - l
    v_ = powers.pop() @ v[..., l:]
    v = v[..., :l]
    v[..., :k] = v[..., :k] + v_

    # Handle reduction for power of 2
    while v.size(-1) > 1:
        v = rearrange(v, '... (z l) -> ... z l', z=2)
        v = v[..., 0, :] + powers.pop() @ v[..., 1, :]
    return I, v.squeeze(-1)

def krylov_toeplitz(L, A, b, c=None):
    """ Specializes to lower triangular Toeplitz matrix A represented by its diagonals

    A : (..., N)
    b : (..., N)
    c : (..., N)

    Returns
    x : (..., N, L)
    x[i, l] = A^l @ b[i]
    """
    x = b.unsqueeze(0) # (1, ..., N)
    A_ = A
    while x.shape[0] < L:
        xx = causal_convolution(A_, x)
        x = torch.cat([x, xx], dim=0) # there might be a more efficient way of ordering axes
        A_ = causal_convolution(A_, A_)
    x = x[:L, ...] # (L, ..., N)
    if c is not None:
        x = torch.einsum('l...n, ...n -> ...l', x, c)
    else:
        x = rearrange(x, 'l ... n -> ... n l')
    x = x.contiguous()
    return x

def krylov_toeplitz_(L, A, b, c=None):
    """ Padded version of krylov_toeplitz that saves some fft's

    TODO currently not faster than original version, not sure why
    """
    N = A.shape[-1]

    x = b.unsqueeze(0) # (1, ..., N)
    x = F.pad(x, (0, N))
    A = F.pad(A, (0, N))
    done = L == 1
    while not done:
        l = x.shape[0]
        # Save memory on last iteration
        if L - l <= l:
            done = True
            _x = x[:L-l]
        else: _x = x
        Af = torch.fft.rfft(A, n=2*N, dim=-1)
        xf = torch.fft.rfft(_x, n=2*N, dim=-1)
        xf_ = Af * xf
        x_ = torch.fft.irfft(xf_, n=2*N, dim=-1)
        x_[..., N:] = 0
        x = torch.cat([x, x_], dim=0) # there might be a more efficient way of ordering axes
        if not done:
            A = torch.fft.irfft(Af*Af, n=2*N, dim=-1)
            A[..., N:] = 0
    x = x[:L, ..., :N] # (L, ..., N)
    if c is not None:
        x = torch.einsum('l...n, ...n -> ...l', x, c)
    else:
        x = rearrange(x, 'l ... n -> ... n l')
    x = x.contiguous()
    return x


def Activation(activation=None, size=None, dim=-1, inplace=False):
    if activation in [ None, 'id', 'identity', 'linear' ]:
        return nn.Identity(inplace)
    elif activation == 'tanh':
        return nn.Tanh(inplace)
    elif activation == 'relu':
        return nn.ReLU(inplace)
    elif activation == 'gelu':
        return nn.GELU()
    elif activation in ['swish', 'silu']:
        return nn.SiLU(inplace)
    elif activation == 'glu':
        return nn.GLU(dim=dim)
    elif activation == 'sigmoid':
        return nn.Sigmoid(inplace)
    else:
        raise NotImplementedError("hidden activation '{}' is not implemented".format(activation))
        
        
class DropoutNd(nn.Module):
    def __init__(self, p: float = 0.5, tie=True, transposed=True):
        """
        tie: tie dropout mask across sequence lengths (Dropout1d/2d/3d)
        """
        super().__init__()
        if p < 0 or p >= 1:
            raise ValueError("dropout probability has to be in [0, 1), " "but got {}".format(p))
        self.p = p
        self.tie = tie
        self.transposed = transposed
        self.binomial = torch.distributions.binomial.Binomial(probs=1-self.p)

    def forward(self, x):
        """ x: (batch, lengths..., dim) """
        if self.training:
            if self.transposed: x = rearrange(x, 'b ... d -> b d ...')
            mask_shape = x.shape[:2] + (1,)*(x.ndim-2) if self.tie else x.shape
            mask = torch.rand(*mask_shape, device=x.device) < 1.-self.p
            x = x * mask * (1.0/(1-self.p))
            if self.transposed: x = rearrange(x, 'b d ... -> b ... d')
            return x
        return x


# class IdentitySSM(SSM):
#     def __init__(self, **kwargs):
#         super().__init__(**kwargs)
        
#     def init_weights(self):
#         self.register('kernel', None, trainable=False)
        
#     def forward(self, u):
#         return u


def init_preprocess_ssm(config):
    if config['method'] == 'differencing':
        ssm = DifferencingSSM
    elif config['method'] == 'ma_residual':
        ssm = MovingAvgResidualSSM
    elif config['method'] == 'residual':
        ssm = ResidualSSM
    elif config['method'] in ['identity', None]:
        return nn.Identity()
    else:
        raise NotImplementedError(f"Preprocessing config method {config['method']} not implemented!")
    return ssm(**config['kwargs'])


class MLP(nn.Module):
    def __init__(self,
                 input_dim: int,                 
                 output_dim: int,
                 activation: str=None,
                 dropout: float=0.,
                 layernorm: bool=False,
                 n_layers: int=1,
                 n_activations: int=0,
                 pre_activation: bool=False,
                 input_shape: str='bld',
                 hidden_dim: int=None,
                 skip_connection: bool=False,
                 average_pool: str=None):
        """
        Fully-connected network 
        """
        super().__init__()
        self.input_dim     = input_dim
        self.hidden_dim    = hidden_dim
        self.output_dim    = output_dim
        self.input_shape   = input_shape
        
        self.activation      = activation
        self.dropout         = dropout
        self.layernorm       = nn.LayerNorm(input_dim) if layernorm else nn.Identity()
        self.n_layers        = n_layers
        self.n_activations   = n_activations
        self.pre_activation  = pre_activation
        self.skip_connection = skip_connection
        self.average_pool    = average_pool
        
        self.initialize_layers()
        
    def initialize_layers(self):
        n_layers_to_init = self.n_layers
        n_activations_to_init = self.n_activations
        
        if self.hidden_dim is None:
            if self.n_layers < 2:
                self.hidden_dim = self.output_dim
            else:
                self.hidden_dim = self.input_dim
            
        # Add layers
        if self.n_activations > self.n_layers or self.pre_activation:
            layers = [Activation(self.activation, inplace=True), self.init_dropout()]
            n_activations_to_init -= 1
        else:
            layers = []
            
        while n_layers_to_init > 0 or n_activations_to_init > 0:
            if n_layers_to_init == self.n_layers:
                layers.append(nn.Linear(self.input_dim, self.hidden_dim))
            elif n_layers_to_init > 1:
                layers.append(nn.Linear(self.hidden_dim, self.hidden_dim))
            elif n_layers_to_init == 1:
                layers.append(nn.Linear(self.hidden_dim, self.output_dim))
            
            if n_activations_to_init > 0:
                layers.append(Activation(self.activation, inplace=True))
            
            n_layers_to_init -= 1
            n_activations_to_init -= 1
            
        self.layers = nn.Sequential(*layers)

        
    def init_dropout(self):
        if self.dropout > 1:  # Dropout hack for now, testing DropoutNd
            return DropoutNd(p=self.dropout-1.)
        elif self.dropout > 0:
            return nn.Dropout(self.dropout)
        else:
            return nn.Identity()
        
        
    def forward(self, x):
        x = self.layernorm(x)
        
        if self.input_shape == 'bdl':
            x = rearrange(x, 'b d l -> b l d')
        
        if self.skip_connection:
            # Layernorm with skip connection
            x = self.layers(x) + x  
        else: 
            x = self.layers(x)
        
        if self.average_pool == 'l':
            x = x.mean(dim=1, keepdim=True)
        return x



def init_ssm(config):
    supported_methods = ['companion', 'closed_loop_companion',
                         'shift', 'closed_loop_shift']
    if config['method'] == 'companion':
        ssm = CompanionSSM
    elif config['method'] == 'closed_loop_companion':
        ssm = ClosedLoopCompanionSSM
    elif config['method'] == 'shift':
        ssm = ShiftSSM
    elif config['method'] == 'closed_loop_shift':
        ssm = ClosedLoopShiftSSM
    else:
        raise NotImplementedError(
            f"SSM config method {config['method']} not implemented! Please choose from {supported_methods}")
    return ssm(**config['kwargs'])


def init_mlp(config):
    if config['method'] == 'mlp':
        return MLP(**config['kwargs'])
    else:
        return nn.Identity()


class Block(OurModule):
    """
    Standard encoder block
    """
    def __init__(self, 
                 input_dim: int,
                 pre_config: str=None,
                 ssm_config: str=None,
                 mlp_config: str=None,
                 skip_connection: bool=False,
                 skip_preprocess: bool=False):
        super().__init__()
        self.input_dim = input_dim
        self.skip_connection = skip_connection
        self.skip_preprocess = skip_preprocess
        
        self.pre = init_preprocess_ssm(pre_config)
        self.ssm = init_ssm(ssm_config)
        self.mlp = init_mlp(mlp_config)
            
    def forward(self, u):
        """
        Input shape: B x L x D
        """
        z = self.pre(u)
        y = self.ssm(z)
        y = self.mlp(y)
        if self.skip_connection and self.skip_preprocess:
            return y + u  # Also skip preprocessing step
        elif self.skip_connection:
            return y + z
        else:
            return y
            

class ClosedLoopBlock(Block):
    """
    Block with a closed-loop SSM. 
    
    In SpaceTime, we only consider using one ClosedLoopBlock 
    as the last-layer in a single-layer decoder. 
    
    However, other architectures can also be explored, e.g., 
    having more "open" blocks on top of the ClosedLoopBlock 
    in a multi-layer decoder.
    """
    def __init__(self, **kwargs):
        kwargs['skip_connection'] = False
        super().__init__(**kwargs)
        
    def forward(self, u):
        z = self.pre(u)
        # Computes layer outputs and next-time-step layer inputs
        y, u_next = self.ssm(z)  
        # Return both layer outputs and prediction + "ground-truth"
        # for next-time-step layer inputs
        return y, (u_next, u)


class Encoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.blocks = self.init_blocks(config)
        
    def init_blocks(self, config):
        blocks = []
        for block in config['blocks']:
            blocks.append(Block(**block))
        return nn.Sequential(*blocks)
    
    def forward(self, x):
        return self.blocks(x)
    
    
class Decoder(nn.Module):
    """
    In SpaceTime, we only consider using one ClosedLoopBlock 
    as the last-layer in a single-layer decoder. 
    
    However, other architectures can also be explored, e.g., 
    having more "open" blocks on top of the ClosedLoopBlock 
    in a multi-layer decoder.
    
    In future, can refactor this class to be more general 
    and support multiple layers. (p easy, just weirdness with
    nn.Sequential and multiple outputs)
    """
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.blocks = self.init_blocks(config)
        
    def init_blocks(self, config):
        return ClosedLoopBlock(**config['blocks'][0])
    
    def forward(self, x):
        return self.blocks(x)  # y, (u_next, u)


class Embedding(nn.Module):
    def __init__(self,
                 input_dim: int,
                 embedding_dim: int):
        """
        Generic class for encoding 
        """
        super().__init__()
        self.input_dim = input_dim
        self.embedding_dim = embedding_dim
        self.initialize_layers()
        
    def initialize_layers(self):
        self.layers = nn.Identity()
        
    def forward(self, x):
        return self.layers(x)
    

class LinearEmbedding(Embedding):
    def __init__(self, input_dim, embedding_dim):
        super().__init__(input_dim, embedding_dim)
        
    def initialize_layers(self):  
        self.layers = nn.Linear(self.input_dim, self.embedding_dim)


class RepeatEmbedding(Embedding):
    def __init__(self, 
                 input_dim: int, 
                 embedding_dim: int=None, 
                 n_heads: int=None,
                 n_kernels: int=None):

        if embedding_dim is None:
            try:
                embedding_dim = input_dim * n_heads * n_kernels
            except Exception as e:
                raise e('If embedding_dim not specified, must specify n_kernels and n_heads')
        else:
            assert embedding_dim % input_dim == 0, 'Embedding_dim should be multiple of input_dim'
        
        super().__init__(input_dim, embedding_dim)
        
    def repeat(self, x):
        return repeat(x, 'b l d -> b l (r d)', 
                      r=self.embedding_dim // self.input_dim)
        
    def initialize_layers(self):  
        self.layers = self.repeat
        


def init_embedding(config):
    methods = ['linear', 'identity', 'repeat']
    if config['method'] == 'linear':
        return LinearEmbedding(**config['kwargs'])
    elif config['method'] == 'repeat':
        return RepeatEmbedding(**config['kwargs'])
    elif config['method'] == 'identity' or config['method'] is None:
        return Embedding(**config['kwargs'])
    else:
        raise NotImplementedError(f"Embedding method {config['method']} not implemented. Please select among {methods}")


src/model/spacetime/ssm.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
# flake8: noqa

# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import contextlib
import copy
import threading
from typing import Iterable

import torch


@torch.no_grad()
def ema_update(ema_model_tuple, current_model_tuple, decay):
    torch._foreach_mul_(ema_model_tuple, decay)
    torch._foreach_add_(
        ema_model_tuple,
        current_model_tuple,
        alpha=(1.0 - decay),
    )


def run_ema_update_cpu(
    ema_model_tuple, current_model_tuple, decay, pre_sync_stream=None
):
    if pre_sync_stream is not None:
        pre_sync_stream.synchronize()

    ema_update(ema_model_tuple, current_model_tuple, decay)


class EMAOptimizer(torch.optim.Optimizer):
    r"""
    EMAOptimizer is a wrapper for torch.optim.Optimizer that computes
    Exponential Moving Average of parameters registered in the optimizer.

    EMA parameters are automatically updated after every step of the optimizer
    with the following formula:

        ema_weight = decay * ema_weight + (1 - decay) * training_weight

    To access EMA parameters, use ``swap_ema_weights()`` context manager to
    perform a temporary in-place swap of regular parameters with EMA
    parameters.

    Notes:
        - EMAOptimizer is not compatible with APEX AMP O2.

    Args:
        optimizer (torch.optim.Optimizer): optimizer to wrap
        device (torch.device): device for EMA parameters
        decay (float): decay factor

    Returns:
        returns an instance of torch.optim.Optimizer that computes EMA of
        parameters

    Example:
        model = Model().to(device)
        opt = torch.optim.Adam(model.parameters())

        opt = EMAOptimizer(opt, device, 0.9999)

        for epoch in range(epochs):
            training_loop(model, opt)

            regular_eval_accuracy = evaluate(model)

            with opt.swap_ema_weights():
                ema_eval_accuracy = evaluate(model)
    """

    def __init__(
        self,
        optimizer: torch.optim.Optimizer,
        device: torch.device,
        decay: float = 0.9999,
        every_n_steps: int = 1,
        current_step: int = 0,
    ):
        self.optimizer = optimizer
        self.decay = decay
        self.device = device
        self.current_step = current_step
        self.every_n_steps = every_n_steps
        self.save_original_optimizer_state = False

        self.first_iteration = True
        self.rebuild_ema_params = True
        self.stream = None
        self.thread = None

        self.ema_params = ()
        self.in_saving_ema_model_context = False

    def all_parameters(self) -> Iterable[torch.Tensor]:
        return (
            param for group in self.param_groups for param in group["params"]
        )

    def step(self, closure=None, **kwargs):
        self.join()

        if self.first_iteration:
            if any(p.is_cuda for p in self.all_parameters()):
                self.stream = torch.cuda.Stream()

            self.first_iteration = False

        if self.rebuild_ema_params:
            opt_params = list(self.all_parameters())

            self.ema_params += tuple(
                copy.deepcopy(param.data.detach()).to(self.device)
                for param in opt_params[len(self.ema_params) :]
            )
            self.rebuild_ema_params = False

        loss = self.optimizer.step(closure)

        if self._should_update_at_step():
            self.update()
        self.current_step += 1
        return loss

    def _should_update_at_step(self) -> bool:
        return self.current_step % self.every_n_steps == 0

    @torch.no_grad()
    def update(self):
        if self.stream is not None:
            self.stream.wait_stream(torch.cuda.current_stream())

        with torch.cuda.stream(self.stream):
            current_model_state = tuple(
                param.data.to(self.device, non_blocking=True)
                for param in self.all_parameters()
            )

            if self.device.type == "cuda":
                ema_update(self.ema_params, current_model_state, self.decay)

        if self.device.type == "cpu":
            self.thread = threading.Thread(
                target=run_ema_update_cpu,
                args=(
                    self.ema_params,
                    current_model_state,
                    self.decay,
                    self.stream,
                ),
            )
            self.thread.start()

    def swap_tensors(self, tensor1, tensor2):
        tmp = torch.empty_like(tensor1)
        tmp.copy_(tensor1)
        tensor1.copy_(tensor2)
        tensor2.copy_(tmp)

    def switch_main_parameter_weights(self, saving_ema_model: bool = False):
        self.join()
        self.in_saving_ema_model_context = saving_ema_model
        for param, ema_param in zip(self.all_parameters(), self.ema_params):
            self.swap_tensors(param.data, ema_param)

    @contextlib.contextmanager
    def swap_ema_weights(self, enabled: bool = True):
        r"""
        A context manager to in-place swap regular parameters with EMA
        parameters.
        It swaps back to the original regular parameters on context manager
        exit.

        Args:
            enabled (bool): whether the swap should be performed
        """

        if enabled:
            self.switch_main_parameter_weights()
        try:
            yield
        finally:
            if enabled:
                self.switch_main_parameter_weights()

    def __getattr__(self, name):
        return getattr(self.optimizer, name)

    def join(self):
        if self.stream is not None:
            self.stream.synchronize()

        if self.thread is not None:
            self.thread.join()

    def state_dict(self):
        self.join()

        if self.save_original_optimizer_state:
            return self.optimizer.state_dict()

        # if we are in the context of saving an EMA model,
        # the EMA weights are in the modules' actual weights
        ema_params = (
            self.ema_params
            if not self.in_saving_ema_model_context
            else list(self.all_parameters())
        )
        state_dict = {
            "opt": self.optimizer.state_dict(),
            "ema": ema_params,
            "current_step": self.current_step,
            "decay": self.decay,
            "every_n_steps": self.every_n_steps,
        }
        return state_dict

    def load_state_dict(self, state_dict):
        self.join()

        self.optimizer.load_state_dict(state_dict["opt"])
        self.ema_params = tuple(
            param.to(self.device) for param in copy.deepcopy(state_dict["ema"])
        )
        self.current_step = state_dict["current_step"]
        self.decay = state_dict["decay"]
        self.every_n_steps = state_dict["every_n_steps"]
        self.rebuild_ema_params = False

    def add_param_group(self, param_group):
        self.optimizer.add_param_group(param_group)
        self.rebuild_ema_params = True

src/optimizer.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import copy
from math import ceil, floor

import numpy as np
import torch

import utils
from model import losses
from utils.model_results import ModelResults


def default_metrics():
    return [
        losses.mse,
        losses.wmape_loss,
        losses.mse_far,
        losses.mae,
        losses.mae_far,
    ]


class TSData:
    def __init__(
        self,
        dataset,
        timestamps,
        obs_flag,
        out_frac,
    ):
        self.dataset = dataset
        self.timestamps = timestamps
        self.obs_flag = obs_flag
        self.out_frac = out_frac

        self._check_init_validity()

    def _check_init_validity(self):
        assert torch.is_tensor(self.dataset)
        assert torch.is_tensor(self.timestamps)
        assert torch.is_tensor(self.obs_flag), f"obs_flag is {self.obs_flag}"

        assert len(self.dataset.shape) == 3
        assert len(self.timestamps.shape) == 2
        assert len(self.obs_flag.shape) == 2

        shp0 = self.dataset.shape[0]
        shp1 = self.dataset.shape[1]

        assert self.dataset.shape[0] == shp0
        assert self.timestamps.shape[0] == shp0
        assert self.obs_flag.shape[0] == shp0

        assert self.dataset.shape[1] == shp1
        assert self.timestamps.shape[1] == shp1
        assert self.obs_flag.shape[1] == shp1

        assert torch.all(self.timestamps[:, :-1] < self.timestamps[:, 1:])

    @property
    def ts_size(self):
        ts_max = torch.max(self.timestamps, dim=1)
        ts_min = torch.min(self.timestamps, dim=1)
        ret_val = (ts_max.values - ts_min.values) * 100
        return ret_val - 1

    def augumented(self, seq_len, aug_step):
        augumented_dataset = utils.augument_dataset(
            self.dataset, seq_len, aug_step
        )
        reshaped_timestamps = self.timestamps.view(
            self.timestamps.shape[0], self.timestamps.shape[1], 1
        )
        augumented_timestamps = utils.augument_dataset(
            reshaped_timestamps, seq_len, aug_step
        ).view(-1, seq_len)
        reshaped_obs_flag = self.obs_flag.view(
            self.obs_flag.shape[0], self.obs_flag.shape[1], 1
        )
        augumented_obs_flag = utils.augument_dataset(
            reshaped_obs_flag, seq_len, aug_step
        ).view(-1, seq_len)

        ret_val = TSData(
            dataset=augumented_dataset,
            timestamps=augumented_timestamps,
            obs_flag=augumented_obs_flag,
            out_frac=self.out_frac,
        )

        return ret_val

    def shuffled(self, cap=100000):
        idxs = np.arange(self.shape[0])
        np.random.shuffle(idxs)
        idxs = torch.tensor(idxs[:cap])

        shuffled_dataset = self.dataset[idxs]
        shuffled_timestamps = self.timestamps[idxs]
        shuffled_obs_flag = (
            None if self.obs_flag is None else self.obs_flag[idxs]
        )

        return TSData(
            dataset=shuffled_dataset,
            timestamps=shuffled_timestamps,
            obs_flag=shuffled_obs_flag,
            out_frac=self.out_frac,
        )

    @property
    def shape(self):
        return self.dataset.shape

    def generate_samples(self, batch_size):
        dataset_size = self.shape[0]
        no_fits = ceil(dataset_size / batch_size)
        batch_size = min(batch_size, dataset_size)
        for itr in range(no_fits):
            base = (itr % (dataset_size // batch_size)) * batch_size
            yielded_dataset = self.dataset[base : base + batch_size].to(
                utils.DEFAULT_DEVICE
            )
            yielded_timestamps = self.timestamps[base : base + batch_size].to(
                utils.DEFAULT_DEVICE
            )
            yielded_obs_flag = (
                None
                if self.obs_flag is None
                else self.obs_flag[base : base + batch_size].to(
                    utils.DEFAULT_DEVICE
                )
            )

            tsdata = TSData(
                dataset=yielded_dataset,
                timestamps=yielded_timestamps,
                obs_flag=yielded_obs_flag,
                out_frac=None,
            )

            if self.out_frac is None:
                yield tsdata
            else:
                raise NotImplementedError

    def truncated_by_obs(self, trunc_len, out_frac):
        """
        returns a TSData with trunc_len first observations
        """
        if trunc_len is None:
            return copy.copy(self)
        else:
            return TSData(
                dataset=self.dataset[:, :trunc_len],
                timestamps=self.timestamps[:, :trunc_len],
                obs_flag=None
                if self.obs_flag is None
                else self.obs_flag[:, :trunc_len],
                out_frac=out_frac,
            )

    def truncated_by_time(self, trunc_len):
        raise NotImplementedError

    def random_removed_at(self, random_start, random_end, out_frac, fill_na):
        np.random.seed(42)
        full_batch_size = self.shape[0]
        work_len = (
            self.shape[1]
            if random_start is None or random_end is None
            else random_end - random_start
        )
        random_start = 0 if random_start is None else random_start
        random_end = self.shape[1] if random_end is None else random_end
        randomized_idxs = np.arange(work_len * full_batch_size)
        np.random.shuffle(randomized_idxs)
        randomized_idxs = torch.tensor(
            randomized_idxs.reshape(full_batch_size, work_len)
        )
        randomized_idxs = torch.sort(randomized_idxs, dim=-1).indices
        out_idxs_len = floor(work_len * out_frac)
        randomized_idxs = randomized_idxs[:, out_idxs_len:]
        randomized_idxs = torch.sort(randomized_idxs).values + random_start

        first_idxs = np.arange(random_start * full_batch_size) % random_start
        first_idxs = torch.tensor(
            first_idxs.reshape(full_batch_size, random_start)
        )

        last_len = self.shape[1] - random_end
        last_idxs = np.arange(last_len * full_batch_size) % last_len
        last_idxs = (
            torch.tensor(last_idxs.reshape(full_batch_size, last_len))
            + random_end
        )

        final_idxs = torch.cat(
            [first_idxs, randomized_idxs, last_idxs], axis=1
        )

        new_dataset = self.dataset.clone() * 0.0 + fill_na
        new_obs_flag = self.dataset[:, :, 0].clone() * 0.0

        for i in range(final_idxs.shape[0]):
            new_dataset[i : i + 1, final_idxs[i]] = self.dataset[
                i : i + 1, final_idxs[i]
            ]
            new_obs_flag[i : i + 1, final_idxs[i]] = 1.0

        return TSData(
            dataset=new_dataset,
            timestamps=self.timestamps,
            obs_flag=new_obs_flag,
            out_frac=None,
        )


class ProblemSpec:
    def __init__(
        self,
        train_dataset,
        test_dataset,
        train_timestamps,
        test_timestamps,
        train_out_frac,
        train_obs_flag=None,
        test_obs_flag=None,
        trunc=None,
        metrics=default_metrics(),
        norm=True,
        less_training_trajectories=None,
    ):
        if train_obs_flag is None:
            train_obs_flag = torch.ones(
                train_timestamps.shape, dtype=torch.float32
            )
        if test_obs_flag is None:
            test_obs_flag = torch.ones(
                test_timestamps.shape, dtype=torch.float32
            )

        self.train_tsdata = TSData(
            train_dataset,
            train_timestamps,
            obs_flag=train_obs_flag,
            out_frac=train_out_frac,
        )
        self.test_tsdata = TSData(
            test_dataset,
            test_timestamps,
            obs_flag=test_obs_flag,
            out_frac=None,
        )
        self.metrics = metrics

        if norm:
            train_dataset_for_scaling = self.train_tsdata.dataset
            reshaped = train_dataset_for_scaling.reshape(
                -1, train_dataset_for_scaling.size(2)
            )

            mean = torch.mean(reshaped, dim=0)
            std = torch.std(reshaped, dim=0)

            self.mean = mean
            self.std = std

            def norm_data(data):
                n_data = data - self.mean
                n_data = n_data / self.std
                return n_data

            self.train_tsdata.dataset = norm_data(self.train_tsdata.dataset)
            self.test_tsdata.dataset = norm_data(self.test_tsdata.dataset)

        self.train_tsdata = self.train_tsdata.truncated_by_obs(
            trunc, train_out_frac
        )
        self.test_tsdata = self.test_tsdata.truncated_by_obs(trunc, None)

        if less_training_trajectories:
            self.train_tsdata = self.train_tsdata.shuffled(cap=less_training_trajectories)

    def random_removed_at(self, random_range, out_frac, fill_na):
        if random_range is None:
            random_start, random_end = None, None
        else:
            random_start, random_end = random_range

        raise Exception(
            f"check if random_start and random_end is right {random_start} - {random_end}"  # noqa
        )
        train_ = self.train_tsdata
        test_ = self.test_tsdata.random_removed_at(
            random_start, random_end, out_frac, fill_na
        )

        ps = ProblemSpec(
            train_dataset=train_.dataset,
            test_dataset=test_.dataset,
            train_timestamps=train_.timestamps,
            test_timestamps=test_.timestamps,
            train_obs_flag=train_.obs_flag,
            test_obs_flag=test_.obs_flag,
            metrics=self.metrics,
            train_out_frac=out_frac,
            norm=False,
        )

        return ps

    def test_a_dataset(
        self,
        trainer,
        batch_size,
        tsdata,
        discrete_updates,
        method,
        atol,
        rtol,
        min_dt,
        sample_vae,
    ):
        assert isinstance(tsdata, TSData)

        no_fits = 0
        computed_metrics_all = []
        steps = []

        len_processed = 0
        for sampled_tsdata in tsdata.generate_samples(batch_size):
            sample_len = sampled_tsdata.shape[0]
            len_processed += sample_len
            no_fits += 1
            result = trainer.model.predict(
                sampled_tsdata,
                discrete_updates=discrete_updates,
                method=method,
                atol=atol,
                rtol=rtol,
                min_dt=min_dt,
                sample_vae=sample_vae,
            )
            pred_x = result["pred_x"]
            if "steps" in result.keys():
                steps_ = result["steps"]
                steps.append(steps_)
            computed_metrics = [
                (
                    metric.__name__,
                    metric(
                        pred_x,
                        sampled_tsdata.dataset,
                        sampled_tsdata.obs_flag,
                        first=trainer.encoded_len,
                    ).item(),
                )
                for metric in self.all_metrics()
            ]
            computed_metrics_all.append(computed_metrics)

        computed_metrics_avg = []

        for i in range(len(computed_metrics_all[0])):
            name = computed_metrics_all[0][i][0]
            sum_ = 0
            for j in range(no_fits):
                sum_ += computed_metrics_all[j][i][1]

            avg = sum_ / no_fits
            computed_metrics_avg.append(
                (
                    name,
                    avg,
                )
            )

        model_results = ModelResults({k: v for k, v in computed_metrics_avg})

        return model_results

    def score_a_small_dataset(
        self,
        trainer,
        dataset,
        samp_ts,
    ):
        dataset = torch.from_numpy(dataset)
        dataset = dataset.to(utils.DEFAULT_DEVICE)
        pred_x, _ = trainer.model.predict(
            dataset,
            samp_ts,
        )
        return pred_x.detach().cpu().numpy()

    def all_metrics(self):
        return self.metrics

    def test(
        self,
        trainer,
        batch_size,
        discrete_updates,
        method,
        data_len,
        atol,
        rtol,
        min_dt,
        sample_vae,
    ):
        return self.test_a_dataset(
            trainer,
            batch_size,
            tsdata=self.test_tsdata.shuffled(cap=data_len),
            discrete_updates=discrete_updates,
            method=method,
            atol=atol,
            rtol=rtol,
            min_dt=min_dt,
            sample_vae=sample_vae,
        )

    def test_train(
        self,
        trainer,
        batch_size,
        discrete_updates,
        method,
        data_len,
        atol,
        rtol,
        min_dt,
        sample_vae,
    ):
        return self.test_a_dataset(
            trainer,
            batch_size,
            tsdata=self.train_tsdata.shuffled(cap=data_len),
            discrete_updates=discrete_updates,
            method=method,
            atol=atol,
            rtol=rtol,
            min_dt=min_dt,
            sample_vae=sample_vae,
        )

src/problem_specs/__init__.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import h5py
import numpy as np
import torch

from problem_specs import ProblemSpec


def get_obs_dim(problem_spec_name):
    if problem_spec_name == "cheetah_medium_dense_1000":
        return 17
    elif problem_spec_name == "cheetah_medium_dense_500":
        return 17
    if (problem_spec_name == "large_cheetah_medium_dense_1000") or (problem_spec_name == "large_cheetah_medium_deter_1000"):
        return 17
    elif problem_spec_name == "large_cheetah_medium_dense_1000_shorter":
        return 17
    elif problem_spec_name == "large_cheetah_medium_dense_500":
        return 17
    elif problem_spec_name == "hopper_medium_dense_1000":
        return 11
    elif problem_spec_name == "hopper_medium_dense_500":
        return 11
    elif (problem_spec_name == "large_hopper_medium_dense_1000") or (problem_spec_name == "large_hopper_medium_deter_1000"):
        return 11
    elif problem_spec_name == "large_hopper_medium_dense_500":
        return 11
    elif problem_spec_name == "ks100_500":
        return 100
    elif problem_spec_name == "ks100_1000":
        return 100
    elif problem_spec_name == "large_ks100_500":
        return 100
    elif problem_spec_name == "large_ks100_1000":
        return 100
    elif problem_spec_name == "large_turbulence_1000":
        return 256
    elif problem_spec_name == "walker_medium_dense_1000":
        return 17
    elif problem_spec_name == "walker_medium_dense_500":
        return 17
    elif (problem_spec_name == "large_walker_medium_dense_1000") or (problem_spec_name == "large_walker_medium_deter_1000"):
        return 17
    elif problem_spec_name == "large_walker_medium_dense_500":
        return 17
    elif problem_spec_name == "electricity":
        return 1
    elif problem_spec_name == "electricity_short":
        return 1
    elif problem_spec_name == "m4_247":
        return 1
    elif problem_spec_name == "m4_short":
        return 1
    elif problem_spec_name == "tourism_228":
        return 1
    elif problem_spec_name == "tourism_short":
        return 1
    elif problem_spec_name == "large_mackey_glass_2000":
        return 1
    elif problem_spec_name == "large_mackey_glass_2000_shorter":
        return 1
    elif problem_spec_name == "large_lorenz_2000":
        return 3
    elif problem_spec_name == "large_mso_2000":
        return 1
    elif problem_spec_name == "large_sspiral_2000":
        return 2
    elif problem_spec_name == "ett_small":
        return 1
    elif problem_spec_name == "ett_small_short":
        return 1
    elif problem_spec_name == "ett_long":
        return 1
    elif problem_spec_name == "ett_long_short":
        return 1
    elif problem_spec_name == "large_cahn_hillard_1000":
        return 256
    elif problem_spec_name == "large_lotka_2000":
        return 2
    elif problem_spec_name == "pems_bay":
        return 325
    elif problem_spec_name == "pems_bay_short":
        return 325
    elif problem_spec_name == "weather":
        return 21
    else:
        raise ValueError(
            f"obs dim unknown for problem spec: {problem_spec_name}"
        )


def problem_spec_from_hdf5_standard(
    file_path,
    trunc=None,
    shorter=None,
):
    f = h5py.File(file_path, "r")

    train_dataset = np.array(f["train_data"])
    test_dataset = np.array(f["test_data"])
    train_timestamps = np.array(f["train_timestamps"])[
        :, : train_dataset.shape[1]
    ]
    test_timestamps = np.array(f["test_timestamps"])[
        :, : test_dataset.shape[1]
    ]
    f.close()

    problem_spec = ProblemSpec(
        train_dataset=torch.from_numpy(train_dataset).float(),
        test_dataset=torch.from_numpy(test_dataset).float(),
        train_timestamps=torch.from_numpy(train_timestamps).float(),
        test_timestamps=torch.from_numpy(test_timestamps).float(),
        train_out_frac=None,
        trunc=trunc,
        less_training_trajectories=shorter,
    )

    return problem_spec

def weather():
    return problem_spec_from_hdf5_standard(
        file_path="./data/weather/dataset.hdf5",
        trunc=1000,
    )

def ett_small():
    return problem_spec_from_hdf5_standard(
        file_path="./data/ett/dataset.hdf5",
        trunc=1440,
    )

def ett_small_short():
    return problem_spec_from_hdf5_standard(
        file_path="./data/ett/dataset.hdf5",
        trunc=1440,
        shorter=1000,
    )

def m4_short():
    return problem_spec_from_hdf5_standard(
        file_path="./data/m4/dataset.hdf5",
        trunc=247,
        shorter=1000,
    )

def ett_long():
    return problem_spec_from_hdf5_standard(
        file_path="./data/ett_long/dataset.hdf5",
        trunc=2000,
    )

def ett_long_short():
    return problem_spec_from_hdf5_standard(
        file_path="./data/ett_long/dataset.hdf5",
        trunc=2000,
        shorter=1000,
    )


def cheetah_medium_dense_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/cheetah_medium_dense/dataset.hdf5",
        trunc=1000,
    )


def cheetah_medium_dense_500():
    return problem_spec_from_hdf5_standard(
        file_path="./data/cheetah_medium_dense/dataset.hdf5",
        trunc=500,
    )


def large_cheetah_medium_dense_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/cheetah_medium_dense/dataset_large.hdf5",
        trunc=1000,
    )

def large_cheetah_medium_dense_1000_shorter():
    return problem_spec_from_hdf5_standard(
        file_path="./data/cheetah_medium_dense/dataset_large.hdf5",
        trunc=1000,
        shorter=1000,
    )

def large_cheetah_medium_deter_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/cheetah_medium_deter/dataset_large.hdf5",
        trunc=1000,
    )

def large_cheetah_medium_dense_500():
    return problem_spec_from_hdf5_standard(
        file_path="./data/cheetah_medium_dense/dataset_large.hdf5",
        trunc=500,
    )


def hopper_medium_dense_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/hopper_medium_dense/dataset.hdf5",
        trunc=1000,
    )


def hopper_medium_dense_500():
    return problem_spec_from_hdf5_standard(
        file_path="./data/hopper_medium_dense/dataset.hdf5",
        trunc=500,
    )


def large_hopper_medium_dense_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/hopper_medium_dense/dataset_large.hdf5",
        trunc=1000,
    )

def large_hopper_medium_deter_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/hopper_medium_deter/dataset_large.hdf5",
        trunc=1000,
    )

def large_hopper_medium_dense_500():
    return problem_spec_from_hdf5_standard(
        file_path="./data/hopper_medium_dense/dataset_large.hdf5",
        trunc=500,
    )


def walker_medium_dense_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/walker_medium_dense/dataset.hdf5",
        trunc=1000,
    )


def walker_medium_dense_500():
    return problem_spec_from_hdf5_standard(
        file_path="./data/walker_medium_dense/dataset.hdf5",
        trunc=500,
    )


def large_walker_medium_dense_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/walker_medium_dense/dataset_large.hdf5",
        trunc=1000,
    )

def large_walker_medium_deter_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/walker_medium_deter/dataset_large.hdf5",
        trunc=1000,
    )

def large_walker_medium_dense_500():
    return problem_spec_from_hdf5_standard(
        file_path="./data/walker_medium_dense/dataset_large.hdf5",
        trunc=500,
    )


def ks100_500():
    return problem_spec_from_hdf5_standard(
        file_path="./data/ks_100_hard/dataset.hdf5",
        trunc=500,
    )


def ks100_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/ks_100_hard/dataset.hdf5",
        trunc=1000,
    )


def large_ks100_500():
    return problem_spec_from_hdf5_standard(
        file_path="./data/ks_100_hard/dataset_large.hdf5",
        trunc=500,
    )


def large_ks100_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/ks_100_hard/dataset_large.hdf5",
        trunc=1000,
    )

def large_turbulence_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/turbulence/dataset_large.hdf5",
        trunc=1000,
    )

def large_cahn_hillard_1000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/cahn_hillard/dataset_large.hdf5",
        trunc=1000,
    )

def electricity():
    return problem_spec_from_hdf5_standard(
        file_path="./data/electricity/dataset.hdf5",
        trunc=1440,
    )


def electricity_short():
    return problem_spec_from_hdf5_standard(
        file_path="./data/electricity/dataset.hdf5",
        trunc=1440,
        shorter=1000,
    )


def pems_bay():
    return problem_spec_from_hdf5_standard(
        file_path="./data/pems_bay/dataset.hdf5",
        trunc=288,
    )


def pems_bay_short():
    return problem_spec_from_hdf5_standard(
        file_path="./data/pems_bay/dataset.hdf5",
        trunc=288,
        shorter=1000,
    )


def m4_247():
    return problem_spec_from_hdf5_standard(
        file_path="./data/m4/dataset.hdf5",
        trunc=247,
    )


def tourism_228():
    return problem_spec_from_hdf5_standard(
        file_path="./data/tourism/dataset.hdf5",
        trunc=228,
    )

def tourism_short():
    return problem_spec_from_hdf5_standard(
        file_path="./data/tourism/dataset.hdf5",
        trunc=228,
        shorter=1000,
    )

def large_mackey_glass_2000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/mackey_glass/dataset_large.hdf5",
        trunc=2000,
    )

def large_mackey_glass_2000_shorter():
    return problem_spec_from_hdf5_standard(
        file_path="./data/mackey_glass/dataset_large.hdf5",
        trunc=2000,
        shorter=1000,
    )

def large_lorenz_2000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/lorenz/dataset_large.hdf5",
        trunc=2000,
    )

def large_mso_2000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/mso/dataset_large.hdf5",
        trunc=2000,
    )

def large_sspiral_2000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/sspiral/dataset_large.hdf5",
        trunc=2000,
    )

def large_lotka_2000():
    return problem_spec_from_hdf5_standard(
        file_path="./data/lotka/dataset_large.hdf5",
        trunc=2000,
    )

src/problem_specs/instances.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
from datetime import datetime

from build_from_configuration import build_from_configuration
from utils import neptune_logging
from utils.parse_args import parse_args

import configurations  # noqa

args = parse_args()

neptune_logging.USE_NEPTUNE = args.use_neptune == 1

n1 = datetime.now()

configuration_fun = eval("configurations.{}".format(args.configuration_name))
configuration = configuration_fun(
    latent_dim=args.latent_dim,
    problem_spec_name=args.problem_spec_name,
    model_instance=args.model_instance,
    enc_dim=args.enc_dim,
    hidden_dim=args.hidden_dim,
    batch_size=args.batch_size,
    pretraining_len=args.pretraining_len,
    dec_depth=args.dec_depth,
    dec_width=args.dec_width,
)

print(configuration)

trainer = build_from_configuration(configuration, args.checkpoint_name)

trainer.load(args.load_checkpoint)

trainer.fit()

n2 = datetime.now()

print("finished in {}".format(n2 - n1))

src/train.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
# flake8: noqa
from .odeint_w_events_trainer import ODEINTWEventsTrainer
src/trainers/__init__.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import os

import torch
from torch import nn
from torch.nn.utils import clip_grad_norm_

import utils
from model import losses
from utils import neptune_logging, count_parameters
from utils.model_results import ModelResults


class BaseTrainer(nn.Module):
    def __init__(
        self,
        problem_spec,
        model,
        problem_spec_name,
        train_seq_len,
        batch_size,
        adjoint,
        max_lr,
        trainer_name,
        train_dir,
        configuration,
        solver_method,
        pretraining_epochs,
        increasing_seq_lens,
        trajectories_per_epoch,
        betas,
        increasing_seq_lens_epochs,
        test_train_data_len,
        test_data_len,
        entropy_loss,
        atol,
        rtol,
        min_dt,
        entropy_factor,
        seq_len_step,
        pretraining_len,
        train_enc_dec_after_pretraining,
        pretraining_ode_lr,
        checkpoint_name,
        aug_step,
        optimizer,
        add_kl,
        sample_vae,
        device=utils.DEFAULT_DEVICE,
    ):
        super(BaseTrainer, self).__init__()

        if not add_kl and sample_vae:
            raise ValueError("cannot sample vae when add_kl is False")

        self.model = model
        self.train_seq_len = train_seq_len
        self.batch_size = batch_size
        self.adjoint = adjoint
        self.max_lr = max_lr
        self.problem_spec = problem_spec
        self.solver_method = solver_method
        self.problem_spec_name = problem_spec_name
        self.trainer_name = trainer_name
        self.device = device
        self.train_dir = train_dir
        self.pretraining_epochs = pretraining_epochs
        self.increasing_seq_lens = increasing_seq_lens
        self.trajectories_per_epoch = trajectories_per_epoch
        self.betas = betas
        self.increasing_seq_lens_epochs = increasing_seq_lens_epochs
        self.test_train_data_len = test_train_data_len
        self.test_data_len = test_data_len
        self.entropy_loss = entropy_loss
        self.atol = atol
        self.rtol = rtol
        self.min_dt = min_dt
        self.entropy_factor = entropy_factor
        self.seq_len_step = seq_len_step
        self.pretraining_len = pretraining_len
        self.train_enc_dec_after_pretraining = train_enc_dec_after_pretraining
        self.pretraining_ode_lr = pretraining_ode_lr
        self.aug_step = aug_step
        self.optimizer = optimizer
        self.add_kl = add_kl
        self.sample_vae = sample_vae
        self.encoded_len = model.rec._rec_len

        self.best_val_loss = 99999999999999.0
        self.last_val_loss = None

        self.relu = torch.nn.functional.relu

        self._run = None

        self.put_mock_best_last_checkpoints()

        if neptune_logging.USE_NEPTUNE:
            self.init_neptune(
                configuration,
                checkpoint_name,
            )

    def put_mock_best_last_checkpoints(self):
        # raise NotImplementedError
        pass

    def init_neptune(self, configuration, checkpoint_name):
        run = neptune_logging.get_neptune_run(
            dict(
                problem_spec=self.problem_spec_name,
                trainer=self.trainer_name,
                configuration=configuration,
                checkpoint_name=checkpoint_name,
                number_of_parameters=count_parameters(self.model)
            ),
            track_files_path="./models_cpth/{}".format(
                configuration["problem"],
            ),
        )

        self._run = run

    @property
    def neptune_run(self):
        return self._run

    def train_run(self, *args, **kwargs):
        (
            train_results,
            train_set_results_hard,
            train_set_results_soft,
            test_set_results_hard,
            test_set_results_soft,
        ) = self.train_run_middle(*args, **kwargs)
        self.last_val_loss = test_set_results_soft["mse"]
        self.train_run_post()
        return (
            train_results,
            train_set_results_hard,
            train_set_results_soft,
            test_set_results_hard,
            test_set_results_soft,
        )

    def train_run_middle(self, *args, **kwargs):
        raise NotImplementedError

    def train_run_post(self):
        if self.last_val_loss < self.best_val_loss:
            self.best_val_loss = self.last_val_loss
            self.save_best_checkpoint()

    def fit_step(
        self,
        tsdata,
        metrics,
        optimizers,
        schedulers,
        method,
    ):
        for optimizer in optimizers:
            optimizer.zero_grad(set_to_none=True)

        result = self.model(
            tsdata=tsdata,
            method=method,
            atol=self.atol,
            rtol=self.rtol,
            min_dt=self.min_dt,
            sample_vae=self.sample_vae,
            discrete_updates=False,
        )
        pred_x = result["pred_x"]
        z0_mean = result["z0_mean"]
        z0_std = result["z0_std"]
        if self.model.model_type in ["standard_node"]:
            steps = result["steps"]

        if self.add_kl:
            error = losses.elbo_loss(
                pred_x,
                tsdata.dataset,
                (
                    z0_mean,
                    z0_std,
                ),
                try_omit_first=self.encoded_len,
            )
        else:
            error = losses.mse(
                pred_x,
                tsdata.dataset,
                obs_flag=tsdata.obs_flag,
                try_omit_first=self.encoded_len,
            )

        criterion = error

        if torch.isnan(criterion):
            criterion_nan = True
        else:
            criterion_nan = False

        metrics = {
            metric.__name__: metric(
                pred_x,
                tsdata.dataset,
                obs_flag=tsdata.obs_flag,
                first=self.encoded_len,
            ).item()
            for metric in metrics
        }

        if self.model.model_type in ["standard_node"]:
            aux_metrics = dict(
                T_su=self.model.T,
                criterion=criterion.item(),
                steps=steps.sum(dim=1).mean().item(),
            )
        else:
            aux_metrics = dict()

        if self.add_kl:
            aux_metrics["kl_divergence"] = (
                losses.kl_divergence(z0_mean, z0_std).mean().item()
            )

        all_metrics = {
            k: v
            for k, v in (list(metrics.items()) + list(aux_metrics.items()))
        }

        if not criterion_nan:
            criterion.backward()

            clip_grad_norm_(self.parameters(), 1.0, "inf", False)

            for optimizer in optimizers:
                optimizer.step()

            for scheduler in schedulers:
                scheduler.step()

        return ModelResults(all_metrics)

    def fit():
        raise NotImplementedError()

    def load(self, load_checkpoint, quiet=False):
        if load_checkpoint == "noload":
            return
        if self.train_dir is not None:
            train_dir = self.train_dir
            if not os.path.exists(train_dir):
                os.makedirs(train_dir)
            init_state_dict = self.state_dict()
            try:
                if load_checkpoint == "load_best":
                    ckpt_path = os.path.join(train_dir, "best_ckpt.pth")
                    checkpoint = torch.load(ckpt_path, map_location="cpu")
                    self.load_state_dict(checkpoint["state_dict"])
                    if not quiet:
                        print("Loaded ckpt from {}".format(ckpt_path))
                elif load_checkpoint == "load_last":
                    ckpt_path = os.path.join(train_dir, "last_ckpt.pth")
                    checkpoint = torch.load(ckpt_path, map_location="cpu")
                    self.load_state_dict(checkpoint["state_dict"])
                    if not quiet:
                        print("Loaded ckpt from {}".format(ckpt_path))
                else:
                    raise ValueError(
                        f"invalid load_checkpoint option: {load_checkpoint}"
                    )
            except Exception as error:
                self.load_state_dict(init_state_dict)
                if not quiet:
                    print(error)
                    print("failed to load")

    def save(self):
        utils.save_model(
            name="last_ckpt.pth",
            train_dir=self.train_dir,
            state_dict=self.state_dict(),
        )

    def save_best_checkpoint(self):
        utils.save_model(
            name="best_ckpt.pth",
            train_dir=self.train_dir,
            state_dict=self.state_dict(),
        )

src/trainers/base_trainer.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import random
from math import ceil
from time import time

import torch

import utils
from optimizer import EMAOptimizer
from utils import neptune_logging

from .base_trainer import BaseTrainer


class ODEINTWEventsTrainer(BaseTrainer):
    def train_run_middle(
        self,
        batch_size,
        ode_optimizer,
        schedulers,
        train_enc_dec,
        random_range,
        method,
        test=True,
        start_only=False,
    ):
        if random_range[1] - 1 == random_range[0]:
            random_seq_len = True
        else:
            random_seq_len = False

        self.model.train()

        t0 = time()
        train_tsdata = self.problem_spec.train_tsdata.shuffled(
            cap=self.trajectories_per_epoch
        )
        augumented_tsdata = train_tsdata.augumented(
            random_range[1] - 1, aug_step=self.aug_step
        )
        shuffled_tsdata = augumented_tsdata.shuffled(cap=200000)
        dataset_size = shuffled_tsdata.shape[0]
        train_batch_size = batch_size
        no_fits = ceil(dataset_size / train_batch_size)

        if train_enc_dec:
            if hasattr(self.model, "rec"):
                self.model.rec.set_freeze(False)
            if hasattr(self.model, "dec"):
                self.model.dec.set_freeze(False)
        else:
            if hasattr(self.model, "rec"):
                self.model.rec.set_freeze(True)
            if hasattr(self.model, "dec"):
                self.model.dec.set_freeze(True)

        epoch_stats = []
        for itr, sampled_tsdata in enumerate(
            shuffled_tsdata.generate_samples(int(train_batch_size))
        ):
            fit_id = itr + 1

            if random_seq_len:
                rand_seq_len = random.randint(random_range[0], random_range[1])
                sampled_tsdata = sampled_tsdata.truncated_by_obs(
                    rand_seq_len, None
                )

            stats = self.fit_step(
                tsdata=sampled_tsdata,
                metrics=self.problem_spec.metrics,
                optimizers=(ode_optimizer,),
                schedulers=schedulers,
                method=method,
            )
            print("fit {}/{}, {}".format(fit_id, no_fits, stats), end="\r")
            epoch_stats.append(stats.results)

        print()

        self.model.eval()

        mses = sum([results["mse"] for results in epoch_stats]) / len(
            epoch_stats
        )

        train_results = utils.model_results.ModelResults(
            {
                "train_mse": mses,
            }
        )

        self.save()

        t1 = time()
        train_set_results_hard = self.problem_spec.test_train(
            self,
            batch_size * 2,
            discrete_updates=True,
            method=method,
            data_len=self.test_train_data_len,
            atol=self.atol,
            rtol=self.rtol,
            min_dt=self.min_dt,
            sample_vae=self.sample_vae,
        )
        if test:
            test_set_results_hard = self.problem_spec.test(
                self,
                batch_size * 2,
                discrete_updates=True,
                method=method,
                data_len=self.test_data_len,
                atol=self.atol,
                rtol=self.rtol,
                min_dt=self.min_dt,
                sample_vae=self.sample_vae,
            )

            t2 = time()
            print("training time:", t1 - t0, "testing time:", t2 - t1)
            return (
                train_results,
                train_set_results_hard,
                train_set_results_hard,
                test_set_results_hard,
                test_set_results_hard,
            )
        else:
            t2 = time()
            print("training time:", t1 - t0, "testing time:", t2 - t1)
            return (
                train_results,
                train_set_results_hard,
                train_set_results_hard,
            )

    def fit(self):
        batch_size = self.batch_size
        ode_params = list(self.model.parameters())

        seq_len = self.train_seq_len

        if self.optimizer == "Adam":
            ode_optimizer = torch.optim.AdamW(
                params=ode_params,
                lr=self.max_lr,
                betas=self.betas,
            )
        elif self.optimizer == "Adamax":
            ode_optimizer = torch.optim.Adamax(
                params=ode_params, lr=self.max_lr
            )
        else:
            raise ValueError(f"unknown optimizer: {self.optimizer}")

        ode_optimizer = EMAOptimizer(
            ode_optimizer, utils.DEFAULT_DEVICE, decay=0.9999
        )
        seq_len_short = self.pretraining_len

        if self.increasing_seq_lens and seq_len_short is None:
            raise ValueError("--pretraining_len is not set")

        while self.increasing_seq_lens and seq_len_short < int(1.0 * seq_len):
            eps = (
                self.pretraining_epochs
                if seq_len_short == self.pretraining_len
                else self.increasing_seq_lens_epochs
            )
            print("eps:", eps)
            for epoch in range(0, eps):
                if seq_len_short == self.pretraining_len:
                    for g in ode_optimizer.param_groups:
                        g["lr"] = self.pretraining_ode_lr
                else:
                    for g in ode_optimizer.param_groups:
                        g["lr"] = self.max_lr

                print(f"small seq_len {seq_len_short} with su")
                print("epoch {}".format(epoch + 1))
                (
                    train_results,
                    train_set_results_hard,
                    train_set_results_soft,
                    test_set_results_hard,
                    test_set_results_soft,
                ) = self.train_run(
                    batch_size=batch_size,
                    random_range=(seq_len_short, seq_len_short + 1),
                    ode_optimizer=ode_optimizer,
                    schedulers=(),
                    train_enc_dec=True
                    if seq_len_short == self.pretraining_len
                    else self.train_enc_dec_after_pretraining,
                    test=True,
                    method=self.solver_method,
                )

                neptune_logging.log_w_prefix(
                    self.neptune_run, "train_short", train_results.results
                )
                neptune_logging.log_w_prefix(
                    self.neptune_run,
                    "train_hard",
                    train_set_results_hard.results,
                )
                neptune_logging.log_w_prefix(
                    self.neptune_run,
                    "train_soft",
                    train_set_results_soft.results,
                )
                neptune_logging.log_w_prefix(
                    self.neptune_run,
                    "test_hard",
                    test_set_results_hard.results,
                )
                neptune_logging.log_w_prefix(
                    self.neptune_run,
                    "test_soft",
                    test_set_results_soft.results,
                )

                print("train_results:", train_results)
                print("train_set_results_hard:", train_set_results_hard)
                print("train_set_results_soft:", train_set_results_soft)
                print("test_set_results_hard:", test_set_results_hard)
                print("test_set_results_soft:", test_set_results_soft)

            seq_len_short += self.seq_len_step

        for g in ode_optimizer.param_groups:
            g["lr"] = self.max_lr

        for epoch in range(1000000):
            print("full seq_len with su, seq_len: {}".format(seq_len))
            print("epoch {}".format(epoch + 1))
            (
                train_results,
                train_set_results_hard,
                train_set_results_soft,
                test_set_results_hard,
                test_set_results_soft,
            ) = self.train_run(
                random_range=(seq_len, seq_len + 1),
                batch_size=batch_size,
                ode_optimizer=ode_optimizer,
                train_enc_dec=True
                if self.train_enc_dec_after_pretraining
                else not self.increasing_seq_lens,
                schedulers=(),
                test=True,
                method=self.solver_method,
            )

            neptune_logging.log_w_prefix(
                self.neptune_run, "train_short", train_results.results
            )
            neptune_logging.log_w_prefix(
                self.neptune_run, "train_hard", train_set_results_hard.results
            )
            neptune_logging.log_w_prefix(
                self.neptune_run, "train_soft", train_set_results_soft.results
            )
            neptune_logging.log_w_prefix(
                self.neptune_run, "test_hard", test_set_results_hard.results
            )
            neptune_logging.log_w_prefix(
                self.neptune_run, "test_soft", test_set_results_soft.results
            )

            print("train_results:", train_results)
            print("train_set_results_hard:", train_set_results_hard)
            print("train_set_results_soft:", train_set_results_soft)
            print("test_set_results_hard:", test_set_results_hard)
            print("test_set_results_soft:", test_set_results_soft)

src/trainers/odeint_w_events_trainer.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import os

import matplotlib
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import torch
from IPython.display import HTML


def count_parameters(model):
    count_parameters = 0

    for param in model.parameters():
        count_parameters += param.contiguous().view(-1).shape[0]

    return count_parameters


class RunningAverageMeter:
    """Computes and stores the average and current value"""

    def __init__(self, momentum=0.99):
        self.momentum = momentum
        self.reset()

    def reset(self):
        self.avg = 0
        self.first = True

    def update(self, val):
        if self.first:
            self.avg = val
            self.first = False
        self.avg = self.avg * self.momentum + val * (1 - self.momentum)


def max_pool_im(ar, div):
    height = ar.shape[-3]
    width = ar.shape[-2]

    new_shape = list(ar.shape)
    new_shape[-3] = height // div
    new_shape[-2] = width // div
    new_shape[-1] = 1
    new_ar = np.zeros(new_shape)

    for i in range(height // div):
        for j in range(width // div):
            avg_slice = ar[
                :, :, i * div : i * div + div, j * div : j * div + div, :
            ]
            avg_slice = np.max(avg_slice, axis=2)
            avg_slice = np.max(avg_slice, axis=2)
            avg_slice = np.max(avg_slice, axis=2)

            new_ar[:, :, i, j, 0] = avg_slice

    return new_ar.astype("float32")


def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda:0")
    # elif torch.backends.mps.is_available():
        # return torch.device("mps")
    else:
        return torch.device("cpu")


DEFAULT_DEVICE = get_device()
DEFAULT_MISSING_VALUE = 0.0


def display_video(frames, framerate=30):
    height, width, _ = frames[0].shape
    dpi = 70
    orig_backend = matplotlib.get_backend()
    matplotlib.use(
        "Agg"
    )  # Switch to headless 'Agg' to inhibit figure rendering.
    fig, ax = plt.subplots(1, 1, figsize=(width / dpi, height / dpi), dpi=dpi)
    matplotlib.use(orig_backend)  # Switch back to the original backend.
    ax.set_axis_off()
    ax.set_aspect("equal")
    ax.set_position([0, 0, 1, 1])
    im = ax.imshow(frames[0])

    def update(frame):
        im.set_data(frame)
        return [im]

    interval = 1000 / framerate
    anim = animation.FuncAnimation(
        fig=fig,
        func=update,
        frames=frames,
        interval=interval,
        blit=True,
        repeat=False,
    )
    return HTML(anim.to_html5_video())


def get_ts(start, seq_len, device=DEFAULT_DEVICE):
    return torch.from_numpy(
        np.linspace(start, start + (seq_len - 1) / 100, seq_len).astype(
            "float32"
        )
    ).to(device)


def transform(ar, mean, scale):
    return (ar - mean) / scale


def inverse_transform(ar, mean, scale):
    return (ar * scale) + mean


def augument_dataset(orig_dataset, seq_len, aug_step):
    assert torch.is_tensor(orig_dataset)
    assert len(orig_dataset.shape) == 3
    assert type(seq_len) == int

    full_len = orig_dataset.shape[1]
    o_len = orig_dataset.shape[0]
    dataset_size = 1

    while dataset_size * aug_step + seq_len < full_len:
        dataset_size += 1

    dataset_size *= o_len
    augumented_dataset = torch.zeros(
        (
            dataset_size,
            seq_len,
            orig_dataset.shape[2],
        )
    )

    i = 0

    while i * aug_step + seq_len < full_len:
        augumented_dataset[i * o_len : (i + 1) * o_len] = orig_dataset[
            :, i * aug_step : i * aug_step + seq_len
        ]
        i += 1

    end_dataset = orig_dataset[:, -seq_len:]
    augumented_dataset[-o_len:] = end_dataset

    return augumented_dataset


def save_model(name, train_dir, state_dict):
    if train_dir is not None:
        if not os.path.exists(train_dir):
            os.makedirs(train_dir)
        ckpt_path = os.path.join(train_dir, name)

        torch.save(
            {
                "state_dict": state_dict,
            },
            ckpt_path,
        )


def load_model():
    raise NotImplementedError

src/utils/__init__.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
class EarlyStopping:
    def __init__(self, patience):
        self.patience = patience
        self.lowest_loss = None
        self.epochs_stagnant = 0

    def consume(self, new_loss):
        if self.lowest_loss is None:
            self.lowest_loss = new_loss
            self.epochs_stagnant = 0
        elif self.lowest_loss <= new_loss:
            self.epochs_stagnant += 1
        else:
            self.lowest_loss = new_loss
            self.epochs_stagnant = 0

    def should_stop(self):
        return self.epochs_stagnant >= self.patience

src/utils/early_stopping.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
class ModelResults:
    def __init__(self, results):
        if not isinstance(results, dict):
            raise ValueError("TODO:", results)
        else:
            self.results = results

    def __str__(self):
        ret_str = ""
        for k, v in self.results.items():
            if ret_str:
                ret_str += " "

            ret_str += "{}: {:.4f}".format(k, v)

        return ret_str

    def __getitem__(self, key):
        return self.results[key]

    def __setitem__(self, key, value):
        self.results[key] = value

src/utils/model_results.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
# flake8: noqa

import neptune.new as neptune

USE_NEPTUNE = True

def project():
    raise NotImplementedError("return the project name")

def api_token():
    raise NotImplementedError("return the api token")

def get_neptune_run(params, track_files_path):
    if not USE_NEPTUNE:
        return None


    run = neptune.init_run(
        project=project(),
        api_token=api_token(),
    )

    run["parameters"] = params
    run["checkpoints"].track_files(track_files_path)

    return run


def stop(run):
    if not USE_NEPTUNE:
        return

    run.stop()


def log_w_prefix(run, prefix, results):
    for name, value in results.items():
        log_w_prefix_single(run, prefix, name, value)


def log_w_prefix_single(run, prefix, name, value):
    if run is None:
        return None

    if not USE_NEPTUNE:
        return None

    run["{}/{}".format(prefix, name)].log(value)

src/utils/neptune_logging.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark
import argparse

import torch


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "configuration_name", type=str, help="configuration name"
    )
    parser.add_argument("checkpoint_name", type=str, help="checkpoint name")
    parser.add_argument(
        "load_checkpoint", type=str, help="noload|load_best|load_last"
    )
    parser.add_argument(
        "--use_neptune", type=bool, default=False, help="use neptune logging"
    )

    parser.add_argument("--latent_dim", type=int, default=None)
    parser.add_argument("--problem_spec_name", type=str, default=None)
    parser.add_argument("--model_instance", type=str, default=None)
    parser.add_argument("--enc_dim", type=int, default=None)
    parser.add_argument("--hidden_dim", type=int, default=None)
    parser.add_argument("--batch_size", type=int, default=None)
    parser.add_argument("--pretraining_len", type=int, default=None)
    parser.add_argument("--dec_depth", type=int, default=None)
    parser.add_argument("--dec_width", type=int, default=None)

    args = parser.parse_args()

    assert args.problem_spec_name is not None
    assert args.model_instance is not None
    assert args.batch_size is not None

    print("Cuda available:", torch.cuda.is_available())
    print("configuration name:", args.configuration_name)
    print("checkpoint_name:", args.checkpoint_name)
    print("load_checkpoint:", args.load_checkpoint)
    print("use neptune:", args.use_neptune)
    print("problem_spec_name:", args.problem_spec_name)
    print("latent_dim:", args.latent_dim)
    print("model_instance:", args.model_instance)

    return args

src/utils/parse_args.py
MIMUW-RL/Unified-Long-Horizon-Time-Series-Benchmark