Updated obs model to use separate obs frac for onshore and offshore regions.

Author: bnb32

sup3r/models/with_obs.py

@@ -14,34 +14,51 @@
 
 
 class Sup3rGanWithObs(Sup3rGan):
-    """Sup3r GAN model which includes mid network observation fixing. This
+    """Sup3r GAN model which includes mid network observation fusion. This
     model is useful for when production runs will be over a domain for which
     observation data is available."""
 
-    def __init__(self, *args, obs_frac=None, loss_obs_weight=None, **kwargs):
+    def __init__(
+        self,
+        *args,
+        onshore_obs_frac=None,
+        offshore_obs_frac=None,
+        loss_obs_weight=None,
+        **kwargs,
+    ):
         """
         Initialize the Sup3rGanWithObs model.
 
         Parameters
         ----------
         args : list
             Positional args for ``Sup3rGan`` parent class.
-        obs_frac : dict
-            Fraction of the batch that should be "fixed" with observations.
-            Should include ``spatial`` key and optionally ``time`` key if this
-            is a spatiotemporal model. The values should correspond roughly to
-            the fraction of the production domain for which observations are
-            available (spatial) and the fraction of the full time period that
-            these cover. For each batch a spatial frac will be selected by
-            uniformly selecting from the range ``(0, obs_frac['spatial'])``
+        onshore_obs_frac : dict
+            Fraction of the batch that should be treated as onshore
+            observations. Should include ``spatial`` key and optionally
+            ``time`` key if this is a spatiotemporal model. The values should
+            correspond roughly to the fraction of the production domain for
+            which onshore observations are available (spatial) and the fraction
+            of the full time period that these cover. For each batch a spatial
+            frac will be selected by uniformly selecting from the range ``(0,
+            obs_frac['spatial'])``
+        offshore_obs_frac : dict
+            Same as ``onshore_obs_frac`` but for offshore observations.
+            Offshore observations are frequently sparser than onshore
+            observations.
         loss_obs_weight : float
             Value used to weight observation locations in extra content loss
             term. e.g. The new content loss will include ``obs_loss_weight *
             MAE(hi_res_gen[~obs_mask], hi_res_true[~obs_mask])``
         kwargs : dict
             Keyword arguments for the ``Sup3rGan`` parent class.
         """
-        self.obs_frac = {} if obs_frac is None else obs_frac
+        self.onshore_obs_frac = (
+            {} if onshore_obs_frac is None else onshore_obs_frac
+        )
+        self.offshore_obs_frac = (
+            {} if offshore_obs_frac is None else offshore_obs_frac
+        )
         self.loss_obs_weight = loss_obs_weight
         super().__init__(*args, **kwargs)
 
@@ -51,12 +68,12 @@ def _get_loss_obs_comparison(self, hi_res_true, hi_res_gen, obs_mask):
         locations."""
 
         hr_true = [
-            hi_res_true[..., self.hr_out_features.index(f)]
+            hi_res_true[..., self.hr_out_features.index(f.replace('_obs', ''))]
             for f in self.obs_features
         ]
         hr_true = tf.stack(hr_true, axis=-1)
         hr_gen = [
-            hi_res_gen[..., self.hr_out_features.index(f)]
+            hi_res_gen[..., self.hr_out_features.index(f.replace('_obs', ''))]
             for f in self.obs_features
         ]
         hr_gen = tf.stack(hr_gen, axis=-1)
@@ -65,23 +82,15 @@ def _get_loss_obs_comparison(self, hi_res_true, hi_res_gen, obs_mask):
         loss_non_obs = MeanAbsoluteError()(hr_true[obs_mask], hr_gen[obs_mask])
         return loss_obs, loss_non_obs
 
-    def _get_obs_mask(self, hi_res, spatial_frac=None, time_frac=None):
-        """Define observation mask for the current batch. This is done
-        with a spatial mask and a temporal mask since often observation data
-        might be very sparse spatially but cover most of the full time period
-        for those locations."""
-        spatial_frac = (
-            self.obs_frac['spatial'] if spatial_frac is None else spatial_frac
-        )
+    def _get_obs_mask(self, hi_res, spatial_frac, time_frac=None):
+        """Get observation mask for a given spatial and temporal obs
+        fraction."""
         obs_mask = RANDOM_GENERATOR.choice(
             [True, False],
             size=hi_res.shape[1:3],
             p=[1 - spatial_frac, spatial_frac],
         )
         if self.is_5d:
-            time_frac = (
-                self.obs_frac['time'] if time_frac is None else time_frac
-            )
             sp_mask = obs_mask.copy()
             obs_mask = RANDOM_GENERATOR.choice(
                 [True, False],
@@ -91,6 +100,30 @@ def _get_obs_mask(self, hi_res, spatial_frac=None, time_frac=None):
             obs_mask[sp_mask] = True
         return np.repeat(obs_mask[None, ...], hi_res.shape[0], axis=0)
 
+    def get_obs_mask(self, hi_res):
+        """Define observation mask for the current batch. This is done
+        with a spatial mask and a temporal mask since often observation data
+        might be very sparse spatially but cover most of the full time period
+        for those locations. This is also divided between onshore and offshore
+        regions"""
+        on_sf = RANDOM_GENERATOR.uniform(
+            low=0, high=self.onshore_obs_frac['spatial']
+        )
+        on_tf = self.onshore_obs_frac.get('time', None)
+        off_tf = self.offshore_obs_frac.get('time', None)
+        obs_mask = self._get_obs_mask(hi_res, on_sf, on_tf)
+        if 'topography' in self.hr_exo_features and self.offshore_obs_frac:
+            topo_idx = len(self.hr_out_features) + self.hr_exo_features.index(
+                'topography'
+            )
+            topo = hi_res[..., topo_idx]
+            off_sf = RANDOM_GENERATOR.uniform(
+                low=0, high=self.offshore_obs_frac['spatial']
+            )
+            offshore_mask = self._get_obs_mask(hi_res, off_sf, off_tf)
+            obs_mask = tf.where(topo > 0, obs_mask, offshore_mask)
+        return obs_mask
+
     @property
     def model_params(self):
         """
@@ -101,19 +134,16 @@ def model_params(self):
         dict
         """
         params = super().model_params
-        params['obs_frac'] = self.obs_frac
+        params['onshore_obs_frac'] = self.onshore_obs_frac
+        params['offshore_obs_frac'] = self.offshore_obs_frac
         params['loss_obs_weight'] = self.loss_obs_weight
         return params
 
     def get_hr_exo_input(self, hi_res_true):
         """Mask high res data to act as sparse observation data. Add this to
         the standard high res exo input"""
         exo_data = super().get_hr_exo_input(hi_res_true)
-        spatial_frac = RANDOM_GENERATOR.uniform(
-            low=0, high=self.obs_frac['spatial']
-        )
-        time_frac = self.obs_frac.get('time', None)
-        obs_mask = self._get_obs_mask(hi_res_true, spatial_frac, time_frac)
+        obs_mask = self.get_obs_mask(hi_res_true)
         for feature in self.obs_features:
             # obs_features can include a _obs suffix to avoid name conflict
             # with fully gridded exo features

tests/training/test_train_conditioned_obs.py

@@ -47,11 +47,11 @@ def test_fixed_wind_obs(gen_config_with_concat_masked):
     model = Sup3rGanWithObs(
         gen_config_with_concat_masked(),
         pytest.S_FP_DISC,
-        obs_frac={'spatial': 0.1},
+        onshore_obs_frac={'spatial': 0.1},
         loss_obs_weight=0.1,
         learning_rate=1e-4,
     )
-    test_mask = model._get_obs_mask(np.zeros((1, 20, 20, 1, 1)))
+    test_mask = model.get_obs_mask(np.zeros((1, 20, 20, 1, 1)))
     frac = 1 - test_mask.sum() / test_mask.size
     assert np.abs(0.1 - frac) < test_mask.size / (2 * np.sqrt(test_mask.size))
     assert model.obs_features == ['u_10m', 'v_10m']