This updates ps costing and fixes minor issues

src/psPlotKit/data_manager/costing_packages/watertap_costing.py

@@ -10,10 +10,14 @@ def __init__(self, costing_block="fs.costing", validation_key="fs.costing.LCOW")
         self.add_parameter("total_investment_factor")
         self.add_parameter("capital_recovery_factor")
         self.add_parameter("maintenance_labor_chemical_factor")
-        self.add_parameter("electricity_cost")
         self.add_parameter("wacc")
         self.add_parameter("plant_lifetime")
-        self.add_flow_cost("electricity", "electricity_cost", units="USD/yr")
+        self.add_flow_cost(
+            "electricity",
+            "electricity_cost",
+            parameter_units="USD/kWh",
+            aggregate_units="USD/yr",
+        )
         self.add_formula(
             "total_capital_cost",
             lambda ek: ek.aggregate_capital_cost * ek.total_investment_factor,

src/psPlotKit/data_manager/ps_costing.py

@@ -55,6 +55,7 @@
 
 import numpy as np
 
+import quantities as qs
 from psPlotKit.data_manager.ps_data import PsData
 from psPlotKit.data_manager.ps_expression import ExpressionNode
 from psPlotKit.util.logger import define_logger
@@ -336,7 +337,14 @@ def add_formula(self, return_key, builder, units=None, assign_units=None):
             }
         )
 
-    def add_flow_cost(self, flow_type, cost_parameter, units=None, assign_units=None):
+    def add_flow_cost(
+        self,
+        flow_type,
+        cost_parameter_key,
+        parameter_units=None,
+        parameter_assign_units=None,
+        aggregate_units="USD/yr",
+    ):
         """Register the cost multiplier for a flow type.
 
         The *cost_parameter* must be a parameter name registered via
@@ -346,15 +354,20 @@ def add_flow_cost(self, flow_type, cost_parameter, units=None, assign_units=None
 
         Args:
             flow_type: Name of the flow type (e.g. ``"electricity"``).
-            cost_parameter: Parameter name giving the unit cost
-                (e.g. ``"electricity_cost"`` in $/kWh).
-            units: Units to convert the flow cost to after evaluation.
-            assign_units: Units to assign to the flow cost result.
+            cost_parameter_key: Parameter key with costs that is imported from the h5 (e.g. ``"electricity_cost"``).
+            parameter_units: Units to convert the cost parameter to
+            parameter_assign_units: Units to assign to cost parameter if there is none.
+            aggregate_units: Units for final aggregate flow after flow is multiplied by cost parameter (generally USD/year).
         """
+        self.add_parameter(
+            cost_parameter_key,
+            units=parameter_units,
+            assign_units=parameter_assign_units,
+        )
         self._flow_costs[flow_type] = {
-            "cost_parameter": cost_parameter,
-            "units": units,
-            "assign_units": assign_units,
+            "cost_parameter": cost_parameter_key,
+            "units": aggregate_units,
+            "assign_units": None,
         }
 
     @property
@@ -725,6 +738,7 @@ def _match_unit_segments(parts, unit_name):
         """
         unit_parts = _split_key(unit_name)
         n = len(unit_parts)
+        matched = False
         for start in range(len(parts) - n + 1):
             matched = True
             for j, up in enumerate(unit_parts):
@@ -741,6 +755,7 @@ def _match_unit_segments(parts, unit_name):
                         break
             if matched:
                 return start + n - 1  # index of last matched segment
+
         return None
 
     @staticmethod
@@ -785,7 +800,8 @@ def _find_matching_unit_keys(available_keys, unit_name, suffix):
                 continue
             remainder_parts = parts[unit_end + 1 :]
             stripped = ".".join(p.split("[")[0] for p in remainder_parts)
-            if stripped == suffix:
+            unstripped = ".".join(remainder_parts)
+            if stripped == suffix or unstripped == suffix:
                 matches.append(key)
         return matches
 
@@ -981,7 +997,11 @@ def _build_flow_expressions(self):
                 cost_param_rk = ("costing", cost_param_name)
                 ek = self.data_manager.get_expression_keys()
                 flow_cost_rk = ("costing", "{}_flow_cost".format(flow_type))
-                expr = ek[agg_flow_rk] * ek[cost_param_rk]
+                expr = (
+                    ek[agg_flow_rk]
+                    * ek[cost_param_rk]
+                    # * self.costing_package.operational_time
+                )
                 self.data_manager.register_expression(
                     expr,
                     return_key=flow_cost_rk,

src/psPlotKit/data_manager/ps_data_manager.py

@@ -183,7 +183,7 @@ def register_data_key(
     def get_expression_keys(self, warn_on_sanitize=False):
         """Return the live :class:`ExpressionKeys` reference for this manager.
 
-        The returned object is kept in sync with the manager — new keys
+        The returned object is kept in sync with the manager - new keys
         added via :meth:`add_data` or :meth:`register_data_key` are
         automatically available without calling this method again.
 
@@ -691,7 +691,7 @@ def sort_idxs(idxs):
             return sorted_idxs, idxs, idx_type
 
         for udir in self.keys():
-            if str(new_directory) not in str(udir) and str(data_key) in str(
+            if new_directory not in udir and str(data_key) in str(
                 self._get_data_key(udir)
             ):
 
@@ -877,7 +877,6 @@ def stack_all_data(self, stack_keys, pad_missing_data, stack_directory=None):
         current_keys = self.data_keys[:]
         for data_key in current_keys:
             if str(stack_directory) not in str(data_key):
-                # print("stack_keys", stack_keys, data_key)
                 self.generate_data_stack(
                     stack_keys,
                     data_key,
@@ -1177,7 +1176,7 @@ def evaluate_expressions(self):
                                 PsData.__new__(PsData)._convert_string_unit(units)
                             )
                     except Exception:
-                        # Dimensions don't match — fall back to labelling only
+                        # Dimensions don't match - fall back to labelling only
                         _effective_units = None
                         _effective_assign = units
                 result = PsData(

src/psPlotKit/data_manager/ps_expression.py

@@ -95,7 +95,7 @@ def _op_node(cls, op, left, right):
     def _as_node(other):
         if isinstance(other, ExpressionNode):
             return other
-        if isinstance(other, (int, float, np.integer, np.floating)):
+        if isinstance(other, (int, float, np.integer, np.floating, qs.Quantity)):
             return ExpressionNode._const_node(other)
         raise TypeError(
             "Cannot combine ExpressionNode with type {}".format(type(other))

src/psPlotKit/data_manager/tests/test_ps_costing.py

@@ -165,7 +165,7 @@ def test_add_flow_cost(self):
         assert pkg.flow_costs == {
             "electricity": {
                 "cost_parameter": "electricity_cost",
-                "units": None,
+                "units": "USD/yr",
                 "assign_units": None,
             }
         }
@@ -573,11 +573,13 @@ def test_flow_cost_electricity(self):
 
         # electricity_flow_cost = 15*0.07, 28*0.07
         flow_cost = dm.get_data("d", ("costing", "electricity_flow_cost"))
-        np.testing.assert_array_almost_equal(flow_cost.data, [1.05, 1.96])
+        np.testing.assert_array_almost_equal(
+            flow_cost.data, [9204.103409, 17180.993031]
+        )
 
         # aggregate_flow_cost = same (only one flow type)
         agg_cost = dm.get_data("d", ("costing", "aggregate_flow_cost"))
-        np.testing.assert_array_almost_equal(agg_cost.data, [1.05, 1.96])
+        np.testing.assert_array_almost_equal(agg_cost.data, [9204.103409, 17180.993031])
 
     def test_aggregate_flow_cost_multiple_types(self):
         """Two flow types should each aggregate separately; total sums them."""
@@ -589,8 +591,10 @@ def test_aggregate_flow_cost_multiple_types(self):
         dm.add_data("d", ("costing", "chemical_cost"), [1.0, 1.0])
 
         pkg = PsCostingPackage()
-        pkg.add_flow_cost("electricity", "electricity_cost")
-        pkg.add_flow_cost("chemicals", "chemical_cost")
+        pkg.add_flow_cost(
+            "electricity", "electricity_cost", aggregate_units="dimensionless"
+        )
+        pkg.add_flow_cost("chemicals", "chemical_cost", aggregate_units="dimensionless")
 
         g = PsCostingGroup("RO")
         cm = PsCostingManager(dm, pkg, [g])
@@ -682,19 +686,19 @@ def test_per_group_flow_expressions(self):
 
         # Per-group flow costs
         pump_fc = dm.get_data("d", ("costing", "Pumps", "electricity_flow_cost"))
-        np.testing.assert_array_almost_equal(pump_fc.data, [0.7, 1.4])
+        np.testing.assert_array_almost_equal(pump_fc.data, [6136.06894, 12272.137879])
         erd_fc = dm.get_data("d", ("costing", "ERD", "electricity_flow_cost"))
-        np.testing.assert_array_almost_equal(erd_fc.data, [0.21, 0.35])
+        np.testing.assert_array_almost_equal(erd_fc.data, [1840.820682, 3068.03447])
 
         # Per-group aggregate flow cost
         pump_afc = dm.get_data("d", ("costing", "Pumps", "flow_cost"))
-        np.testing.assert_array_almost_equal(pump_afc.data, [0.7, 1.4])
+        np.testing.assert_array_almost_equal(pump_afc.data, [6136.06894, 12272.137879])
         erd_afc = dm.get_data("d", ("costing", "ERD", "flow_cost"))
-        np.testing.assert_array_almost_equal(erd_afc.data, [0.21, 0.35])
+        np.testing.assert_array_almost_equal(erd_afc.data, [1840.820682, 3068.03447])
 
         # Global aggregate should still be the sum
         agg = dm.get_data("d", ("costing", "aggregate_flow_cost"))
-        np.testing.assert_array_almost_equal(agg.data, [0.91, 1.75])
+        np.testing.assert_array_almost_equal(agg.data, [7976.889621, 15340.172349])
 
     # ------------------------------------------------------------------
     # Per-group formula expressions
@@ -785,7 +789,9 @@ def test_per_group_formula_with_flow(self):
         dm.add_data("d", "capital_recovery_factor", [0.1, 0.1])
 
         pkg = PsCostingPackage()
-        pkg.add_flow_cost("electricity", "electricity_cost")
+        pkg.add_flow_cost(
+            "electricity", "electricity_cost", aggregate_units="dimensionless"
+        )
         pkg.add_formula(
             "total_cost",
             lambda ek: (
@@ -821,7 +827,9 @@ def test_per_group_formula_missing_capex_uses_zero(self):
         dm.add_data("d", "capital_recovery_factor", [0.1, 0.1])
 
         pkg = PsCostingPackage()
-        pkg.add_flow_cost("electricity", "electricity_cost")
+        pkg.add_flow_cost(
+            "electricity", "electricity_cost", aggregate_units="dimensionless"
+        )
         pkg.add_formula(
             "total_cost",
             lambda ek: (
@@ -879,7 +887,7 @@ def test_flow_from_data_manager(self):
         np.testing.assert_array_almost_equal(agg_flow.data, [10.0, 20.0])
 
         flow_cost = dm.get_data("d", ("costing", "electricity_flow_cost"))
-        np.testing.assert_array_almost_equal(flow_cost.data, [0.5, 1.0])
+        np.testing.assert_array_almost_equal(flow_cost.data, [4382.906385, 8765.812771])
 
     def test_formula_aggregate_flow_cost_per_type(self):
         """Formulas can reference individual flow-type costs via
@@ -891,7 +899,9 @@ def test_formula_aggregate_flow_cost_per_type(self):
 
         pkg = PsCostingPackage()
         pkg.add_parameter("product_flow", file_key="dummy")
-        pkg.add_flow_cost("electricity", "electricity_cost")
+        pkg.add_flow_cost(
+            "electricity", "electricity_cost", aggregate_units="dimensionless"
+        )
         pkg.add_formula(
             "levelized_electricity",
             lambda ek: ek.aggregate_flow_cost_electricity / ek.product_flow,

src/psPlotKit/data_plotter/ps_break_down_plotter.py

@@ -137,8 +137,13 @@ def check_key_in_dir(self, udir, key):
                     return True
             else:
                 for di in d:
-                    if key == di:
-                        return True
+                    if isinstance(di, str):
+                        if key == di:
+                            return True
+                    else:
+                        for _di in di:
+                            if key == _di:
+                                return True
         return False
 
     def _get_group_options(self, selected_keys, xdata, ydata):
@@ -159,8 +164,6 @@ def _get_group_options(self, selected_keys, xdata, ydata):
                         opts["hatch"] = self.hatch_options[i]
 
                     self.hatch_groups[key] = copy.deepcopy(opts)
-                    # if key not in self.line_indexes:
-                    #     self.line_indexes[key] = {"idx": 0, "auto": True}
             for akey in self.area_groups:
                 _label = None
                 if isinstance(akey, dict):
@@ -212,12 +215,13 @@ def _get_group_options(self, selected_keys, xdata, ydata):
                         cur_line["color"] = color
                         cur_line["label"] = plot_label
                         self.plot_areas[akey] = cur_line
+
         self.plot_order = []
         for akey in self.area_groups:
             if isinstance(akey, dict):
                 akey, item = list(akey.items())[0]
             for key in self.plot_areas.keys():
-                if akey in key:
+                if akey == key:
                     self.plot_order.append(key)
         # assert False
 
@@ -271,7 +275,6 @@ def plot_imported_data(self):
             self.fig.plot_area([], [], **items)
         old_data = 0
         current_data = None
-
         for linelabel in self.plot_order:
             line = self.plot_areas[linelabel]
             if line.get("label") in plotted_legend: