Updating multi pannel plotting capability in fig_generator

AGENTS.md

@@ -0,0 +1,103 @@
+# psPlotKit — Agent Instructions
+
+## Project Overview
+
+psPlotKit is a Python plotting/analysis toolkit for data generated by the **Parameter Sweep (PS) tool** and **Loop tool** — primarily HDF5 (`.h5`) and JSON files produced by WaterTAP/IDAES process simulations. It imports hierarchical simulation results, manages unit conversions, computes levelized costs, and produces publication-quality matplotlib figures.
+
+## Architecture — Two-Layer Pipeline
+
+```
+data_manager/          →  data_plotter/
+  PsDataImport             FigureGenerator  (low-level matplotlib wrapper)
+  PsData                   linePlotter, boxPlotter, BreakDownPlotter, MapPlotter
+  PsDataManager (dict)     (each plotter consumes a PsDataManager)
+  PsCosting
+  PsExpression
+  PsDataExporter
+```
+
+1. **`PsDataImport`** — reads `.h5` / `.json`, auto-discovers directories (parameter sweep runs), indexes data keys, and produces `PsData` objects.
+2. **`PsData`** — wraps a numpy array with physical units (via the `quantities` library), supports unit conversion (`to_units`), unit assignment (`assign_units`), feasibility masking, and label generation for matplotlib.
+3. **`PsDataManager(dict)`** — central data store. Keys are `(directory_tuple, data_key)` tuples. Provides `load_data`, `register_data_key`, `reduce_data` (stacking/min-reduction across sweep axes), `normalize_data`, `eval_function`, `get_costing`, and data selection helpers.
+4. **`PsCosting`** — computes CAPEX/OPEX breakdowns and levelized costs per device group from costing blocks in the simulation data.
+5. **`PsExpression`** — evaluates mathematical expressions on data stored in `PsDataManager`.
+6. **`PsDataExporter`** — exports data from `PsDataManager` to CSV, JSON, or HDF5.
+7. **Plotters** (`linePlotter`, `BreakDownPlotter`, `boxPlotter`, `MapPlotter`) — high-level plotting classes that accept a `PsDataManager`, select data, and delegate rendering to `FigureGenerator`.
+8. **`FigureGenerator`** — thin matplotlib wrapper providing `plot_line`, `plot_bar`, `plot_map`, `plot_area`, axis styling, colorbars, and figure save/show.
+
+## Key Conventions
+
+### Naming & Deprecation Pattern
+Every class has a **PascalCase canonical name** and a **legacy camelCase/lowercase subclass alias** that logs a deprecation warning. Always use the canonical names:
+- `PsData` (not `psData`), `PsDataManager` (not `psDataManager`)
+- `FigureGenerator` (not `figureGenerator`), `BreakDownPlotter` (not `breakDownPlotter`)
+- `PsCosting` (not `psCosting`), `PsExpression` (not `psExpression`)
+- `PsDataExporter` (not `psDataExporter`)
+
+### Data Key Dictionaries
+When specifying data to import, keys are passed as dicts with specific fields:
+```python
+{"filekey": "fs.costing.LCOW", "return_key": "LCOW", "units": "USD/m**3"}
+# Optional fields: "assign_units", "conversion_factor", "directories"
+```
+
+### Unit Handling
+Units flow through `quantities` (aliased as `qs`). Custom units (`USD`, `PPM`) are registered in `CustomUnits`. String unit normalization happens in `PsData._convert_string_unit` (e.g., `"gal"` → `"US_liquid_gallon"`, `"°C"` → `"*degC"`, `"USD/a"` → `"USD/year"`).
+
+### Logging
+Every module creates its logger via `psPlotKit.util.logger.define_logger(__name__, "DisplayName")`. Use this pattern — do not use `logging.getLogger` directly.
+
+### Composite Tuple Keys
+`PsDataManager` stores data under composite tuple keys built from directory labels and data keys (e.g., `(("erd_type", "pressure_exchanger"), "membrane_cost", "LCOW")`). The `add_data(dir_key, data_key, PsData)` method constructs these automatically.
+
+## Development Setup
+
+```bash
+conda env create -f psplotkit.yml   # Python 3.11, pytest, editable install
+conda activate psPlotKit
+```
+
+Dependencies: `h5py`, `quantities`, `matplotlib`, `numpy`, `scipy`, `sigfig`, `pyyaml`.
+
+## Testing
+
+```bash
+pytest src/psPlotKit/data_manager/tests/
+```
+
+Tests live alongside the modules they cover (e.g., `data_manager/tests/`). Test fixtures load `.h5` files from the same test directory using `os.path.dirname(os.path.abspath(__file__))`.
+
+## File Layout
+
+- `src/psPlotKit/data_manager/` — data import, storage, costing logic, expressions, exporter
+- `src/psPlotKit/data_plotter/` — all plotting classes and figure generation
+- `src/psPlotKit/util/` — shared logger and filesystem helpers
+- `src/psPlotKit/examples/` — runnable example scripts with sample data in `examples/data/`
+- `docs/` — mkdocs documentation (API reference, guides, examples)
+- `legacy_data_collator.py` — legacy monolithic class; do not extend, prefer the modular pipeline above
+
+## Common Tasks
+
+### Adding a New Data Key
+1. Create data key dict with `filekey`, `return_key`, `units`
+2. Register with `PsDataManager.register_data_key(data_key_dict)`
+3. Load data with `PsDataManager.load_data()`
+
+### Adding a New Plotter
+1. Create new plotter class in `src/psPlotKit/data_plotter/`
+2. Inherit from base plotter pattern (accept `PsDataManager` in constructor)
+3. Use `FigureGenerator` for low-level matplotlib operations
+4. Add to `src/psPlotKit/data_plotter/__init__.py`
+
+### Adding a New Costing Package
+1. Create new module in `src/psPlotKit/data_manager/costing_packages/`
+2. Implement costing computation functions
+3. Register in `PsCosting` or use directly
+
+## Important Notes
+
+- **Never modify `legacy_data_collator.py`** — it is deprecated
+- **Always use canonical PascalCase class names** in new code
+- **Unit handling** must go through `quantities` library
+- **Logging** must use `define_logger` from `psPlotKit.util.logger`
+- **Tests** should be placed in `tests/` subdirectory alongside the module

src/psPlotKit/data_manager/costing_packages/tests/test_watertap_costing.py

@@ -299,7 +299,9 @@ def test_complex():
 
     # lets create costing pacakage
     package = WaterTapCostingPackage(
-        costing_block="costing", validation_key="costing.LCOW"
+        costing_block="costing",
+        validation_key="costing.LCOW",
+        sec_validation_key="costing.SEC",
     )
     package.register_product_flow("avg_product_flow_rate")
 

src/psPlotKit/data_manager/costing_packages/watertap_costing.py

@@ -2,7 +2,12 @@
 
 
 class WaterTapCostingPackage(PsCostingPackage):
-    def __init__(self, costing_block="fs.costing", validation_key="fs.costing.LCOW"):
+    def __init__(
+        self,
+        costing_block="fs.costing",
+        validation_key="fs.costing.LCOW",
+        sec_validation_key="fs.costing.SEC",
+    ):
         super().__init__(costing_block=costing_block)
         self.costing_package_name = "watertap"
 
@@ -41,7 +46,9 @@ def __init__(self, costing_block="fs.costing", validation_key="fs.costing.LCOW")
             lambda ek: ek.total_capital_cost * ek.capital_recovery_factor
             + ek.total_operating_cost,
         )
+
         self.add_validation("LCOW", file_key=validation_key, rtol=1e-4)
+        self.add_validation("SEC", file_key=sec_validation_key, rtol=1e-4)
         self.register_fraction("LCOW")
         self.register_fraction("LCOW_opex", "LCOW")
         self.register_fraction("LCOW_capex", "LCOW")
@@ -50,6 +57,11 @@ def register_product_flow(
         self, file_key="fs.product.properties[0.0].flow_vol_phase[Liq]"
     ):
         self.add_parameter("product_flow", file_key=file_key)
+        self.add_formula(
+            "SEC",
+            lambda ek: (ek.aggregate_electricity_flow) / (ek.product_flow),
+            units="kWh/m**3",
+        )
         self.add_formula(
             "LCOW",
             lambda ek: (ek.total_annualized_cost)

src/psPlotKit/data_manager/ps_costing.py

@@ -52,6 +52,7 @@
 """
 
 import difflib
+import time
 
 import numpy as np
 
@@ -526,34 +527,33 @@ def build(self, load_data=True, error_on_validation_failure=True):
                 after key registration.  Set to *False* if data is already
                 loaded.
         """
+
+        ts = time.time()
         self.data_manager._load_registered_data_files()
         self._discover_keys()
         self._check_discovery_status()
         self._register_parameters()
         self._register_discovered_keys()
         self._register_validation_keys()
-
         if load_data:
             self.data_manager.load_data(evaluate_expressions=False)
-
         self._register_zero_sentinel()
+
         self._build_group_expressions()
         self._build_flow_expressions()
         self._build_per_group_flow_expressions()
         self._build_formula_expressions()
         self._build_per_group_formula_expressions()
         self._build_total_formula_expressions()
         self._build_fraction_expressions()
-
         self.data_manager.evaluate_expressions()
-
         _logger.info(
-            "Costing build complete - {} group(s), {} formula(e).".format(
+            "Costing build complete took {} seconds - {} group(s), {} formula(e).".format(
+                round(time.time() - ts, 2),
                 len(self.costing_groups),
                 len(self.costing_package.formulae),
             )
         )
-
         self._validate(error_on_failure=error_on_validation_failure)
 
     # ------------------------------------------------------------------
@@ -1166,12 +1166,17 @@ def _build_per_group_formula_expressions(self):
                     flow_type in self.costing_package.flow_costs
                     and group_flow_types.get(flow_type)
                 )
+
                 alias_map["{}_flow_cost".format(flow_type)] = (
                     gft_rk if has_cost else None
                 )
                 alias_map["aggregate_flow_cost_{}".format(flow_type)] = (
                     gft_rk if has_cost else None
                 )
+                gft_rk = ("costing", group.name, "{}_flow".format(flow_type))
+                alias_map["aggregate_{}_flow".format(flow_type)] = (
+                    gft_rk if has_cost else None
+                )
 
             for fdef in self.costing_package.formulae:
                 ek = self.data_manager.get_expression_keys()
@@ -1199,7 +1204,7 @@ def _build_per_group_formula_expressions(self):
                 self._per_group_formula_keys.setdefault(fdef["return_key"], []).append(
                     group_rk
                 )
-                _logger.info("Registered per-group formula '{}'.".format(group_rk))
+                _logger.debug("Registered per-group formula '{}'.".format(group_rk))
 
     def _build_total_formula_expressions(self):
         """Build total (sum across groups) expressions for each formula.

src/psPlotKit/data_manager/ps_data_manager.py

@@ -825,8 +825,15 @@ def sort_idxs(idxs):
                     map_idxs = temp_map_idxs
                 units = np.unique(map_units)
                 if len(units) > 1:
-                    _logger.info("Units are inconsistent, using dimensionless")
-                    units = "dimensionless"
+                    _map_units = []
+                    for k in range(len(map_units)):
+                        true_sum = np.nansum(map_data[k])
+                        if true_sum != 0:
+                            _map_units.append(map_units[k])
+                    units = np.unique(_map_units)
+                    if len(units) > 1:
+                        _logger.info("Units are inconsistent, using dimensionless")
+                        units = "dimensionless"
 
                 else:
                     units = units[0]
@@ -867,9 +874,10 @@ def sort_idxs(idxs):
                     new_keys.append(data_key)
                     new_keys.append(stack_keys)
                 except:
-                    _logger.error(
-                        "Could not stack data for dir {}, key {}".format(uq, data_key)
-                    )
+                    pass
+                    # _logger.error(
+                    #     "Could not stack data for dir {}, key {}".format(uq, data_key)
+                    # )
         return unique_dirs, keys_to_process
 
     def add_mask(self, directory, indexes, data_shape=None, shape="1D"):