123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778 |
- import reservoirpy as rpy
- import torch
- from numpy.linalg import eigvals
- from reservoirpy.nodes import Reservoir, Ridge
- import numpy as np
- from dataclasses import dataclass
- @dataclass
- class KoopmanESNConfig:
- units: int
- lr: float
- sr: float
- sp: float
- ridge: float
- train_start: int
- train_len: int
- train_warmup: int
- predict_warmup: int
- predict_len: int
- state_dim: int
- class KoopmanESN:
- def __init__(self, A, config: KoopmanESNConfig):
- super().__init__()
- self.esn_model = None
- self.config = config
- self.A = A/100
- def W_initial(self, A):
- W_koop = A
- W = W_koop
- # 5. 调整谱半径
- current_radius = np.max(np.abs(eigvals(W)))
- W = (self.config.sr / current_radius) * W
- return W
- def koopmanesn_train(self, data_train):
- rpy.verbosity(0)
- rpy.set_seed(42)
- W = self.W_initial(self.A)
- Win = np.eye(self.config.units, self.config.units)
- reservoir = Reservoir(
- units=self.config.units,
- lr=self.config.lr,
- sr=self.config.sr,
- W=W,
- input_scaling=1.0,
- Win=Win,
- 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
|