some fixes to gewex computations

yori/aggrtools.py

@@ -59,33 +59,97 @@ def prepare_vars(var_add, group_batch, ncfile, granule_edges, gewex=False):
         # for nonzero entries
 #        if 'min_pixel_counts' in grp:
 #            var_add[grp]['min_pixel_counts'] = ncfile[grp]['min_pixel_counts']
+        _cae = ['CAE', 'CAEH', 'CAEM', 'CAEL', 'CAEW', 'CAEI', 'CAEIH']
+        _wp = ['ALWP', 'AIWP', 'AIWPH']
+        if grp in _cae:
+            if 'Pixel_Counts' in ncfile['CA'].variables:
+                cnt = ncfile['CA/Pixel_Counts'][:]
+                x, y = cnt.nonzero()
+                if len(x) == 0:
+                    del var_add[grp]
+                    continue
+                x_slice = slice(x.min(), x.max() + 1)
+                y_slice = slice(y.min(), y.max() + 1)
+                var_add[grp]['Pixel_Counts'] = cnt[x_slice, y_slice]
+                var_add[grp]['_slices'] = (x_slice, y_slice)
+                var_add[grp]['Num_Days'] = compute_valid_days(cnt)[x_slice, y_slice]
+            elif granule_edges:
+                x_slice, y_slice = granule_edges
+                var_add[grp]['_slices'] = (x_slice, y_slice)
+            else:
+                x_slice = slice(None)  # gotta read the whole grid...
+                y_slice = slice(None)
+
+        elif grp in _wp:
+            if 'Pixel_Counts' in ncfile['CLWP'].variables:
+                cnt = ncfile['CLWP/Pixel_Counts'][:]
+                x, y = cnt.nonzero()
+                if len(x) == 0:
+                    del var_add[grp]
+                    continue
+                x_slice = slice(x.min(), x.max() + 1)
+                y_slice = slice(y.min(), y.max() + 1)
+                var_add[grp]['Pixel_Counts'] = cnt[x_slice, y_slice]
+                var_add[grp]['_slices'] = (x_slice, y_slice)
+                var_add[grp]['Num_Days'] = compute_valid_days(cnt)[x_slice, y_slice]
+            elif granule_edges:
+                x_slice, y_slice = granule_edges
+                var_add[grp]['_slices'] = (x_slice, y_slice)
+            else:
+                x_slice = slice(None)  # gotta read the whole grid...
+                y_slice = slice(None)
 
-        if 'Pixel_Counts' in ncfile[grp].variables:
-            cnt = ncfile[grp]['Pixel_Counts'][:]
-            x, y = cnt.nonzero()
-            if len(x) == 0:
-                del var_add[grp]
-                continue
-            x_slice = slice(x.min(), x.max() + 1)
-            y_slice = slice(y.min(), y.max() + 1)
-            var_add[grp]['Pixel_Counts'] = cnt[x_slice, y_slice]
-            var_add[grp]['_slices'] = (x_slice, y_slice)
-            var_add[grp]['Num_Days'] = compute_valid_days(cnt)[x_slice, y_slice]
-        elif granule_edges:
-            x_slice, y_slice = granule_edges
-            var_add[grp]['_slices'] = (x_slice, y_slice)
         else:
-            x_slice = slice(None)  # gotta read the whole grid...
-            y_slice = slice(None)
-
-        for w in ncfile[grp].variables:
-            if w != ['Pixel_Counts', 'min_pixel_counts', 'min_valid_days']:
-                var_add[grp][w] = ncfile[grp][w][x_slice, y_slice]
-            if gewex is True:
-                var_add[grp]['Sum_Mean'] = ncfile[grp]['Mean'][x_slice, y_slice]
-                var_add[grp]['Sum_Squares_Mean'] = (
-                    ncfile[grp]['Mean'][x_slice, y_slice]**2)
-
+            if 'Pixel_Counts' in ncfile[grp].variables:
+                cnt = ncfile[grp]['Pixel_Counts'][:]
+                x, y = cnt.nonzero()
+                if len(x) == 0:
+                    del var_add[grp]
+                    continue
+                x_slice = slice(x.min(), x.max() + 1)
+                y_slice = slice(y.min(), y.max() + 1)
+                var_add[grp]['Pixel_Counts'] = cnt[x_slice, y_slice]
+                var_add[grp]['_slices'] = (x_slice, y_slice)
+                var_add[grp]['Num_Days'] = compute_valid_days(cnt)[x_slice, y_slice]
+            elif granule_edges:
+                x_slice, y_slice = granule_edges
+                var_add[grp]['_slices'] = (x_slice, y_slice)
+            else:
+                x_slice = slice(None)  # gotta read the whole grid...
+                y_slice = slice(None)
+
+        if grp in _cae:
+            for w in ncfile['CA'].variables:
+                suffix = grp.replace('CAE', '')
+                if w != ['Pixel_Counts', 'min_pixel_counts', 'min_valid_days']:
+                    var_add[f'CAE{suffix}'][w] = ncfile[f'CEM{suffix}'][w][x_slice, y_slice]
+                    var_add[f'CAE{suffix}/Sum_Mean'] = (
+                            ncfile[f'CA{suffix}/Mean'][x_slice, y_slice] *
+                            ncfile[f'CEM{suffix}/Mean'][x_slice, y_slice])
+                    var_add[f'CAE{suffix}/Sum_Square_Mean'] = (
+                            ncfile[f'CA{suffix}/Mean'][x_slice, y_slice] *
+                            ncfile[f'CEM{suffix}/Mean'][x_slice, y_slice])**2
+        elif grp in _wp:
+            wp_dic = {'ALWP': 'CAL', 'AIWP': 'CAI', 'AIWPH': 'CAIH'}
+            for w in ncfile['CLWP'].variables:
+                suffix = grp.replace('A', 'C')
+                ca_name = wp_dic[grp]
+                if w != ['Pixel_Counts', 'min_pixel_counts', 'min_valid_days']:
+                    var_add[grp][w] = ncfile[suffix][w][x_slice, y_slice]
+                    var_add[f'{grp}/Sum_Mean'] = (
+                            ncfile[f'{suffix}/Mean'][x_slice, y_slice] *
+                            ncfile[f'{ca_name}/Mean'][x_slice, y_slice])
+                    var_add[f'{grp}/Sum_Square_Mean'] = (
+                            ncfile[f'{suffix}/Mean'][x_slice, y_slice] *
+                            ncfile[f'{ca_name}/Mean'][x_slice, y_slice])**2
+        else:
+            for w in ncfile[grp].variables:
+                if w != ['Pixel_Counts', 'min_pixel_counts', 'min_valid_days']:
+                    var_add[grp][w] = ncfile[grp][w][x_slice, y_slice]
+                if gewex is True:
+                    var_add[grp]['Sum_Mean'] = ncfile[grp]['Mean'][x_slice, y_slice]
+                    var_add[grp]['Sum_Squares_Mean'] = (
+                        ncfile[grp]['Mean'][x_slice, y_slice]**2)
     return var_add
 
 
@@ -94,6 +158,28 @@ def initialize_arrays(group_batch, ncfile, latkey, lonkey, fv, gewex=False):
     for grp in group_batch:
         if grp == latkey or grp == lonkey:
             pass
+        elif grp in ['CAE', 'CAEH', 'CAEM', 'CAEL', 'CAEW', 'CAEI', 'CAEIH']:
+            var_in[grp] = {}
+            var_in[grp]['Mean'] = ncfile[f'CA/Mean'][:]*0
+            var_in[grp]['Standard_Deviation'] = ncfile[f'CA/Standard_Deviation'][:]*0
+            var_in[grp]['Pixel_Counts'] = ncfile[f'CA/Pixel_Counts'][:]*0
+            var_in[grp]['Sum'] = ncfile[f'CA/Sum'][:]*0
+            var_in[grp]['Sum_Squares'] = ncfile[f'CA/Sum_Squares'][:]*0
+            # var_in[grp]['Histogram_Counts'] = ncfile[f'CA/Histogram_Counts'][:]*0
+            var_in[grp]['Sum_Mean'] = ncfile[f'CA/Sum'][:]*0
+            var_in[grp]['Sum_Squares_Mean'] = ncfile[f'CA/Sum_Squares'][:]*0
+            var_in[grp]['Num_Days'] = ncfile[f'CA/Pixel_Counts'][:]*0
+        elif grp in ['ALWP', 'AIWP', 'AIWPH']:
+            var_in[grp] = {}
+            var_in[grp]['Mean'] = ncfile[f'CLWP/Mean'][:]*0
+            var_in[grp]['Standard_Deviation'] = ncfile[f'CLWP/Standard_Deviation'][:]*0
+            var_in[grp]['Pixel_Counts'] = ncfile[f'CLWP/Pixel_Counts'][:]*0
+            var_in[grp]['Sum'] = ncfile[f'CLWP/Sum'][:]*0
+            var_in[grp]['Sum_Squares'] = ncfile[f'CLWP/Sum_Squares'][:]*0
+            # var_in[grp]['Histogram_Counts'] = ncfile[f'CA/Histogram_Counts'][:]*0
+            var_in[grp]['Sum_Mean'] = ncfile[f'CLWP/Sum'][:]*0
+            var_in[grp]['Sum_Squares_Mean'] = ncfile[f'CLWP/Sum_Squares'][:]*0
+            var_in[grp]['Num_Days'] = ncfile[f'CLWP/Pixel_Counts'][:]*0
         else:
             var_in[grp] = {}
             if hasattr(ncfile[grp], 'min_pixel_counts'):

yori/yori_aggregate.py

@@ -29,6 +29,17 @@ def aggregate(flist, fname_out, satellite='', daily='',
 
     f0 = nc.Dataset(flist[0])
     group_list = list(f0.groups)
+    if gewex is True:
+        group_list.append('CAE')
+        group_list.append('CAEH')
+        group_list.append('CAEM')
+        group_list.append('CAEL')
+        group_list.append('CAEW')
+        group_list.append('CAEI')
+        group_list.append('CAEIH')
+        group_list.append('ALWP')
+        group_list.append('AIWP')
+        group_list.append('AIWPH')
 
     latdim = f0.dimensions['latitude'].size
     londim = f0.dimensions['longitude'].size
@@ -177,6 +188,7 @@ def aggregate(flist, fname_out, satellite='', daily='',
             var_add = {}
             var_add = tools.prepare_vars(var_add, group_batch, file_in,
                                          granule_edges, gewex)
+
             for k in group_batch:
                 if k == latkey or k == lonkey or k not in var_add:
                     pass