Merge branch 'ewquon-f/update_toof'

mmctools/coupling/sowfa.py

@@ -316,7 +316,8 @@ def __init__(self,
                  name='patch',
                  dateref=None,
                  datefrom=None,
-                 dateto=None):
+                 dateto=None,
+                 verbose=True):
         """
         Initialize SOWFA input object. This should be called for _each_
         inflow/outflow boundary.
@@ -344,6 +345,7 @@ def __init__(self,
             with the last timestamp in df; only used if dateref is
             specified
         """
+        self.verbose = verbose
         self.name = name
         self.dpath = os.path.join(dpath, name)
         # Create folder dpath if needed
@@ -381,20 +383,25 @@ def _check_xarray_dataset(self,
         """Do all sanity checks here"""
         for dim in self.ds.dims:
             # dimension coordinates
-            assert dim in expected_dims
+            assert dim in expected_dims, f'missing dim {dim}'
             coord = self.ds.coords[dim]
-            assert (coord.dims[0] == dim) and (len(coord.dims) == 1)
+            assert (coord.dims[0] == dim) and (len(coord.dims) == 1), \
+                    f'{dim} is not a dimension coordinate'
         # Only handle a single boundary plane at a time; boundaries
         # should be aligned with the Cartesian axes
-        dims = expected_dims.copy()
-        for dim in self.ds.dims:
-            dims.remove(dim)
-        assert (len(dims) == 1)
-        constdim = dims[0]
-        print('Input is an {:s}-boundary at {:g}'.format(constdim,
-                                                         self.ds.coords[constdim].values))
+        constdims = [dim for dim in self.ds.dims if self.ds.dims[dim]==1]
+        if len(constdims) > 0:
+            assert (len(constdims) == 1), 'more than one constant dim'
+            constdim = constdims[0]
+        else:
+            nodimcoords = [coord for coord in self.ds.coords if len(self.ds.coords[coord].dims)==0]
+            assert (len(nodimcoords) == 1), 'more than one selected dim'
+            constdim = nodimcoords[0]
+        print('Input is a {:s}-boundary at {:g}'.format(constdim,
+                                                        float(self.ds.coords[constdim])))
+        self.constdim = constdim
 
-    def write(self, fields, binary=False, gzip=False):
+    def write(self, fields, points=True, binary=False, gzip=False):
         """
         Write surface boundary conditions to SOWFA-readable input files
         for the solver in constant/boundaryData
@@ -408,9 +415,13 @@ def write(self, fields, binary=False, gzip=False):
             field name, values corresponding to dataset data variables;
             values may be a single variable (scalar) or a list/tuple of
             variables (vector)
+        points : bool, optional
+            Write out points definition for this patch
         binary : bool, optional
             Write out actual data (coordinates, scalars, vectors) in
             binary for faster I/O
+        gzip : bool, optional
+            Write out compressed data to save disk space
         """
         # check output options
         if binary and gzip:
@@ -423,10 +434,14 @@ def write(self, fields, binary=False, gzip=False):
         # make sure ordering of bnd_dims is correct
         dims = list(self.ds.dims)
         dims.remove('datetime')
-        self.bndry_dims = [dim for dim in ['x','y','height'] if dim in dims]
-        assert (len(self.bndry_dims) == 2)
+        self.bndry_dims = [
+            dim for dim in ['x','y','height']
+            if (dim in dims) and self.ds.dims[dim] > 1
+        ]
+        assert (len(self.bndry_dims) == 2), f'boundary patch dims: {str(self.bndry_dims)}'
         # write out patch/points
-        self._write_points(binary=binary, gzip=gzip)
+        if points:
+            self._write_points(binary=binary, gzip=gzip)
         # write out patch/*/field
         for fieldname,dvars in fields.items():
             if isinstance(dvars, (list,tuple)):
@@ -470,10 +485,15 @@ def _write_points(self,fname='points',binary=False,gzip=False):
         else:
             with self._open(fpath, 'w', gzip=gzip) as f:
                 np.savetxt(f, pts, fmt='(%g %g %g)', header=header, footer=')', comments='')
-        print('Wrote',N,'points to',fpath)
+        if self.verbose: 
+            print('Wrote',N,'points to',fpath)
 
     def _write_boundary_vector(self,fname,components,binary=False,gzip=False):
-        ds = self.ds.copy()
+        if self.constdim in self.ds.dims:
+            assert self.ds.dims[self.constdim] == 1
+            ds = self.ds.isel({self.constdim:0})
+        else:
+            ds = self.ds
         # add missing dimensions, if any
         for dim in self.bndry_dims:
             for var in components:
@@ -516,10 +536,15 @@ def _write_boundary_vector(self,fname,components,binary=False,gzip=False):
                         f.write(b'\n(0 0 0) // average value')
                     else:
                         f.write('\n(0 0 0) // average value')
-            print('Wrote',N,'vectors to',fpath,'at',str(tstamp))
+            if self.verbose: 
+                print('Wrote',N,'vectors to',fpath,'at',str(tstamp))
 
     def _write_boundary_scalar(self,fname,var,binary=False,gzip=False):
-        ds = self.ds.copy()
+        if self.constdim in self.ds.dims:
+            assert self.ds.dims[self.constdim] == 1
+            ds = self.ds.isel({self.constdim:0})
+        else:
+            ds = self.ds
         # add missing dimensions, if any
         for dim in self.bndry_dims:
             if dim not in ds[var].dims:
@@ -554,5 +579,6 @@ def _write_boundary_scalar(self,fname,var,binary=False,gzip=False):
                         f.write(b'\n0 // average value')
                     else:
                         f.write('\n0 // average value')
-            print('Wrote',N,'scalars to',fpath,'at',str(tstamp))
+            if self.verbose: 
+                print('Wrote',N,'scalars to',fpath,'at',str(tstamp))
 

mmctools/wrf/ts.py

@@ -188,7 +188,8 @@ def _read_towers(self):
             agl=True,
             verbose=self.verbose
         )
-        self.ds = self.ds.swap_dims({'nz':'height'}) # to facilitate ds.interp()
+        self.ds = self.ds.swap_dims({'nz':'z'}) # dimension coordinate to enable ds.interp()
+        self.ds = self.ds.rename({'z':'height'})
         if self.verbose:
             print('... done reading ts outputs')
 
@@ -229,18 +230,24 @@ def interp_to_latlon(self,latlon):
         assert (tgtlat >= wrflat[i,j]) and (tgtlat < wrflat[i,j+1])
         assert (tgtlon >= wrflon[i,j]) and (tgtlon < wrflon[i+1,j])
         # bilinear interpolation
-        f00 = self.ds.sel(nx=i  ,ny=j)
-        f10 = self.ds.sel(nx=i+1,ny=j)
-        f01 = self.ds.sel(nx=i  ,ny=j+1)
-        f11 = self.ds.sel(nx=i+1,ny=j+1)
-        finterp = f00 * (wrflon[i+1,j] - tgtlon     ) * (wrflat[i,j+1] - tgtlat     ) + \
-                  f10 * (tgtlon        - wrflon[i,j]) * (wrflat[i,j+1] - tgtlat     ) + \
-                  f01 * (wrflon[i+1,j] - tgtlon     ) * (tgtlat        - wrflat[i,j]) + \
-                  f11 * (tgtlon        - wrflon[i,j]) * (tgtlat        - wrflat[i,j])
+        f00 = self.ds.sel(nx=i  ,ny=j).drop_vars(['x','y'])
+        f10 = self.ds.sel(nx=i+1,ny=j).drop_vars(['x','y'])
+        f01 = self.ds.sel(nx=i  ,ny=j+1).drop_vars(['x','y'])
+        f11 = self.ds.sel(nx=i+1,ny=j+1).drop_vars(['x','y'])
+        coef00  = (wrflon[i+1,j] - tgtlon     ) * (wrflat[i,j+1] - tgtlat     )
+        coef10  = (tgtlon        - wrflon[i,j]) * (wrflat[i,j+1] - tgtlat     )
+        coef01  = (wrflon[i+1,j] - tgtlon     ) * (tgtlat        - wrflat[i,j])
+        coef11  = (tgtlon        - wrflon[i,j]) * (tgtlat        - wrflat[i,j])
+        finterp = coef00*f00 + coef10*f10 + coef01*f01 + coef11*f11
         finterp = finterp / ((wrflon[i+1,j] - wrflon[i,j]) * (wrflat[i,j+1] - wrflat[i,j]))
-        # note: y and z coordinates don't get interpolated
         finterp = finterp.assign_coords({'lon':tgtlon,'lat':tgtlat})
-        return finterp.drop_vars(['y','z'])
+        # fix height coordinate if needed
+        if 'height' not in finterp.coords:
+            if 'z' in finterp.coords:
+                finterp = finterp.rename_vars({'z':'height'})
+            else:
+                print('WARNING: no z or height coordinate found')
+        return finterp
 
 
     def map_to_boundary(self,i=None,j=None,k=None,allpts=False):
@@ -252,20 +259,25 @@ def map_to_boundary(self,i=None,j=None,k=None,allpts=False):
         assert np.count_nonzero([idx is not None for idx in [i,j,k]])==1, \
                 'Specify i, j, or k'
         if allpts:
-            print('WARNING: current implementation of allpts is likely to result in extreme memory usage and may crash')
+            print('WARNING: Current implementation of allpts can result in extreme memory usage')
         # interpolate to selected lat/lon
         selected_x, selected_y, selected_lat, selected_lon = \
                 self._select_boundary_latlon(i,j,k,allpts)
         if self.verbose:
             print('selected lat:',selected_lat)
             print('selected lon:',selected_lon)
+        # get all selected profiles
         dslist = self._get_datasets_at_locations(
                 selected_x,selected_y,selected_lat,selected_lon)
         # combine all interpolated profiles
         boundarydata = self._create_dataset_from_list(i,j,k,allpts,dslist)
+        if boundarydata.chunks:
+            from dask.diagnostics import ProgressBar
+            with ProgressBar():
+                boundarydata = boundarydata.compute()
         return boundarydata
 
-    def _get_datasets_at_locations(self,selected_x, selected_y, selected_lat, selected_lon):
+    def _get_datasets_at_locations(self, selected_x, selected_y, selected_lat, selected_lon):
         dslist = []
         for x,y,lat,lon in zip(selected_x, selected_y, selected_lat, selected_lon):
             ds = self.interp_to_latlon((lat,lon))
@@ -281,14 +293,12 @@ def _create_dataset_from_list(self,i,j,k,allpts,dslist):
         elif (j is not None):
             mydim = 'y'
             idx = j
-        elif (k is not None):
-            mydim = 'height'
-            idx = k
         if ((i is not None) or (j is not None)) and allpts:
-            # if allpts, interpolate side boundary profiles to exact domain heights
+            # lateral boundaries: if allpts, interpolate side boundary
+            # profiles to exact domain heights
             ds = ds.interp(height=self.domain.z)
         elif k is not None:
-            # if horizontal boundary, interpolate to constant z
+            # horizontal boundaries: interpolate to constant z
             if self.verbose:
                 print('interpolating to',self.domain.z[k])
             ds = ds.interp(height=self.domain.z[k])

mmctools/wrf/utils.py

@@ -752,7 +752,7 @@ def extract_column_from_wrfdata(fpath, coords,
                                 spatial_filter='interpolate',L_filter=0.0,
                                 additional_fields=[],
                                 verbose=False,
-                               ):
+                                **kwargs):
     """
     Extract a column of time-height data for a specific site from a
     4-dimensional WRF output file
@@ -788,7 +788,7 @@ def extract_column_from_wrfdata(fpath, coords,
             'Spatial filtering type "'+spatial_filter+'" not recognised'
 
     # Load WRF data
-    ds = xr.open_dataset(fpath)
+    ds = xr.open_dataset(fpath,**kwargs)
     tdim, zdim, ydim, xdim = get_wrf_dims(ds)