KoopmanESN/TrajTrans/StateOnly/ESN/BaseModel/ESNModel.py

import reservoirpy as rpy
from KoopmanESNModel import KoopmanESNConfig
from reservoirpy.nodes import Reservoir, Ridge
import numpy as np


class ESN:
    def __init__(self, config: KoopmanESNConfig):
        super().__init__()

        self.esn_model = None
        self.config = config

    def esn_train(self, data_train):
        rpy.verbosity(0)
        rpy.set_seed(42)

        reservoir = Reservoir(
            units=self.config.units,
            lr=self.config.lr,
            sr=self.config.sr,
            rc_connectivity=self.config.sp
        )

        readout = Ridge(ridge=self.config.ridge)

        esn_model = reservoir >> readout

        data_train_in = data_train[:-1, :]
        data_train_out = data_train[1:, :]

        self.esn_model = esn_model.fit(
            data_train_in,
            data_train_out,
            warmup=self.config.train_warmup
        )

    def predict(self, data_warm):
        warmup_data = self.esn_model.run(data_warm, reset=True)
        x = warmup_data[-1, :].reshape(1, -1)
        data_pre = np.empty((self.config.predict_len, self.config.state_dim))
        for i in range(self.config.predict_len):
            data_pre[i, :] = x
            x = self.esn_model(x)
        return data_pre