Dev/reservoir experiments

external/fv3fit/fv3fit/reservoir/adapters.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
+from dataclasses import asdict
+import fsspec
 import numpy as np
 import os
 import typing
-from typing import Iterable, Hashable, Sequence, Union, Mapping
+from typing import Iterable, Hashable, Sequence, Union, Mapping, Optional
 import xarray as xr
+import yaml
 
 import fv3fit
 from fv3fit import Predictor
 from fv3fit._shared import io
+from fv3fit.reservoir.config import ClipZConfig
 from .model import (
     HybridReservoirComputingModel,
     ReservoirComputingModel,
@@ -66,7 +70,10 @@ def output_array_to_ds(
         ).transpose(*output_dims)
 
     def input_dataset_to_arrays(
-        self, inputs: xr.Dataset, variables: Iterable[Hashable]
+        self,
+        inputs: xr.Dataset,
+        variables: Iterable[Hashable],
+        clip_config: Optional[Mapping[Hashable, ClipZConfig]] = None,
     ) -> Sequence[np.ndarray]:
         # Converts from xr dataset to sequence of variable ndarrays expected by encoder
         # Make sure the xy dimensions match the rank divider
@@ -80,6 +87,10 @@ def input_dataset_to_arrays(
             da = transposed_inputs[variable]
             if "z" not in da.dims:
                 da = da.expand_dims("z", axis=-1)
+            if clip_config is not None and variable in clip_config:
+                da = da.isel(
+                    z=slice(clip_config[variable].start, clip_config[variable].stop)
+                )
             input_arrs.append(da.values)
         return input_arrs
 
@@ -118,6 +129,9 @@ def input_overlap(self):
     def is_hybrid(self):
         return False
 
+    def dump_state(self, path: str):
+        self.model.reservoir.dump_state(path)
+
     def predict(self, inputs: xr.Dataset) -> xr.Dataset:
         # inputs arg is not used, but is required by Predictor signature and prog run
         prediction_arr = self.model.predict()
@@ -127,7 +141,7 @@ def predict(self, inputs: xr.Dataset) -> xr.Dataset:
 
     def increment_state(self, inputs: xr.Dataset):
         xy_input_arrs = self.model_adapter.input_dataset_to_arrays(
-            inputs, self.input_variables
+            inputs, self.input_variables,
         )  # x, y, feature dims
         self.model.increment_state(xy_input_arrs)
 
@@ -159,12 +173,14 @@ def load(cls, path: str) -> "ReservoirDatasetAdapter":
 @io.register("hybrid-reservoir-adapter")
 class HybridReservoirDatasetAdapter(Predictor):
     MODEL_DIR = "hybrid_reservoir_model"
+    CLIP_CONFIG_FILE = "clip_config.yaml"
 
     def __init__(
         self,
         model: HybridReservoirComputingModel,
         input_variables: Iterable[Hashable],
         output_variables: Iterable[Hashable],
+        clip_config: Optional[Mapping[Hashable, ClipZConfig]] = None,
     ) -> None:
         """Wraps a hybrid reservoir model to take in and return xarray datasets.
         The initialization args for input and output variables are not used and
@@ -183,6 +199,10 @@ def __init__(
             input_variables=self.input_variables,
             output_variables=model.output_variables,
         )
+        self.clip_config = clip_config
+
+    def dump_state(self, path: str):
+        self.model.reservoir.dump_state(path)
 
     @property
     def input_overlap(self):
@@ -195,7 +215,7 @@ def is_hybrid(self):
 
     def predict(self, inputs: xr.Dataset) -> xr.Dataset:
         xy_input_arrs = self.model_adapter.input_dataset_to_arrays(
-            inputs, self.model.hybrid_variables
+            inputs, self.model.hybrid_variables, clip_config=self.clip_config
         )  # x, y, feature dims
 
         prediction_arr = self.model.predict(xy_input_arrs)
@@ -205,7 +225,7 @@ def predict(self, inputs: xr.Dataset) -> xr.Dataset:
 
     def increment_state(self, inputs: xr.Dataset):
         xy_input_arrs = self.model_adapter.input_dataset_to_arrays(
-            inputs, self.model.input_variables
+            inputs, self.model.input_variables, clip_config=self.clip_config
         )  # x, y, feature dims
         self.model.increment_state(xy_input_arrs)
 
@@ -224,14 +244,30 @@ def get_model_from_subdomain(
 
     def dump(self, path):
         self.model.dump(os.path.join(path, self.MODEL_DIR))
+        if self.clip_config is not None:
+            clip_config_dict = {
+                var: asdict(var_config) for var, var_config in self.clip_config.items()
+            }
+            with fsspec.open(os.path.join(path, self.CLIP_CONFIG_FILE), "w") as f:
+                yaml.dump(clip_config_dict, f)
 
     @classmethod
     def load(cls, path: str) -> "HybridReservoirDatasetAdapter":
         model = HybridReservoirComputingModel.load(os.path.join(path, cls.MODEL_DIR))
+        try:
+            with fsspec.open(os.path.join(path, cls.CLIP_CONFIG_FILE), "r") as f:
+                clip_config_dict = yaml.safe_load(f)
+            clip_config: Optional[Mapping[Hashable, ClipZConfig]] = {
+                var: ClipZConfig(**var_config)
+                for var, var_config in clip_config_dict.items()
+            }
+        except FileNotFoundError:
+            clip_config = None
         adapter = cls(
             input_variables=model.input_variables,
             output_variables=model.output_variables,
             model=model,
+            clip_config=clip_config,
         )
         return adapter
 

external/fv3fit/fv3fit/reservoir/config.py

@@ -1,6 +1,6 @@
 import dacite
 from dataclasses import dataclass, asdict
-from typing import Sequence, Optional, Set, Tuple
+from typing import Sequence, Optional, Set, Tuple, Mapping, Hashable
 import fsspec
 import yaml
 from .._shared.training_config import Hyperparameters
@@ -65,6 +65,16 @@ class TransformerConfig:
     hybrid: Optional[str] = None
 
 
+@dataclass
+class ClipZConfig:
+    """ Vertical levels **between** start and stop are kept,
+    levels outside start/stop are clipped off.
+    """
+
+    start: Optional[int] = None
+    stop: Optional[int] = None
+
+
 @dataclass
 class ReservoirTrainingConfig(Hyperparameters):
     """
@@ -101,11 +111,13 @@ class ReservoirTrainingConfig(Hyperparameters):
     n_timesteps_synchronize: int
     input_noise: float
     seed: int = 0
+    zero_fill_clipped_output_levels: bool = False
     transformers: Optional[TransformerConfig] = None
     n_jobs: Optional[int] = 1
     square_half_hidden_state: bool = False
     hybrid_variables: Optional[Sequence[str]] = None
     mask_variable: Optional[str] = None
+    clip_config: Optional[Mapping[Hashable, ClipZConfig]] = None
     _METADATA_NAME = "reservoir_training_config.yaml"
 
     def __post_init__(self):

external/fv3fit/fv3fit/reservoir/train.py

@@ -19,6 +19,8 @@
     assure_txyz_dims,
     SynchronziationTracker,
     get_standard_normalizing_transformer,
+    clip_batch_data,
+    zero_fill_clipped_output_levels,
 )
 from .transformers import TransformerGroup, Transformer
 from .._shared import register_training_function
@@ -48,6 +50,8 @@ def _add_input_noise(arr: np.ndarray, stddev: float) -> np.ndarray:
 def _get_transformers(
     sample_batch: Mapping[str, tf.Tensor], hyperparameters: ReservoirTrainingConfig
 ) -> TransformerGroup:
+    clipped_sample_batch = clip_batch_data(sample_batch, hyperparameters.clip_config)
+
     # Load transformers with specified paths
     transformers = {}
     for variable_group in ["input", "output", "hybrid"]:
@@ -58,25 +62,25 @@ def _get_transformers(
     # If input transformer not specified, always create a standard norm transform
     if "input" not in transformers:
         transformers["input"] = get_standard_normalizing_transformer(
-            hyperparameters.input_variables, sample_batch
+            hyperparameters.input_variables, clipped_sample_batch
         )
 
-    # If output transformer not specified and output_variables != input_variables,
-    # create a separate standard norm transform
+    # Output is not clipped, so use the original sample batch
+    if hyperparameters.zero_fill_clipped_output_levels:
+        sample_batch = zero_fill_clipped_output_levels(
+            sample_batch, hyperparameters.clip_config
+        )
     if "output" not in transformers:
-        if hyperparameters.output_variables != hyperparameters.input_variables:
-            transformers["output"] = get_standard_normalizing_transformer(
-                hyperparameters.output_variables, sample_batch
-            )
-        else:
-            transformers["output"] = transformers["input"]
+        transformers["output"] = get_standard_normalizing_transformer(
+            hyperparameters.output_variables, sample_batch
+        )
 
     # If hybrid variables transformer not specified, and hybrid variables are defined,
     # create a separate standard norm transform
     if "hybrid" not in transformers:
         if hyperparameters.hybrid_variables is not None:
             transformers["hybrid"] = get_standard_normalizing_transformer(
-                hyperparameters.hybrid_variables, sample_batch
+                hyperparameters.hybrid_variables, clipped_sample_batch
             )
         else:
             transformers["hybrid"] = transformers["input"]
@@ -115,9 +119,17 @@ def train_reservoir_model(
         train_batches if isinstance(train_batches, Sequence) else [train_batches]
     )
     sample_batch = next(iter(train_batches_sequence[0]))
-    sample_X = get_ordered_X(sample_batch, hyperparameters.input_variables)
+    if hyperparameters.zero_fill_clipped_output_levels:
+        sample_batch = zero_fill_clipped_output_levels(
+            sample_batch, hyperparameters.clip_config
+        )
+
+    # Clipping is done inside this function to preserve full length outputs
+    transformers = _get_transformers(sample_batch, hyperparameters,)
+
+    clipped_sample_batch = clip_batch_data(sample_batch, hyperparameters.clip_config)
+    sample_X = get_ordered_X(clipped_sample_batch, hyperparameters.input_variables)
 
-    transformers = _get_transformers(sample_batch, hyperparameters)
     subdomain_config = hyperparameters.subdomain
 
     # sample_X[0] is the first data variable, shape elements 1:-1 are the x,y shape
@@ -131,7 +143,7 @@ def train_reservoir_model(
 
     if hyperparameters.mask_variable is not None:
         input_mask_array: Optional[np.ndarray] = _get_input_mask_array(
-            hyperparameters.mask_variable, sample_batch, rank_divider
+            hyperparameters.mask_variable, clipped_sample_batch, rank_divider
         )
     else:
         input_mask_array = None
@@ -157,9 +169,12 @@ def train_reservoir_model(
         )
 
         for b, batch_data in enumerate(train_batches):
+            batch_data_clipped = clip_batch_data(
+                batch_data, hyperparameters.clip_config
+            )
             input_time_series = process_batch_data(
                 variables=hyperparameters.input_variables,
-                batch_data=batch_data,
+                batch_data=batch_data_clipped,
                 rank_divider=rank_divider,
                 autoencoder=transformers.input,
                 trim_halo=False,
@@ -171,6 +186,11 @@ def train_reservoir_model(
             _output_rank_divider_with_overlap = rank_divider.get_new_zdim_rank_divider(
                 z_feature_size=transformers.output.n_latent_dims
             )
+            # don't pass in clipped data here, as clipping is not enabled for outputs
+            if hyperparameters.zero_fill_clipped_output_levels:
+                batch_data = zero_fill_clipped_output_levels(
+                    batch_data, hyperparameters.clip_config
+                )
             output_time_series = process_batch_data(
                 variables=hyperparameters.output_variables,
                 batch_data=batch_data,
@@ -196,7 +216,7 @@ def train_reservoir_model(
 
                 hybrid_time_series = process_batch_data(
                     variables=hyperparameters.hybrid_variables,
-                    batch_data=batch_data,
+                    batch_data=batch_data_clipped,
                     rank_divider=_hybrid_rank_divider_w_overlap,
                     autoencoder=transformers.hybrid,
                     trim_halo=True,
@@ -279,6 +299,7 @@ def train_reservoir_model(
             model=model,
             input_variables=model.input_variables,
             output_variables=model.output_variables,
+            clip_config=hyperparameters.clip_config,
         )
 
     if validation_batches is not None and wandb.run is not None:
@@ -287,6 +308,7 @@ def train_reservoir_model(
                 model,
                 val_batches=validation_batches,
                 n_synchronize=hyperparameters.n_timesteps_synchronize,
+                clip_config=hyperparameters.clip_config,
             )
             log_rmse_z_plots(ds_val, model.output_variables)
             log_rmse_scalar_metrics(ds_val, model.output_variables)