From ebac92ccab6d0f6096c5d39fbdf5ab15b4a8f00e Mon Sep 17 00:00:00 2001
From: ktehranchi <ktehranchi@stanford.edu>
Date: Fri, 22 Nov 2024 10:59:42 -0800
Subject: [PATCH] Fix variable scenario filtering, and availability TS
 multiplier

---
 src/r2x/parser/plexos.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/r2x/parser/plexos.py b/src/r2x/parser/plexos.py
index eeee647a..c5788ad3 100644
--- a/src/r2x/parser/plexos.py
+++ b/src/r2x/parser/plexos.py
@@ -424,6 +424,7 @@ def _construct_branches(self, default_model=MonitoredLine):
             pl.col("parent_class_name") == ClassEnum.System.value
         )
         system_lines = self._get_model_data(system_lines)
+
         lines_pivot = system_lines.pivot(  # noqa: PD010
             index=DEFAULT_INDEX,
             on="property_name",
@@ -613,6 +614,7 @@ def _construct_generators(self):  # noqa: C901
         system_generators_filter = (pl.col("child_class_name") == str(ClassEnum.Generator)) & (
             pl.col("parent_class_name") == str(ClassEnum.System)
         )
+
         system_generators = self._get_model_data(system_generators_filter)
         if self.config.feature_flags.get("plexos-csv", None):
             system_generators.write_csv("generators.csv")
@@ -1169,7 +1171,7 @@ def _set_unit_capacity(self, record):  # noqa: C901
         At this point of the code, all the properties should be either a single value, or a `SingleTimeSeries`
         """
         if not (availability := record.get("available")):
-            logger.warning("Unit `{}` is not activated on the model.", record["name"])
+            logger.debug("Unit `{}` is not activated on the model.", record["name"])
             return
         record["active_power_limits"] = self._get_active_power_limits(record)
 
@@ -1233,7 +1235,7 @@ def _get_active_power_limits(self, record) -> MinMax:
         if active_power_max := record.get("max_active_power"):
             if isinstance(active_power_max, SingleTimeSeries):
                 active_power_max = np.nanmax(active_power_max.data)
-        active_power_limits_max = active_power_max or record["base_power"]
+        active_power_limits_max = (active_power_max or record["base_power"]) * record["available"]
         return MinMax(active_power_limits_min, active_power_limits_max)
 
     def _plexos_table_data(self) -> list[tuple]:
@@ -1244,10 +1246,22 @@ def _plexos_table_data(self) -> list[tuple]:
     def _polarize_data(self, object_data: list[tuple]) -> pl.DataFrame:
         data = pl.from_records(object_data, schema=SIMPLE_QUERY_COLUMNS_SCHEMA)
 
+        # Filtering Variables to select only the correct band-id which should be used
+        # when accessing Variable data. We may also need to apply this filtering to
+        # other Enum classes.
+        variables = data.filter(pl.col("child_class_name") == ClassEnum.Variable.name)
+        variables = variables.group_by("object_id", maintain_order=True).map_groups(
+            lambda groupdf: groupdf.filter(pl.col("band") == pl.col("band").min())
+        )
+
+        data = data.filter(pl.col("child_class_name") != ClassEnum.Variable.name)
+        data = pl.concat([data, variables])
+
         # Create a lookup map to find nested objects easily
         object_map = data.select(["object_id", "tag_datafile", "tag_datafile_object_id"]).to_dict(
             as_series=False
         )
+
         self.id_to_name = dict(zip(object_map["object_id"], object_map["tag_datafile"]))
         self.id_to_tag_id = dict(zip(object_map["object_id"], object_map["tag_datafile_object_id"]))
         return data
@@ -1647,6 +1661,7 @@ def _handle_record(self, record: dict[str, Any], prop_name, prop_value, unit):
                 if action:
                     value = self._apply_action(action, prop_value, value)
                 value = self._parse_value(value, variable_name=mapped_property_name, unit=unit)
+
             case {"tag_timeslice": str()}:
                 value = self._parse_value(prop_value, variable_name=mapped_property_name, unit=unit)
             case _: