Merge pull request #163 from noaa-ocs-modeling/feature/sflux_file

Add sflux forcing from existing file functionality for SCHISM setup

Author: SorooshMani-NOAA

coupledmodeldriver/client/initialize_schism.py

@@ -25,6 +25,7 @@
     TimestepForcingJSON,
     WaveForcingJSON,
     WindForcingJSON,
+    SfluxFileForcingJSON,
 )
 from coupledmodeldriver.generate import (
     SCHISMRunConfiguration,
@@ -38,6 +39,7 @@ class ForcingConfigurations(Enum):
     tidal = TidalForcingJSON
     besttrack = BestTrackForcingJSON
     nwm = NationalWaterModelFocringJSON
+    sfluxfiles = SfluxFileForcingJSON
 
 
 FORCING_NAMES = list(entry.name for entry in ForcingConfigurations)

coupledmodeldriver/configure/forcings/base.py

@@ -24,6 +24,7 @@
 from pyschism.forcing.bctides.tpxo import TPXO_VELOCITY as PySCHISMTPXO_VEL
 from pyschism.forcing.bctides.tides import TidalDatabase as PySCHISMTidalDatabase
 from pyschism.forcing.nws import BestTrackForcing as PySCHISMBestTrackForcing
+from pyschism.forcing.nws.nws2 import NWS2 as PySCHISMSfluxForcing
 from pyschism.forcing.base import ModelForcing as PySCHISMForcing
 from pyschism.forcing import NWM as PySCHISMNWM
 from pyschism.forcing.source_sink.nwm import NWMElementPairings
@@ -50,13 +51,11 @@
 PYSCHISM_FORCINGS = {
     'Tides': 'TidalForcingJSON',
     'BestTrackForcing': 'BestTrackForcingJSON',
-    'NationalWaterModel': 'NationalWaterModelFocringJSON'
-    # 'NWM : ,
-    # 'GFS, etc.': 'ATMESHForcingJSON',
-    # 'AtmosphericMeshForcing': 'ATMESHForcingJSON',
+    'NationalWaterModel': 'NationalWaterModelFocringJSON',
+    'NWS2': 'SfluxFileForcingJSON',
 }
 
-PYSCHISM_FORCING_CLASSES = (PySCHISMTides, PySCHISMForcing, PySCHISMNWM)
+PYSCHISM_FORCING_CLASSES = (PySCHISMTides, PySCHISMForcing, PySCHISMNWM, PySCHISMSfluxForcing)
 
 
 class TidalSource(IntEnum):
@@ -845,7 +844,7 @@ def __init__(
         kwargs['fields'].update(NationalWaterModelFocringJSON.field_types)
 
         HydrologyForcingJSON.__init__(self, **kwargs)
-        FileForcingJSON.__init__(self, resource=resource, **kwargs)
+        FileGenForcingJSON.__init__(self, resource=resource, **kwargs)
 
     @property
     def adcircpy_forcing(self) -> None:
@@ -894,3 +893,99 @@ def from_pyschism(cls, forcing: PySCHISMNWM) -> 'ForcingJSON':
             # sink_json=sink_json_path,
             # pairing_hgrid=pairing_hgrid_path
         )
+
+
+class SfluxFileForcingJSON(WindForcingJSON, FileForcingJSON):
+
+    name = 'Sflux'
+    default_filename = f'configure_sfluxfiles.json'
+    default_nws = 2  # SCHISM NWS
+    default_sflux_1_glob = '*_1.*'
+    default_sflux_2_glob = '*_2.*'
+    field_types = {
+        'sflux_1_glob': str,
+        'sflux_2_glob': str,
+    }
+
+    def __init__(
+        self, resource: PathLike, sflux_1_glob: str = None, sflux_2_glob: str = None, **kwargs,
+    ):
+
+        # TODO: Add windrot?
+        if sflux_1_glob is None:
+            sflux_1_glob = self.default_sflux_1_glob
+        if sflux_2_glob is None:
+            sflux_2_glob = self.default_sflux_2_glob
+
+        if 'fields' not in kwargs:
+            kwargs['fields'] = {}
+        kwargs['fields'].update(SfluxFileForcingJSON.field_types)
+
+        WindForcingJSON.__init__(self, **kwargs)
+        FileForcingJSON.__init__(self, resource=resource, **kwargs)
+
+        self['sflux_1_glob'] = sflux_1_glob
+        self['sflux_2_glob'] = sflux_2_glob
+
+    @property
+    def adcircpy_forcing(self) -> None:
+        raise NotImplementedError('ADCIRC does NOT support Sflux forcing!')
+
+    @classmethod
+    def from_adcircpy(cls, forcing: None) -> 'None':
+        raise NotImplementedError('ADCIRC does NOT support Sflux forcing!')
+
+    @property
+    def pyschism_forcing(self) -> PySCHISMSfluxForcing:
+
+        return PySCHISMSfluxForcing.read(
+            path=self['resource'],
+            sflux_1_glob=self['sflux_1_glob'],
+            sflux_2_glob=self['sflux_2_glob'],
+        )
+
+    @classmethod
+    def from_pyschism(cls, forcing: PySCHISMSfluxForcing) -> 'ForcingJSON':
+        sf1_paths = forcing.sflux_1.resource
+        if isinstance(sf1_paths, (str, Path)):
+            sf1_paths = [sf1_paths]
+        sf2_paths = forcing.sflux_2.resource
+        if isinstance(sf2_paths, (str, Path)):
+            sf2_paths = [sf2_paths]
+
+        path = cls._find_shared_root(*sf1_paths, *sf2_paths)
+        sf1_pat = cls._find_simple_pattern(sf1_paths, path)
+        sf2_pat = cls._find_simple_pattern(sf2_paths, path)
+
+        return cls(resource=path, sflux_1_glob=sf1_pat, sflux_2_glob=sf2_pat,)
+
+    def _find_shared_root(*paths):
+        roots = {Path(i).parent for i in paths}
+        while len(roots) != 0:
+            roots = {Path(i).parent for i in roots}
+
+        return roots.pop()
+
+    def _find_simple_pattern(paths, root):
+        # Note: CANNOT match complicated patterns, depth MUST be the same
+
+        if len(paths) == 0:
+            return ''
+
+        relpaths = [Path(p).relative_to(root) for p in paths]
+        if len(relpaths) == 1:
+            return relpaths[0]
+
+        exts = {rp.suffix for rp in relpaths}
+        ext = '.*'
+        if len(exts) == 1:
+            ext = exts.pop()
+
+        ns_parts = {len(rp.parts) - 1 for rp in relpaths}
+        if len(ns_parts) > 1:
+            raise ValueError('Cannot deduce pattern for different directory depths!')
+        n_parts = ns_parts.pop()
+
+        pat = ('*/' * n_parts) + '*' + ext
+
+        return pat

coupledmodeldriver/generate/schism/base.py

@@ -429,6 +429,9 @@ def pyschism_driver(self) -> ModelDriver:
             nws=meteo,
             source_sink=hydrology,
         )
+        # To avoid excessive memory use, especially when writing ensembles
+        self.base_mesh = self['hgrid_path']
+        self.pyschism_mesh = self['hgrid_path']
         # Hacky way to set the resource for tide
         if tides is not None:
             config.bctides.tides = tides

coupledmodeldriver/generate/schism/configure.py

@@ -22,6 +22,7 @@
     TidalForcingJSON,
     BestTrackForcingJSON,
     NationalWaterModelFocringJSON,
+    SfluxFileForcingJSON,
 )
 from coupledmodeldriver.generate.schism.base import SCHISMJSON
 from coupledmodeldriver.platforms import Platform
@@ -42,6 +43,7 @@ class SCHISMRunConfiguration(RunConfiguration):
         TidalForcingJSON,
         BestTrackForcingJSON,
         NationalWaterModelFocringJSON,
+        SfluxFileForcingJSON,
     }
 
     def __init__(