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, activation='identity' ) 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