only have water source shares for withdrawal; fix queries; hack to capture transnational basins within the USA

Author: thurber

tethys/datareader/easy_query.py

@@ -1,7 +1,7 @@
 import gcamreader
 
 
-def easy_query(variable, year_axis=None, **kwargs):
+def easy_query(variable, year_axis=None, replace_filters=False, **kwargs):
     """Build a query for a GCAM database
 
     kwargs act as filters on nodes, eg. "sector='Beef'" is converted to the xpath "[@type='sector' and @name='Beef']"
@@ -24,10 +24,11 @@ def easy_query(variable, year_axis=None, **kwargs):
     :return: gcamreader Query object
     """
 
-    if not kwargs:
-        filters = dict(sector=None)  # match nodes where @type='sector' by default if no filters are specified
+    if replace_filters:
+        filters = kwargs
     else:
-        filters = kwargs  # update the filters with user-provided kwargs
+        filters = dict(sector=None)  # match nodes where @type='sector' by default
+        filters.update(kwargs)  # update the filters with user-provided kwargs
 
     # filters are separated by "//*" to match any descendant
     xpath = "*" + "//*".join(handle_filter(k, v) for k, v in filters.items()) + f"//{variable}/node()"

tethys/model.py

@@ -206,7 +206,7 @@ def downscale(self, distribution, inputs, region_masks):
 
         return out
 
-    def disaggregate_source(self, downscaled):
+    def disaggregate_source(self):
         """Disaggregate water demand by source (runoff, groundwater, desal)"""
 
         print('Disaggregating Source')
@@ -221,11 +221,7 @@ def disaggregate_source(self, downscaled):
         # apply shares to the gridded region masks
         gridded_shares = get_shares.generate_gridded_shares(shares_df)
 
-        # disaggregate water supply source for each downscaled demand in each grid cell
-        downscaled_disaggregated_sw = (downscaled * gridded_shares['runoff']).compute()
-        downscaled_disaggregated_gw = (downscaled * gridded_shares['groundwater']).compute()
-
-        return gridded_shares, downscaled_disaggregated_sw, downscaled_disaggregated_gw
+        return gridded_shares
 
 
     def run_model(self):
@@ -238,6 +234,17 @@ def run_model(self):
         self._load_region_masks()
         self._load_inputs()
 
+        # disaggregate water supply source
+        if self.source_disaggregation:
+            gshares = self.disaggregate_source()
+
+            # store the disaggregated results
+            self.griddedshares = gshares
+
+            if self.output_dir is not None:
+                filename = os.path.join(self.output_dir, f'gridded_runoff_shares.nc')
+                gshares['runoff'].to_netcdf(filename)
+
         for supersector, rules in self.downscaling_rules.items():
 
             print(f'Downscaling {supersector}')
@@ -284,16 +291,6 @@ def run_model(self):
                 downscaled.to_dataset(dim='sector').to_netcdf(filename, encoding=encoding)
                 downscaled = xr.open_dataset(filename).to_array(dim='sector')  # hopefully this keeps dask happy
 
-            # disaggregate water supply source and update the objects
-            if self.source_disaggregation:
-                gshares, sw, gw = self.disaggregate_source(downscaled)
-
-                # store the disaggregated results
-                self.griddedshares = gshares
-                self.disaggregated_sw = sw
-                self.disaggregated_gw = gw
-
-
             if self.temporal_config is not None:
                 # calculate the monthly distributions (share of annual) for each year
 
@@ -315,7 +312,7 @@ def run_model(self):
                 if self.output_dir is not None:
                     filename = os.path.join(self.output_dir, f'{supersector}_{self.demand_type}_monthly.nc')
                     encoding = {sector: {'zlib': True, 'complevel': 5} for sector in downscaled.sector.data}
-                    downscaled.to_dataset(dim='sector').to_netcdf(filename, encoding=encoding)
+                    downscaled.to_dataset(dim='sector').compute().to_netcdf(filename, encoding=encoding)
                     downscaled = xr.open_dataset(filename).to_array(dim='sector')  # hopefully this keeps dask happy
 
             self.outputs.update(downscaled.to_dataset(dim='sector'))

tethys/utils/source_disaggregation.py

@@ -34,9 +34,9 @@ def calculate_shares(self):
 
         :return: Processed DataFrame with calculated shares.
         """
-        # query GCAM database for water used by source for the demand type (W or C)
-        query_type = f'*_water {self.demand_type}'
-        shares_df = self.conn.runQuery(easy_query('production', resource=query_type))
+        # query GCAM database for water used by source for the demand type (W only)
+        query_type = f'*_water withdrawals'
+        shares_df = self.conn.runQuery(easy_query('production', replace_filters=True, resource=query_type))
 
         # extract and clean resource names (e.g., remove '_water withdrawals')
         shares_df['resource'] = shares_df['resource'].apply(extract_resource_name)
@@ -54,6 +54,15 @@ def calculate_shares(self):
             ['scenario', 'region', 'region_resourcemap', 'year']
         )['value'].transform(lambda x: x / x.sum() if x.sum() != 0 else 0)
 
+        # capture basins that are mostly belonging to other countries
+        if basin_name_mapping == 'basin_name_mapping_im3':
+            shares_df = shares_df.replace({'region_resourcemap': {
+                'Canada_FraserR': 'USA_FraserR',
+                'Canada_GreatLakes': 'USA_GreatLakes',
+                'Mexico_MexCstNW': 'USA_MexCstNW',
+                'Canada_NelsonR': 'USA_NelsonR',
+            }})
+
         # filter shares based on the region masks from the config file, otherwise all regions would be included
         shares_df = shares_df[shares_df['region_resourcemap'].isin(self.region_masks.region.values)]