more mods for wrf

ldcpy/calcs.py

@@ -1155,7 +1155,8 @@ def lag1(self) -> xr.DataArray:
             if 'dayofyear' in self._ds.attrs.keys():
                 key = f'{self._time_dim_name}.dayofyear'
             else:
-                key = 'time'
+                key = f'{self._time_dim_name}'
+
             grouped = self._ds.groupby(key, squeeze=False)
             if self._time_dim_name in self._ds.attrs.keys():
                 self._deseas_resid = grouped - grouped.mean(dim=self._time_dim_name)
@@ -1220,7 +1221,7 @@ def fft2(self) -> xr.DataArray:
             # if hasattr(self._ds, 'units'):
             #    self._fft2.attrs['units'] = f'{self._ds.units}'
             self._fft2 = self._fft2.rename(
-                {'dim_0': 'time', 'dim_1': self._lat_dim_name, 'dim_2': self._lon_dim_name}
+                {'dim_0': self._time_dim_name, 'dim_1': self._lat_dim_name, 'dim_2': self._lon_dim_name}
             )
             self._fft2 = self._fft2.assign_coords(
                 {

ldcpy/plot.py

@@ -323,7 +323,7 @@ def spatial_plot(self, da_sets, titles, data_type):
         cmin = []
 
         # lat/lon could be 1 or 2d and have different names
-        # TODO  - will need to adjust for WRF for U and V?
+        #print(da_sets[0].cf.coordinates)
         lon_coord_name = da_sets[0].cf.coordinates['longitude'][0]
         lat_coord_name = da_sets[0].cf.coordinates['latitude'][0]
 
@@ -660,7 +660,7 @@ def time_series_plot(
         )
 
 
-        print(group_string)
+        #print(group_string)
         for i in range(da_sets.sets.size):
             if self._group_by is not None:
                 plt.plot(
@@ -962,7 +962,6 @@ class in ldcpy.plot, for more information about the available calcs see ldcpy.Da
             - rms
             - sum
             - sum_squared
-            - corr_lag1
             - quantile
             - lag1
             - standardized_mean
@@ -1175,14 +1174,23 @@ class in ldcpy.plot, for more information about the available calcs see ldcpy.Da
         if ds.data_type == 'cam-fv':  # 1D
             mp.title_lat = subsets[0][lat_coord_name].data[0]
             mp.title_lon = subsets[0][lon_coord_name].data[0] - 180
-        elif ds.data_type == 'pop' or 'wrf':  # 2D
+        elif ds.data_type == 'pop' :  # 2D
             # lon should be 0- 360
             mylat = subsets[0][lat_coord_name].data[0]
             mylon = subsets[0][lon_coord_name].data[0]
             if mylon < 0:
                 mylon = mylon + 360
             mp.title_lat = mylat
             mp.title_lon = mylon
+        elif ds.data_type == 'wrf':
+            #for wrf, we don't want the x and y values, we need to
+            # convert to lat and long (to do)
+            # so for now, we output user's input for the title
+            mylat = lat
+            mylon = lon
+            mp.title_lat = mylat
+            mp.title_lon = mylon
+
         else:
             print('ERROR: unknown data type')
 

ldcpy/util.py

@@ -10,7 +10,7 @@
 from .calcs import Datasetcalcs, Diffcalcs
 
 
-def collect_datasets(data_type, varnames, list_of_ds, labels, coords_ds=None, **kwargs):
+def collect_datasets(data_type, varnames, list_of_ds, labels, coords_ds=None, file_sizes=None, **kwargs):
     """
     Concatonate several different xarray datasets across a new
     "collection" dimension, which can be accessed with the specified
@@ -31,6 +31,8 @@ def collect_datasets(data_type, varnames, list_of_ds, labels, coords_ds=None, **
         The respective label to access data from each dataset (also used in plotting fcns)
     coords_ds : xarray dataset
         (optional) Specify an additional file that contains lat/lon corrds (common for WRF data)
+    file_sizes : list
+         (optional) sizes of files that each dataset corresponds to (used to print in compare_stats table
     **kwargs :
         (optional) – Additional arguments passed on to xarray.concat(). A list of available arguments can
         be found here: https://xarray-test.readthedocs.io/en/latest/generated/xarray.concat.html
@@ -61,6 +63,10 @@ def collect_datasets(data_type, varnames, list_of_ds, labels, coords_ds=None, **
     indx = np.unique(sz)
     assert indx.size == 1, 'ERROR: all datasets must have the same length time dimension'
 
+    #file sizes?
+    if file_sizes is not None:
+        assert len(file_sizes) == len(labels), 'ERROR::collect_dataset dataset list and file sizes arguments must be the same length'
+
     # wrf data must contain lat/lon info in same file if a coord_file is not specified
     if data_type == 'wrf':
         if coords_ds is None:
@@ -119,6 +125,13 @@ def preprocess_vars(ds, varnames):
     full_ds.attrs['data_type'] = data_type
     full_ds.attrs['file_size'] = None
 
+    #file sizes?    
+    if file_sizes is not None:
+        file_size_dict = {}
+        for i, myfile in enumerate(list_of_ds):
+            file_size_dict[labels[i]]  = file_sizes[i]
+        full_ds.attrs['file_size'] = file_size_dict
+
     # from other copy of this function
     for v in varnames[:-1]:
         new_ds = []
@@ -723,7 +736,7 @@ def subset_data(
     lev=None,
     start=None,
     end=None,
-    time_dim_name='time',
+    time_dim_name=None,
     vertical_dim_name=None,
     lat_coord_name=None,
     lon_coord_name=None,
@@ -733,7 +746,9 @@ def subset_data(
     """
     ds_subset = ds
 
-    # print(ds.cf.describe())
+
+
+    #print(ds.cf.describe())
 
     if lon_coord_name is None:
         lon_coord_name = ds.cf.coordinates['longitude'][0]
@@ -750,6 +765,10 @@ def subset_data(
 
     # print(lat_coord_name, lon_coord_name, vertical_dim_name)
 
+    if time_dim_name is None:
+        time_dim_name = ds.cf.coordinates['time'][0]
+
+
     latdim = ds_subset.cf[lon_coord_name].ndim
     # need dim names
     dd = ds_subset.cf['latitude'].dims
@@ -789,33 +808,48 @@ def subset_data(
 
     elif latdim == 2:
 
-        # print(ds_subset)
 
         if lat is not None:
             if lon is not None:
 
-                # lat is -90 to 90
-                # lon should be 0- 360
-                ad_lon = lon
-                if ad_lon < 0:
-                    ad_lon = ad_lon + 360
-
-                mlat = ds_subset[lat_coord_name].compute()
-                mlon = ds_subset[lon_coord_name].compute()
-                # euclidean dist for now....
-                di = np.sqrt(np.square(ad_lon - mlon) + np.square(lat - mlat))
-                index = np.where(di == np.min(di))
-                xmin = index[0][0]
-                ymin = index[1][0]
-
-                # Don't want if it's a land point
-                check = ds_subset.isel(nlat=xmin, nlon=ymin, time=1).compute()
-                if np.isnan(check):
-                    print(
-                        'You have chosen a lat/lon point with Nan values (i.e., a land point). Plot will not make sense.'
+                if ds_subset.data_type == "pop":
+
+                    # lat is -90 to 90
+                    # lon should be 0- 360
+                    ad_lon = lon
+                    if ad_lon < 0:
+                        ad_lon = ad_lon + 360
+
+                    mlat = ds_subset[lat_coord_name].compute()
+                    mlon = ds_subset[lon_coord_name].compute()
+                    # euclidean dist for now....
+                    di = np.sqrt(np.square(ad_lon - mlon) + np.square(lat - mlat))
+                    index = np.where(di == np.min(di))
+                    xmin = index[0][0]
+                    ymin = index[1][0]
+
+                    # Don't want if it's a land point
+                    check = ds_subset.isel(nlat=xmin, nlon=ymin, time=1).compute()
+                    if np.isnan(check):
+                        print('You have chosen a lat/lon point with Nan values (i.e., a land point). Plot will not make sense.'
                     )
-                ds_subset = ds_subset.isel({lat_dim_name: [xmin], lon_dim_name: [ymin]})
-
+                    ds_subset = ds_subset.isel({lat_dim_name: [xmin], lon_dim_name: [ymin]})
+
+                elif ds_subset.data_type == 'wrf':
+                    #print(lat_dim_name, lon_dim_name)
+                    ad_lon = lon
+                    mlat = ds_subset[lat_coord_name].compute()
+                    mlon = ds_subset[lon_coord_name].compute()
+                    di = np.sqrt(np.square(ad_lon - mlon) + np.square(lat - mlat))
+                    index = np.where(di == np.min(di))
+                    xmin = index[0][0]
+                    ymin = index[1][0]
+                    #print(xmin, ymin)
+                    #TO DO: add error checking for if it's out of bounds
+                    #check = ds_subset.isel({lat_dim_name: [xmin], lon_dim_name : [ymin], time_dim_name: [1]}).compute()
+                    #print(check)
+                    ds_subset = ds_subset.isel({lat_dim_name: [xmin], lon_dim_name : [ymin]})
+
                 # ds_subset.compute()
 
     return ds_subset