Add ability to select subset of upgrades and skip monthly simulation in buildstock_visualizer

buildstock_query/aggregate_query.py

@@ -71,6 +71,7 @@ def aggregate_annual(self, *,
         restrict = [(self._bsq._bs_completed_status_col, [self._bsq.db_schema.completion_values.success])] + restrict
         query = self._bsq._add_join(query, join_list)
         query = self._bsq._add_restrict(query, restrict)
+        query = self._bsq._add_avoid(query, params.avoid)
         query = self._bsq._add_group_by(query, group_by_selection)
         query = self._bsq._add_order_by(query, group_by_selection if params.sort else [])
 
@@ -190,6 +191,7 @@ def aggregate_timeseries(self, params: TSQuery):
             params.restrict = list(params.restrict) + [(self._bsq._ts_upgrade_col, [upgrade_id])]
 
         query = self._bsq._add_restrict(query, params.restrict)
+        query = self._bsq._add_avoid(query, params.avoid)
         query = self._bsq._add_group_by(query, group_by_selection)
         query = self._bsq._add_order_by(query, group_by_selection if params.sort else [])
         query = query.limit(params.limit) if params.limit else query

buildstock_query/aggregate_query.pyi

@@ -21,6 +21,7 @@ class BuildStockAggregate:
                          join_list: Sequence[tuple[AnyTableType, AnyColType, AnyColType]] = [],
                          weights: Sequence[Union[str, tuple]] = [],
                          restrict: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
+                         avoid: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
                          get_quartiles: bool = False,
                          get_nonzero_count: bool = False,
                          ) -> str:
@@ -36,6 +37,7 @@ class BuildStockAggregate:
                          join_list: Sequence[tuple[AnyTableType, AnyColType, AnyColType]] = [],
                          weights: Sequence[Union[str, tuple]] = [],
                          restrict: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
+                         avoid: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
                          get_quartiles: bool = False,
                          get_nonzero_count: bool = False,
                          ) -> pd.DataFrame:
@@ -51,6 +53,7 @@ class BuildStockAggregate:
                          join_list: Sequence[tuple[AnyTableType, AnyColType, AnyColType]] = [],
                          weights: Sequence[Union[str, tuple]] = [],
                          restrict: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
+                         avoid: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
                          get_quartiles: bool = False,
                          get_nonzero_count: bool = False,
                          ) -> Union[pd.DataFrame, str]:
@@ -79,6 +82,8 @@ class BuildStockAggregate:
 
             restrict: The list of where condition to restrict the results to. It should be specified as a list of tuple.
                     Example: `[('state',['VA','AZ']), ("build_existing_model.lighting",['60% CFL']), ...]`
+            avoid: Just like restrict, but the opposite. It will only include rows that do not match (any of) the
+                   conditions.
             get_quartiles: If true, return the following quartiles in addition to the sum for each enduses:
                         [0, 0.02, .25, .5, .75, .98, 1]. The 0% quartile is the minimum and the 100% quartile
                         is the maximum.
@@ -109,6 +114,7 @@ class BuildStockAggregate:
                              join_list: Sequence[tuple[AnyTableType, AnyColType, AnyColType]] = [],
                              weights: Sequence[Union[str, tuple]] = [],
                              restrict: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
+                             avoid: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
                              split_enduses: bool = False,
                              collapse_ts: bool = False,
                              timestamp_grouping_func: Optional[str] = None,
@@ -125,6 +131,7 @@ class BuildStockAggregate:
                              join_list: Sequence[tuple[AnyTableType, AnyColType, AnyColType]] = [],
                              weights: Sequence[Union[str, tuple]] = [],
                              restrict: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
+                             avoid: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
                              split_enduses: bool = False,
                              collapse_ts: bool = False,
                              timestamp_grouping_func: Optional[str] = None,
@@ -143,6 +150,7 @@ class BuildStockAggregate:
                              join_list: Sequence[tuple[AnyTableType, AnyColType, AnyColType]] = [],
                              weights: Sequence[Union[str, tuple]] = [],
                              restrict: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
+                             avoid: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = [],
                              split_enduses: bool = False,
                              collapse_ts: bool = False,
                              timestamp_grouping_func: Optional[str] = None,

buildstock_query/main.py

@@ -40,11 +40,12 @@ def __init__(self,
                  table_name: Union[str, tuple[str, Optional[str], Optional[str]]],
                  db_schema: Optional[str] = None,
                  buildstock_type: Literal['resstock', 'comstock'] = 'resstock',
-                 sample_weight: Optional[Union[int, float]] = None,
+                 sample_weight_override: Optional[Union[int, float]] = None,
                  region_name: str = 'us-west-2',
                  execution_history: Optional[str] = None,
                  skip_reports: bool = False,
                  athena_query_reuse: bool = True,
+                 **kwargs,
                  ) -> None:
         """A class to run Athena queries for BuildStock runs and download results as pandas DataFrame.
 
@@ -60,8 +61,8 @@ def __init__(self,
                 It is also different between the version in OEDI and default version from BuildStockBatch. This argument
                 controls the assumed schema. Allowed values are 'resstock_default', 'resstock_oedi', 'comstock_default'
                 and 'comstock_oedi'. Defaults to 'resstock_default' for resstock and 'comstock_default' for comstock.
-            sample_weight (str, optional): Specify a custom sample_weight. Otherwise, the default is 1 for ComStock and
-                uses sample_weight in the run for ResStock.
+            sample_weight_override (str, optional): Specify a custom sample_weight. Otherwise, the default is 1 for
+                ComStock and uses sample_weight in the run for ResStock.
             region_name (str, optional): the AWS region where the database exists. Defaults to 'us-west-2'.
             execution_history (str, optional): A temporary file to record which execution is run by the user,
                 to help stop them. Will use .execution_history if not supplied. Generally, not required to supply a
@@ -71,6 +72,7 @@ def __init__(self,
             athena_query_reuse (bool, optional): When true, Athena will make use of its built-in 7 day query cache.
                 When false, it will not. Defaults to True. One use case to set this to False is when you have modified
                 the underlying s3 data or glue schema and want to make sure you are not using the cached results.
+            kargs: Any other extra keyword argument supported by the QueryCore can be supplied here
         """
         db_schema = db_schema or f"{buildstock_type}_default"
         self.params = BSQParams(
@@ -79,7 +81,7 @@ def __init__(self,
             buildstock_type=buildstock_type,
             table_name=table_name,
             db_schema=db_schema,
-            sample_weight_override=sample_weight,
+            sample_weight_override=sample_weight_override,
             region_name=region_name,
             execution_history=execution_history,
             athena_query_reuse=athena_query_reuse

buildstock_query/query_core.py

@@ -972,7 +972,7 @@ def _simple_label(self, label: str):
         label = label.removeprefix(self.db_schema.column_prefix.output)
         return label
 
-    def _add_restrict(self, query, restrict, bs_only=False):
+    def _add_restrict(self, query, restrict, *, bs_only=False):
         if not restrict:
             return query
         where_clauses = []
@@ -988,6 +988,22 @@ def _add_restrict(self, query, restrict, bs_only=False):
         query = query.where(*where_clauses)
         return query
 
+    def _add_avoid(self, query, avoid, *, bs_only=False):
+        if not avoid:
+            return query
+        where_clauses = []
+        for col_str, criteria in avoid:
+            col = self._get_column(col_str, table_name=self.bs_table) if bs_only else self._get_column(col_str)
+            if isinstance(criteria, (list, tuple)):
+                if len(criteria) > 1:
+                    where_clauses.append(self._get_column(col).not_in(criteria))
+                    continue
+                else:
+                    criteria = criteria[0]
+            where_clauses.append(col != criteria)
+        query = query.where(*where_clauses)
+        return query
+
     def _get_name(self, col):
         if isinstance(col, tuple):
             return col[1]

buildstock_query/schema/query_params.py

@@ -11,6 +11,7 @@ class AnnualQuery(BaseModel):
     sort: bool = True
     join_list: Sequence[tuple[AnyTableType, AnyColType, AnyColType]] = Field(default_factory=list)
     restrict: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = Field(default_factory=list)
+    avoid: Sequence[tuple[AnyColType, Union[str, int, Sequence[Union[int, str]]]]] = Field(default_factory=list)
     weights: Sequence[Union[str, tuple, AnyColType]] = Field(default_factory=list)
     get_quartiles: bool = False
     get_nonzero_count: bool = False

buildstock_query/schema/run_params.py

@@ -8,7 +8,7 @@ class RunParams(BaseModel):
     db_name: str
     table_name: Union[str, tuple[str, Optional[str], Optional[str]]]
     buildstock_type: Literal['resstock', 'comstock'] = 'resstock'
-    db_schema: Optional[str] = 'resstock_raw'
+    db_schema: Optional[str] = None
     sample_weight_override: Optional[Union[int, float]] = None
     region_name: str = 'us-west-2'
     execution_history: Optional[str] = None

buildstock_query/schema/utilities.py

@@ -2,13 +2,13 @@
 from typing import Union, Any, Sequence
 from pydantic import BaseModel
 import sqlalchemy as sa
-from sqlalchemy.sql.elements import Label
+from sqlalchemy.sql.elements import Label, ColumnElement
 from sqlalchemy.sql.selectable import Subquery
 
 # from buildstock_query import BuildStockQuery  # can't import due to circular import
 
 
-SACol = sa.Column
+SACol = Union[sa.Column, ColumnElement]
 SALabel = Label
 DBColType = Union[SALabel,  SACol]
 DBTableType = sa.Table