def cli(lstm_path, train_data_ncfiles, pred_length, train_length, params=None):
"""Create lstm-based prediction. Results are stored in
`lstm_{train_data_nc.stem}.npy` and skipped if already present.
"""
import torch
from tqdm import tqdm
import numpy as np
import netCDF4 as nc
import torsk
from torsk.scripts.pred_perf import sort_filenames
from torsk.models.torch_lstm import LSTM
if params is None:
params = torsk.Params()
params.dtype = "float32"
params.train_length = 200
params.pred_length = 300
params.batch_size = 32
params.hidden_size = 10000
params.input_shape = [30, 30]
else:
params = torsk.Params(json_path=params)
input_size = params.input_shape[0] * params.input_shape[1]
lstm_model = LSTM(input_size, params.hidden_size)
lstm_model.load_state_dict(torch.load(lstm_path))
train_data_ncfiles = sort_filenames(train_data_ncfiles)
for train_data_nc in tqdm(train_data_ncfiles):
assert "train" in train_data_nc.as_posix()
name = train_data_nc.stem.replace("train", "pred")
outfile = train_data_nc.parent / f"lstm_{name}.npy"
if not outfile.exists():
with nc.Dataset(train_data_nc, "r") as src:
labels = src["labels"][-train_length:]
labels = torch.Tensor(labels.reshape([1, train_length, -1]))
lstm_pred = lstm_model.predict(labels, steps=pred_length)
lstm_pred = lstm_pred.detach().numpy().reshape([
pred_length,
] + params.input_shape)
name = train_data_nc.stem.replace("train", "pred")
outfile = train_data_nc.parent / f"lstm_{name}.npy"
tqdm.write(f"Saving lstm-based prediction at {outfile}")
np.save(outfile, lstm_pred)
else:
tqdm.write(f"Skipping {outfile}")
def lstm_params(hidden_size):
hp = torsk.Params()
hp.dtype = "float32"
hp.train_length = 200
hp.pred_length = 300
hp.batch_size = 32
hp.hidden_size = hidden_size
return hp
def test_image_dataset():
params = torsk.Params(params={
"input_shape": [30, 30],
"input_map_specs": [
{"type": "pixels", "size": [10, 10], "input_scale": 3}],
"reservoir_representation": "dense",
"spectral_radius": 2.0,
"density": 1e-1,
"train_length": 10,
"pred_length": 5,
"transient_length": 5,
"train_method": "pinv_svd",
"dtype": "float64",
"backend": "numpy",
"debug": False,
"imed_loss": True
})
# test dtypes
for dtype_str in ["float32", "float64"]:
params.dtype = dtype_str
dtype = np.dtype(params.dtype)
height = params.input_shape[0]
width = params.input_shape[1]
images = np.zeros(
[params.pred_length + params.train_length + 2, height, width],
dtype=params.dtype)
images[:, 1:3, 1:3] = 2.
images[:, 0, 0] = -0.3
dataset = NumpyImageDataset(images, params)
assert len(dataset) == 2
inputs, labels, pred_labels = dataset[1]
assert inputs.shape == (params.train_length,) + images.shape[1:]
assert labels.shape == (params.train_length,) + images.shape[1:]
assert pred_labels.shape == (params.pred_length,) + images.shape[1:]
for arr in dataset[0]:
assert np.all(arr <= 1.)
assert np.all(arr >= -1.)
assert arr.dtype == dtype
unscaled = dataset.unscale(inputs)
assert np.isclose(unscaled.max(), 2.)
assert np.isclose(unscaled.min(), -0.3)
def esn_params(hidden_size):
params = torsk.Params()
params.input_map_specs = [{
"type": "random_weights",
"size": [hidden_size],
"input_scale": 1.
}]
params.spectral_radius = 1.5
params.density = 0.1
params.input_shape = [1, 1]
params.train_length = 2200
params.pred_length = 300
params.transient_length = 200
params.dtype = "float64"
params.reservoir_representation = "dense"
params.backend = "numpy"
params.train_method = "pinv_svd"
params.tikhonov_beta = 2.0
params.debug = False
params.imed_loss = False
return params
def test_numpy_save_load(tmpdir):
params_string = """{
"input_shape": [10, 10],
"input_map_specs": [
{"type":"pixels", "size":[10, 10], "input_scale":3}],
"reservoir_representation": "dense",
"spectral_radius" : 2.0,
"density": 1e-1,
"train_length": 800,
"pred_length": 300,
"transient_length": 100,
"train_method": "pinv_svd",
"dtype": "float64",
"backend": "numpy",
"debug": false,
"imed_loss": true,
"timing_depth": 1
}
"""
params_json = tmpdir.join("params.json")
with open(params_json, "w") as dst:
dst.write(params_string)
params = torsk.Params(params_json)
model = NumpyESN(params)
inputs = bh.random.uniform(size=[2, 10, 10]).astype(bh.float64)
state = bh.random.uniform(size=[100]).astype(bh.float64)
_, out1 = model.forward(inputs, state)
torsk.save_model(tmpdir, model)
model = torsk.load_model(str(tmpdir))
_, out2 = model.forward(inputs, state)
assert bh.all(out1 == out2)
import pathlib
import numpy as np
import netCDF4 as nc
import matplotlib.pyplot as plt
import torsk
from torsk.imed import imed_metric
from torsk.data.utils import resample2d_sequence
from torsk.visualize import animate_double_imshow
logger = logging.getLogger(__file__)
logging.basicConfig(level=logging.INFO)
np.random.seed(0)
params = torsk.Params()
params.input_map_specs = [
{
"type": "pixels",
"size": [30, 30],
"input_scale": 3.
},
{
"type": "conv",
"mode": "same",
"size": [5, 5],
"kernel_type": "gauss",
"input_scale": 2.
},
{
"type": "conv",
