feat: enabled forward passes with obs models to run without obs data; added nn_fill on ouput data outside of allowed range, instead of clipping; added collection of netcdf spatial chunks.

Author: bnb32

sup3r/models/abstract.py

@@ -36,16 +36,15 @@
 
 logger = logging.getLogger(__name__)
 
-
-SUP3R_LAYERS = (
-    Sup3rAdder,
-    Sup3rConcat,
+SUP3R_OBS_LAYERS = (
     Sup3rConcatObs,
     Sup3rConcatEmbeddedObs,
     Sup3rConcatWeightedObs,
     Sup3rConcatWeightedObsWithEmbedding,
 )
 
+SUP3R_LAYERS = (Sup3rAdder, Sup3rConcat, *SUP3R_OBS_LAYERS)
+
 
 # pylint: disable=E1101,W0201,E0203
 class AbstractSingleModel(ABC, TensorboardMixIn):
@@ -1041,7 +1040,15 @@ def generate(
         try:
             for i, layer in enumerate(self.generator.layers[1:]):
                 layer_num = i + 1
-                if isinstance(layer, SUP3R_LAYERS):
+                is_obs_layer = isinstance(layer, SUP3R_OBS_LAYERS)
+                is_exo_layer = isinstance(layer, SUP3R_LAYERS)
+                if is_obs_layer and layer.name not in exogenous_data:
+                    msg = (f'Observation data not given for {layer.name}. '
+                           'Will run forward pass without it.')
+                    logger.warning(msg)
+                    warn(msg)
+                    hi_res = layer(hi_res)
+                elif is_exo_layer:
                     msg = (
                         f'layer.name = {layer.name} does not match any '
                         'features in exogenous_data '

sup3r/postprocessing/collectors/nc.py

@@ -6,9 +6,11 @@
 import logging
 import os
 
+import xarray as xr
 from rex.utilities.loggers import init_logger
 
 from sup3r.preprocessing.cachers import Cacher
+from sup3r.preprocessing.names import Dimension
 from sup3r.utilities.utilities import xr_open_mfdataset
 
 from .base import BaseCollector
@@ -32,6 +34,12 @@ def collect(
     ):
         """Collect data files from a dir to one output file.
 
+        TODO: This assumes that if there is any spatial chunking it is split
+        by latitude. This should be generalized to allow for any spatial
+        chunking and any dimension. This will either require a new file
+        naming scheme with a spatial index for both latitude and
+        longitude or checking each chunk to see how they are split.
+
         Filename requirements:
          - Should end with ".nc"
 
@@ -76,10 +84,19 @@ def collect(
             logger.info(f'overwrite=True, removing {out_file}.')
             os.remove(out_file)
 
+        spatial_chunks = collector.group_spatial_chunks()
+
         tmp_file = out_file + '.tmp'
         if not os.path.exists(tmp_file):
-            res_kwargs = res_kwargs or {}
-            out = xr_open_mfdataset(collector.flist, **res_kwargs)
+            res_kwargs = res_kwargs or {
+                'combine': 'nested',
+                'concat_dim': Dimension.TIME,
+            }
+            for s_idx in spatial_chunks:
+                spatial_chunks[s_idx] = xr_open_mfdataset(
+                    spatial_chunks[s_idx], **res_kwargs
+                )
+            out = xr.concat(spatial_chunks.values(), dim=Dimension.SOUTH_NORTH)
             Cacher.write_netcdf(tmp_file, data=out, features=features)
 
         os.replace(tmp_file, out_file)
@@ -88,12 +105,12 @@ def collect(
         logger.info('Finished file collection.')
 
     def group_spatial_chunks(self):
-        """Group same spatial chunks together so each chunk has same spatial
+        """Group same spatial chunks together so each entry has same spatial
         footprint but different times"""
         chunks = {}
         for file in self.flist:
-            s_chunk = file.split('_')[0]
-            dirname = os.path.dirname(file)
-            s_file = os.path.join(dirname, f's_{s_chunk}.nc')
-            chunks[s_file] = [*chunks.get(s_file, []), s_file]
+            _, s_idx = self.get_chunk_indices(file)
+            chunks[s_idx] = [*chunks.get(s_idx, []), file]
+        for k, v in chunks.items():
+            chunks[k] = sorted(v)
         return chunks

sup3r/preprocessing/accessor.py

@@ -372,10 +372,11 @@ def interpolate_na(self, **kwargs):
                         'use_coordinate', False
                     )
                 self._ds[feat] = self._ds[feat].interpolate_na(**kwargs)
-            else:
+            elif np.isnan(self._ds[feat]).any():
                 msg = (
                     'No dim given for interpolate_na. This will use nearest '
-                    'neighbor fill, which could take some time.'
+                    f'neighbor fill to interpolate {feat}, which could take '
+                    'some time.'
                 )
                 logger.warning(msg)
                 warn(msg)

sup3r/preprocessing/data_handlers/exo.py

@@ -97,7 +97,7 @@ def __init__(self, steps):
                     warn(msg)
                 steps_list = entry.get('steps', [entry])
                 for i, step in enumerate(steps_list):
-                    msg = (f'ExoData entry for {feat}, step #{i+1} has no '
+                    msg = (f'ExoData entry for {feat}, step #{i + 1}, has no '
                            '"combine_type" key. Assuming this is for a '
                            'layer combination.')
                     if 'combine_type' not in step:

sup3r/utilities/utilities.py

@@ -156,6 +156,10 @@ def enforce_limits(features, data, nn_fill=False):
         max_val = OUTPUT_ATTRS[dset_name].get('max', np.inf)
         min_val = OUTPUT_ATTRS[dset_name].get('min', -np.inf)
         enforcing_msg = f'Enforcing range of ({min_val}, {max_val}) for "{fn}"'
+        if nn_fill:
+            enforcing_msg += ' with nearest neighbor interpolation.'
+        else:
+            enforcing_msg += ' with clipping.'
 
         f_max = data[..., fidx].max()
         f_min = data[..., fidx].min()