Merged NWM hindcast inventory class into nwm

schism-dev · Jul 28, 2021 · 001445d · 001445d
1 parent 1c1db31
commit 001445d
Show file tree

Hide file tree

Showing 10 changed files with 204 additions and 5 deletions.
diff --git a/.idea/.gitignore b/.idea/.gitignore
diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml
diff --git a/.idea/modules.xml b/.idea/modules.xml
diff --git a/.idea/pyschism.iml b/.idea/pyschism.iml
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/pyschism/forcing/hycom/hycom.py b/pyschism/forcing/hycom/hycom.py
@@ -14,7 +14,6 @@
 from matplotlib.transforms import Bbox
 
 from pyschism.mesh import Hgrid
-from pyschism.dates import localize_datetime, nearest_cycle_date, pivot_time
 
 logger = logging.getLogger(__name__)
 

diff --git a/pyschism/forcing/hydrology/nwm.py b/pyschism/forcing/hydrology/nwm.py
@@ -14,6 +14,7 @@
 import boto3
 from botocore import UNSIGNED
 from botocore.config import Config
+import pandas as pd
 import geopandas as gpd
 from netCDF4 import Dataset
 import numpy as np
@@ -110,6 +111,7 @@ def __init__(self, hgrid):
                 'No National Water model intersections found on the mesh.')
         intersection = gpd.GeoDataFrame(data, crs=hgrid.crs)
         del data
+        print(intersection)
 
         # 2) Generate element centroid KDTree
         centroids = []
@@ -422,6 +424,133 @@ def requested_product(self):
         }[self.product]
 
 
+class AWSHindcastInventory:
+
+    def __init__(
+            self,
+            start_date: datetime = None,
+            rnday: Union[int, float, timedelta] = timedelta(days=5.),
+            product='CHRTOUT_DOMAIN1.comp',
+            verbose=False,
+            fallback=True,
+    ):
+        """This will download the National Water Model retro data.
+        A 26-year (January 1993 through December 2018) retrospective 
+        simulation using version 2.0 of the NWM. 
+
+        NetCDF files are saved to the system's temporary directory.
+        """
+        self.start_date = dates.nearest_cycle() if start_date is None \
+            else dates.nearest_cycle(dates.localize_datetime(start_date))
+        # self.start_date = self.start_date.replace(tzinfo=None)
+        self.rnday = rnday if isinstance(rnday, timedelta) \
+            else timedelta(days=rnday)
+        self.product = product
+        self.fallback = fallback
+        self._files = {_: None for _ in np.arange(
+            self.start_date,
+            self.start_date + self.rnday + self.output_interval,
+            self.output_interval
+        ).astype(datetime)}
+
+        paginator=self.s3.get_paginator('list_objects_v2')
+        pages=paginator.paginate(Bucket=self.bucket,
+                Prefix=f'full_physics/{self.start_date.year}')
+
+        self.data=[]
+        for page in pages:
+            for obj in page['Contents']:
+                self.data.append(obj)
+
+        self.file_metadata = list(sorted([
+            _['Key'] for _ in self.data if 'CHRTOUT_DOMAIN1.comp' in _['Key']
+        ]))
+
+        timevector=np.arange(datetime(self.start_date.year, 1, 1),
+            datetime(self.start_date.year+1,1, 1),
+            np.timedelta64(1, 'h'),
+            dtype='datetime64')
+
+        timefile= {pd.to_datetime(str(timevector[i])): self.file_metadata[i] 
+            for i in range(len(timevector))}
+
+        for requested_time, _ in self._files.items():
+            #print(f'Requesting NWM data for time {requested_time}')
+            key=timefile.get(requested_time)
+            print(f'Requesting NWM data for time {requested_time}, {key}')
+            logger.info(f'Requesting NWM data for time {requested_time}')
+            self._files[requested_time] = self.request_data(key, requested_time)
+
+    def request_data(self, key, request_time):
+
+        #print(key)
+        filename = pathlib.Path(self.tmpdir.name) / key
+        filename.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(filename, 'wb') as f:
+            #logger.info(f'Downloading file {key}, ')
+            print(f'Downloading file {key}, ')
+            self.s3.download_fileobj(self.bucket, key, f)
+        return filename
+
+    def get_nc_pairing_indexes(self, pairings: NWMElementPairings):
+        nc_feature_id = Dataset(self.files[0])['feature_id'][:]
+
+        def get_aggregated_features(features):
+            aggregated_features = []
+            for source_feats in features:
+                aggregated_features.extend(list(source_feats))
+            in_file = np.where(
+                np.in1d(nc_feature_id, aggregated_features,
+                assume_unique=True))[0]
+            in_file_2 = []
+            sidx = 0
+            for source_feats in features:
+                eidx = sidx + len(source_feats)
+                in_file_2.append(in_file[sidx:eidx].tolist())
+                sidx = eidx
+            return in_file_2
+
+        sources = get_aggregated_features(pairings.sources.values())
+        sinks = get_aggregated_features(pairings.sinks.values())
+        return sources, sinks
+
+    @property
+    def bucket(self):
+        return 'noaa-nwm-retro-v2.0-pds'
+
+    @property
+    def nearest_cycle(self) -> datetime:
+        return dates.nearest_cycle(self.start_date)
+
+    @property
+    def output_interval(self) -> timedelta:
+        return {
+            'CHRTOUT_DOMAIN1.comp': timedelta(hours=1)
+        }[self.product]
+
+    @property
+    def s3(self):
+        try:
+            return self._s3
+        except AttributeError:
+            self._s3 = boto3.client(
+                's3', config=Config(signature_version=UNSIGNED))
+            return self._s3
+
+    @property
+    def tmpdir(self):
+        try:
+            return self._tmpdir
+        except AttributeError:
+            self._tmpdir = tempfile.TemporaryDirectory()
+            return self._tmpdir
+
+    @property
+    def files(self):
+        return sorted(list(pathlib.Path(self.tmpdir.name).glob('**/*.comp'))) 
+
+
 class NationalWaterModel(Hydrology):
 
     def __init__(self, aggregation_radius=None):

diff --git a/pyschism/mesh/vgrid.py b/pyschism/mesh/vgrid.py
@@ -14,7 +14,7 @@ class Vgrid:
 
     def __init__(self):
         """Represents a SCHISM vertical grid."""
-        #self._vgrid = self._get_2D_string()
+        self._vgrid = self._get_2D_string()
         pass
 
     @classmethod
@@ -63,7 +63,7 @@ def read_vgrid(fname):
             lines=lines[2:]
             kbp=np.array([int(i.split()[1])-1 for i in lines])
             NP=len(kbp)
-            print(NP)
+            #print(NP)
             sigma=-np.ones([NP,nvrt])
             for i, line in enumerate(lines):
                 sigma[i,kbp[i]:]=np.array(line.strip().split()[2:]).astype('float')

diff --git a/pyschism/outputs/outputs.py b/pyschism/outputs/outputs.py
@@ -28,7 +28,6 @@
 )
 from pyschism.mesh.base import Gr3
 
-
 def get_stack_id_by_datetime(dt: datetime, flattened_timevector, stacks) -> str:
     output_stack = None
     for stack_id, stack_data in stacks.items():
@@ -188,6 +187,43 @@ def animate(index):
 
         return anim
 
+    def get_zcor_interp_coefficient(self, z0, zcor, level):
+        '''
+        z0: water elevation (elev)
+        zcor: schout outputs zcor[np,nvrt] for each time record.
+        level: fixed depth (< 0, e.g., 4.5 m below the surface, level=[-4.5]) where values is interplated at.
+        '''
+        NP=zcor.shape[0]
+        nvrt=zcor.shape[1]
+        k1=np.full((NP), np.nan)
+        coeff=np.full((NP), np.nan)
+        zinter=np.ones(NP)*level+z0
+
+        #calculate kbp
+
+
+        #surface 
+        idxs=zinter>=zcor[:,-1]
+        k1[idxs]=nvrt-2
+        coeff[idxs]=1.0
+
+        #bottom
+        idxs=zinter<zcor[:,0]
+        k1[idxs]=kbp[idxs]
+        coeff[idxs]=0.0
+
+        for k in np.arange(nvrt-1):
+            idxs=(zinter>=zcor[:,k])*(zinter<zcor[:,k+1])
+            k1[idxs]=k
+            coeff[idxs]=(zinter[idxs]-zcor[idxs,k])/(zcor[idxs,k+1]-zcor[idxs,k])
+        if sum(np.isnan(np.r_[k1,coeff])) != 0:
+            sys.exit('Check vertical interpolation')
+        return np.array(k1).astype('int'), np.array(coeff)
+
+    def interpolate_to_levels(self):
+        pass
+
+
     def aggregate(self, dt: Union[datetime, timedelta, int]):
         if isinstance(dt, datetime):
             return self.aggregate_by_datetime(dt)

diff --git a/pyschism/outputs/plot.py b/pyschism/outputs/plot.py
@@ -130,7 +130,7 @@ def animate(index):
         if save:
             anim.save(
                 pathlib.Path(save),
-                writer='imagemagick',
+                writer='ffmpeg',
                 fps=fps
             )