From 038a35635c5222da73a38f1b6e064455cbb9175a Mon Sep 17 00:00:00 2001
From: avdudchenko <33663878+avdudchenko@users.noreply.github.com>
Date: Fri, 19 Jun 2026 00:38:28 -0400
Subject: [PATCH 1/5] adding primary key

---
 .../data_manager/ps_data_exporter.py          | 43 ++++++++++++++++-
 src/psPlotKit/data_manager/ps_data_manager.py |  7 ++-
 .../tests/test_ps_data_exporter.py            | 47 +++++++++++++++++++
 3 files changed, 94 insertions(+), 3 deletions(-)

diff --git a/src/psPlotKit/data_manager/ps_data_exporter.py b/src/psPlotKit/data_manager/ps_data_exporter.py
index f94321a..bd75f09 100644
--- a/src/psPlotKit/data_manager/ps_data_exporter.py
+++ b/src/psPlotKit/data_manager/ps_data_exporter.py
@@ -18,6 +18,10 @@ class PsDataExporter:
     save_location : str
         File path (single directory) or folder path (multiple directories)
         where the CSV output will be written.
+    first_key : str, optional
+        Data key that should always appear as the first column in the exported
+        CSV files. If provided and the key exists in the data, it will be
+        moved to the front of the column order.
 
     Notes
     -----
@@ -31,9 +35,10 @@ class PsDataExporter:
       :meth:`PsData.set_label`).
     """
 
-    def __init__(self, ps_data_manager, save_location):
+    def __init__(self, ps_data_manager, save_location, first_key=None):
         self.ps_data_manager = ps_data_manager
         self.save_location = save_location
+        self.first_key = first_key
 
     @staticmethod
     def _ensure_csv_extension(path):
@@ -101,6 +106,40 @@ def _group_data_by_directory(self):
 
         return grouped
 
+    def _reorder_with_first(self, data_items):
+        """Reorder data items to place first_key at the front if specified.
+
+        Parameters
+        ----------
+        data_items : list[tuple]
+            List of (data_key, PsData) pairs.
+
+        Returns
+        -------
+        list[tuple]
+            Reordered list with first_key at the front if applicable.
+        """
+        if not self.first_key:
+            return data_items
+
+        reordered = []
+        first_item = None
+        for item in data_items:
+            if item[0] == self.first_key:
+                first_item = item
+            else:
+                reordered.append(item)
+
+        if first_item:
+            reordered.insert(0, first_item)
+            _logger.info("Placed '{}' as first column".format(self.first_key))
+        else:
+            _logger.warning(
+                "first_key '{}' not found in data keys".format(self.first_key)
+            )
+
+        return reordered
+
     def _build_header(self, data_items):
         """Build a CSV header row from a list of (data_key, PsData) pairs.
 
@@ -208,6 +247,7 @@ def _export_single(self, grouped):
 
         _logger.info("Single directory detected, exporting to single CSV file")
         data_items = list(grouped.values())[0]
+        data_items = self._reorder_with_first(data_items)
         headers = self._build_header(data_items)
         rows = self._build_rows(data_items)
         _logger.info(
@@ -273,6 +313,7 @@ def _export_multiple(self, grouped):
         for dir_key, data_items in grouped.items():
             filename = self._dir_key_to_filename(dir_key) + ".csv"
             file_path = os.path.join(folder, filename)
+            data_items = self._reorder_with_first(data_items)
             headers = self._build_header(data_items)
             rows = self._build_rows(data_items)
             _logger.info(
diff --git a/src/psPlotKit/data_manager/ps_data_manager.py b/src/psPlotKit/data_manager/ps_data_manager.py
index a053ccf..6722ab2 100644
--- a/src/psPlotKit/data_manager/ps_data_manager.py
+++ b/src/psPlotKit/data_manager/ps_data_manager.py
@@ -465,7 +465,7 @@ def get_data_keys(self, selected_keys):
             data_keys = self.data_keys[:]
         return data_keys
 
-    def export_data_to_csv(self, save_location):
+    def export_data_to_csv(self, save_location, first_key=None):
         """Export all loaded data to CSV file(s).
 
         Convenience wrapper around :class:`PsDataExporter`.  If this
@@ -478,6 +478,9 @@ def export_data_to_csv(self, save_location):
         save_location : str
             File path for single-directory export (e.g. ``"results.csv"``)
             or folder path for multi-directory export.
+        first_key : str, optional
+            Data key that should always appear as the first column in the
+            exported CSV files.
 
         Returns
         -------
@@ -486,7 +489,7 @@ def export_data_to_csv(self, save_location):
         """
         from psPlotKit.data_manager.ps_data_exporter import PsDataExporter
 
-        exporter = PsDataExporter(self, save_location)
+        exporter = PsDataExporter(self, save_location, first_key)
         return exporter.export()
 
     def display_loaded_contents(self):
diff --git a/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py b/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
index bd014d0..58952af 100644
--- a/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
+++ b/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
@@ -352,3 +352,50 @@ def test_manual_single_via_manager(self, manual_single_dm, tmp_path):
         assert len(headers) == 2
         assert len(rows) == 3
         assert float(rows[0][0]) == pytest.approx(1.0)
+
+
+class TestFirstKeyParameter:
+    def test_first_key_places_column_first_single_dir(self, manual_single_dm, tmp_path):
+        save_path = str(tmp_path / "first_key.csv")
+        exporter = PsDataExporter(manual_single_dm, save_path, first_key="metric_2")
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert "metric_2" in headers[0]
+        assert "metric_1" in headers[1]
+
+    def test_first_key_places_column_first_multi_dir(self, manual_multi_dm, tmp_path):
+        save_folder = str(tmp_path / "first_key_export")
+        exporter = PsDataExporter(manual_multi_dm, save_folder, first_key="metric_2")
+        written = exporter.export()
+
+        for path in written:
+            with open(path, "r") as f:
+                reader = csv.reader(f)
+                headers = next(reader)
+            assert "metric_2" in headers[0]
+
+    def test_first_key_not_found_logs_warning(self, manual_single_dm, tmp_path):
+        save_path = str(tmp_path / "missing_key.csv")
+        exporter = PsDataExporter(manual_single_dm, save_path, first_key="nonexistent")
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert len(headers) == 2
+        assert "metric_1" in headers[0]
+
+    def test_first_key_via_manager_method(self, manual_single_dm, tmp_path):
+        save_path = str(tmp_path / "manager_first_key.csv")
+        written = manual_single_dm.export_data_to_csv(save_path, first_key="metric_2")
+
+        with open(written[0], "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert "metric_2" in headers[0]

From 35e5ee18314980e94b2f9c7df72c6fcc6764ce74 Mon Sep 17 00:00:00 2001
From: avdudchenko <33663878+avdudchenko@users.noreply.github.com>
Date: Fri, 19 Jun 2026 14:29:52 -0400
Subject: [PATCH 2/5] added exporting rules

---
 .../data_manager/ps_data_exporter.py          | 249 ++++++++++++++++--
 src/psPlotKit/data_manager/ps_data_manager.py |  10 +-
 .../tests/test_ps_data_exporter.py            | 192 ++++++++++++++
 3 files changed, 430 insertions(+), 21 deletions(-)

diff --git a/src/psPlotKit/data_manager/ps_data_exporter.py b/src/psPlotKit/data_manager/ps_data_exporter.py
index bd75f09..894119a 100644
--- a/src/psPlotKit/data_manager/ps_data_exporter.py
+++ b/src/psPlotKit/data_manager/ps_data_exporter.py
@@ -22,6 +22,15 @@ class PsDataExporter:
         Data key that should always appear as the first column in the exported
         CSV files. If provided and the key exists in the data, it will be
         moved to the front of the column order.
+    export_keys : list, optional
+        List of keys to include in export. Each key is matched against data keys:
+        - For string keys, a match occurs if the key is in the list
+        - For tuple keys, a match occurs if any element of the tuple is in the list
+        If None, all non-internal keys are exported.
+    exact_keys : list, optional
+        List of keys to include in export with exact matching. The entire key
+        (including all tuple elements) must match exactly. If specified, this
+        takes precedence over export_keys.
 
     Notes
     -----
@@ -33,12 +42,26 @@ class PsDataExporter:
     * Column headers are built from each :class:`PsData` object's
       ``data_label`` and ``mpl_units`` attributes (set by
       :meth:`PsData.set_label`).
+    * Internal keys (e.g., ``_zero_sentinel``) are automatically excluded
+      from export.
+    * For tuple data keys (e.g., ``("costing", "pump")``), multi-row headers
+      are generated with each tuple element in a separate row, making the
+      structure clearer when reading CSVs in spreadsheet applications.
     """
 
-    def __init__(self, ps_data_manager, save_location, first_key=None):
+    def __init__(
+        self,
+        ps_data_manager,
+        save_location,
+        first_key=None,
+        export_keys=None,
+        exact_keys=None,
+    ):
         self.ps_data_manager = ps_data_manager
         self.save_location = save_location
         self.first_key = first_key
+        self.export_keys = export_keys
+        self.exact_keys = exact_keys
 
     @staticmethod
     def _ensure_csv_extension(path):
@@ -84,6 +107,8 @@ def _group_data_by_directory(self):
         composite keys are plain strings) and multi-directory managers
         (where composite keys are tuples).
 
+        Internal keys like ``_zero_sentinel`` are excluded from export.
+
         Returns
         -------
         dict
@@ -100,12 +125,77 @@ def _group_data_by_directory(self):
                 dir_key = tuple(dir_key)
             data_key = dm._get_data_key(composite_key)
 
+            # Skip internal keys like _zero_sentinel
+            if self._is_internal_key(data_key):
+                continue
+
+            # Skip keys not in export filter
+            if not self._should_export_key(data_key):
+                continue
+
             if dir_key not in grouped:
                 grouped[dir_key] = []
             grouped[dir_key].append((data_key, ps_data))
 
         return grouped
 
+    def _should_export_key(self, data_key):
+        """Check if a data key should be exported based on export_keys/exact_keys filters.
+
+        Parameters
+        ----------
+        data_key : str or tuple
+            The data key to check (may contain nested tuples).
+
+        Returns
+        -------
+        bool
+            True if the key should be included in export.
+        """
+        # Flatten nested tuples for matching
+        flat_key = (
+            self._flatten_key(data_key) if isinstance(data_key, tuple) else data_key
+        )
+
+        if self.exact_keys is not None:
+            return data_key in self.exact_keys
+
+        if self.export_keys is not None:
+            if isinstance(flat_key, tuple):
+                # For tuple keys, match if ANY element is in export_keys
+                return any(elem in self.export_keys for elem in flat_key)
+            else:
+                return flat_key in self.export_keys
+
+        return True
+
+    def _is_internal_key(self, data_key):
+        """Check if a data key is an internal (non-exportable) key.
+
+        Parameters
+        ----------
+        data_key : str or tuple
+            The data key to check (may contain nested tuples).
+
+        Returns
+        -------
+        bool
+            True if the key should be excluded from export.
+        """
+
+        def check_item(item):
+            if isinstance(item, tuple):
+                return any("_zero_sentinel" in str(part) for part in item)
+            return "_zero_sentinel" in str(item)
+
+        if isinstance(data_key, tuple):
+            # Check all elements including nested ones
+            for elem in data_key:
+                if check_item(elem):
+                    return True
+            return False
+        return check_item(data_key)
+
     def _reorder_with_first(self, data_items):
         """Reorder data items to place first_key at the front if specified.
 
@@ -125,7 +215,14 @@ def _reorder_with_first(self, data_items):
         reordered = []
         first_item = None
         for item in data_items:
-            if item[0] == self.first_key:
+            data_key = item[0]
+            # For tuple keys, also check flattened version for matching
+            key_to_check = (
+                self._flatten_key(data_key) if isinstance(data_key, tuple) else data_key
+            )
+            if data_key == self.first_key or (
+                isinstance(key_to_check, tuple) and self.first_key in key_to_check
+            ):
                 first_item = item
             else:
                 reordered.append(item)
@@ -141,11 +238,11 @@ def _reorder_with_first(self, data_items):
         return reordered
 
     def _build_header(self, data_items):
-        """Build a CSV header row from a list of (data_key, PsData) pairs.
+        """Build CSV header rows from a list of (data_key, PsData) pairs.
 
-        Each column header is formatted as ``"label (units)"`` using the
-        ``data_label`` and ``mpl_units`` attributes produced by
-        :meth:`PsData.set_label`.
+        For string data keys, returns a single header row. For tuple data keys,
+        returns multiple header rows with the tuple components in separate rows,
+        making multi-level keys easier to read.
 
         Parameters
         ----------
@@ -154,8 +251,29 @@ def _build_header(self, data_items):
 
         Returns
         -------
-        list[str]
-            Column header strings.
+        list[list[str]]
+            List of header rows (each row is a list of strings).
+            Single row if all keys are strings, multiple rows for tuple keys.
+        """
+        has_tuple_keys = any(isinstance(key, tuple) for key, _ in data_items)
+
+        if not has_tuple_keys:
+            return self._build_single_row_header(data_items)
+
+        return self._build_multi_row_header(data_items)
+
+    def _build_single_row_header(self, data_items):
+        """Build a single header row from data items.
+
+        Parameters
+        ----------
+        data_items : list[tuple]
+            List of (data_key, PsData) pairs.
+
+        Returns
+        -------
+        list[list[str]]
+            Single row as a list containing one list of header strings.
         """
         headers = []
         for _, ps_data in data_items:
@@ -166,7 +284,99 @@ def _build_header(self, data_items):
             else:
                 header = str(label)
             headers.append(header)
-        return headers
+        return [headers]
+
+    @staticmethod
+    def _flatten_key(data_key):
+        """Flatten nested tuples into a single-level tuple.
+
+        For example, ("costing", ("stage 1", "pump"), "LCOW") becomes
+        ("costing", "stage 1", "pump", "LCOW").
+
+        Parameters
+        ----------
+        data_key : tuple
+            The data key potentially containing nested tuples.
+
+        Returns
+        -------
+        tuple
+            Flattened tuple with all elements at single level.
+        """
+        flattened = []
+        for elem in data_key:
+            if isinstance(elem, tuple):
+                flattened.extend(elem)
+            else:
+                flattened.append(elem)
+        return tuple(flattened)
+
+    def _build_multi_row_header(self, data_items):
+        """Build multi-row headers for tuple data keys.
+
+        Parameters
+        ----------
+        data_items : list[tuple]
+            List of (data_key, PsData) pairs.
+
+        Returns
+        -------
+        list[list[str]]
+            List of header rows. Shorter tuples are aligned to the bottom,
+            so all labels appear in the final row.
+        """
+        # Flatten nested tuples to determine max depth
+        flattened_keys = []
+        for key, _ in data_items:
+            if isinstance(key, tuple):
+                flattened_keys.append(self._flatten_key(key))
+            else:
+                flattened_keys.append(key)
+
+        max_depth = max(
+            (len(key) for key in flattened_keys if isinstance(key, tuple)), default=1
+        )
+
+        header_rows = [[] for _ in range(max_depth)]
+
+        for idx, (data_key, ps_data) in enumerate(data_items):
+            label = ps_data.data_label
+            units = getattr(ps_data, "mpl_units", "-")
+
+            flat_key = flattened_keys[idx]
+
+            # If label equals data_key, use the last tuple element as label
+            if label == data_key and isinstance(data_key, tuple):
+                label = flat_key[-1]
+
+            if units and units != "-":
+                full_label = "{} ({})".format(label, units)
+            else:
+                full_label = str(label)
+
+            if isinstance(flat_key, tuple):
+                key_len = len(flat_key)
+                # Calculate starting row (align shorter tuples to bottom)
+                start_row = max_depth - key_len
+                for row_idx in range(max_depth):
+                    if row_idx < start_row:
+                        header_rows[row_idx].append("")
+                    elif row_idx == max_depth - 1:
+                        # Last row gets the full label
+                        header_rows[row_idx].append(full_label)
+                    else:
+                        # Middle rows get tuple elements starting from start_row
+                        elem_idx = row_idx - start_row
+                        header_rows[row_idx].append(str(flat_key[elem_idx]))
+            else:
+                # String keys go in the last row only
+                for row_idx in range(max_depth):
+                    if row_idx == max_depth - 1:
+                        header_rows[row_idx].append(full_label)
+                    else:
+                        header_rows[row_idx].append("")
+
+        return header_rows
 
     def _build_rows(self, data_items):
         """Build row data from a list of (data_key, PsData) pairs.
@@ -203,15 +413,15 @@ def _build_rows(self, data_items):
             rows.append(row)
         return rows
 
-    def _write_csv(self, file_path, headers, rows):
+    def _write_csv(self, file_path, header_rows, rows):
         """Write headers and rows to a CSV file.
 
         Parameters
         ----------
         file_path : str
             Destination file path.
-        headers : list[str]
-            Column header strings.
+        header_rows : list[list[str]]
+            List of header row lists (supports multi-level headers).
         rows : list[list]
             Row-major data.
         """
@@ -221,7 +431,8 @@ def _write_csv(self, file_path, headers, rows):
 
         with open(file_path, "w", newline="") as f:
             writer = csv.writer(f)
-            writer.writerow(headers)
+            for header_row in header_rows:
+                writer.writerow(header_row)
             writer.writerows(rows)
 
         _logger.info("Saved CSV: {}".format(file_path))
@@ -248,14 +459,14 @@ def _export_single(self, grouped):
         _logger.info("Single directory detected, exporting to single CSV file")
         data_items = list(grouped.values())[0]
         data_items = self._reorder_with_first(data_items)
-        headers = self._build_header(data_items)
+        header_rows = self._build_header(data_items)
         rows = self._build_rows(data_items)
         _logger.info(
             "Writing {} columns and {} rows to {}".format(
-                len(headers), len(rows), save_path
+                len(header_rows[0]), len(rows), save_path
             )
         )
-        self._write_csv(save_path, headers, rows)
+        self._write_csv(save_path, header_rows, rows)
         return [save_path]
 
     def _dir_key_to_filename(self, dir_key):
@@ -314,14 +525,14 @@ def _export_multiple(self, grouped):
             filename = self._dir_key_to_filename(dir_key) + ".csv"
             file_path = os.path.join(folder, filename)
             data_items = self._reorder_with_first(data_items)
-            headers = self._build_header(data_items)
+            header_rows = self._build_header(data_items)
             rows = self._build_rows(data_items)
             _logger.info(
                 "Directory '{}': writing {} columns and {} rows to {}".format(
-                    dir_key, len(headers), len(rows), filename
+                    dir_key, len(header_rows[0]), len(rows), filename
                 )
             )
-            self._write_csv(file_path, headers, rows)
+            self._write_csv(file_path, header_rows, rows)
             written.append(file_path)
 
         _logger.info(
diff --git a/src/psPlotKit/data_manager/ps_data_manager.py b/src/psPlotKit/data_manager/ps_data_manager.py
index 6722ab2..dc92d59 100644
--- a/src/psPlotKit/data_manager/ps_data_manager.py
+++ b/src/psPlotKit/data_manager/ps_data_manager.py
@@ -465,7 +465,7 @@ def get_data_keys(self, selected_keys):
             data_keys = self.data_keys[:]
         return data_keys
 
-    def export_data_to_csv(self, save_location, first_key=None):
+    def export_data_to_csv(self, save_location, first_key=None, export_keys=None, exact_keys=None):
         """Export all loaded data to CSV file(s).
 
         Convenience wrapper around :class:`PsDataExporter`.  If this
@@ -481,6 +481,12 @@ def export_data_to_csv(self, save_location, first_key=None):
         first_key : str, optional
             Data key that should always appear as the first column in the
             exported CSV files.
+        export_keys : list, optional
+            List of keys to include in export. For tuple keys, any matching
+            element will be included. If None, all non-internal keys are exported.
+        exact_keys : list, optional
+            List of keys to include in export with exact matching. The entire
+            key must match exactly. Takes precedence over export_keys.
 
         Returns
         -------
@@ -489,7 +495,7 @@ def export_data_to_csv(self, save_location, first_key=None):
         """
         from psPlotKit.data_manager.ps_data_exporter import PsDataExporter
 
-        exporter = PsDataExporter(self, save_location, first_key)
+        exporter = PsDataExporter(self, save_location, first_key, export_keys, exact_keys)
         return exporter.export()
 
     def display_loaded_contents(self):
diff --git a/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py b/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
index 58952af..18d876d 100644
--- a/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
+++ b/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
@@ -399,3 +399,195 @@ def test_first_key_via_manager_method(self, manual_single_dm, tmp_path):
             headers = next(reader)
 
         assert "metric_2" in headers[0]
+
+
+class TestNestedTupleProcessing:
+    def test_nested_tuple_keys_flattened_in_header(self, tmp_path):
+        """Nested tuple keys are flattened for multi-row headers."""
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            ("costing", ("stage 1", "pump"), "LCOW"),
+            PsData(("costing", ("stage 1", "pump"), "LCOW"), "test", [150.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            "simple_key",
+            PsData("Simple", "test", [50.0]),
+        )
+
+        save_path = str(tmp_path / "nested.csv")
+        exporter = PsDataExporter(dm, save_path)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header_rows = [next(reader) for _ in range(4)]
+
+        assert len(header_rows) == 4
+        assert header_rows[0] == ["costing", ""]
+        assert header_rows[1] == ["stage 1", ""]
+        assert header_rows[2] == ["pump", ""]
+        assert "LCOW" in header_rows[3][0] and "Simple" in header_rows[3][1]
+
+    def test_nested_tuple_key_matching_in_export_keys(self, tmp_path):
+        """export_keys matches elements within nested tuples."""
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            ("costing", ("stage 1", "pump"), "LCOW"),
+            PsData("pump LCOW", "test", [150.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            ("data", ("stage 2", "other"), "value"),
+            PsData("other value", "test", [200.0]),
+        )
+
+        save_path = str(tmp_path / "nested_filter.csv")
+        exporter = PsDataExporter(dm, save_path, export_keys=["stage 1"])
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header_rows = [next(reader) for _ in range(4)]
+
+        assert len(header_rows[3]) == 1
+        assert "LCOW" in header_rows[3][0]
+
+
+class TestKeyFiltering:
+    def test_export_keys_filters_by_partial_match(self, tmp_path):
+        """export_keys filters keys by partial match (tuple elements)."""
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "water_recovery",
+            PsData("Water recovery", "test", [0.85]),
+        )
+        dm.add_data(
+            "dir_a",
+            ("costing", "pump"),
+            PsData("pump", "test", [100.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            ("data_a", "data_b"),
+            PsData("data_b", "test", [200.0]),
+        )
+
+        save_path = str(tmp_path / "filtered.csv")
+        exporter = PsDataExporter(
+            dm, save_path, export_keys=["water_recovery", "costing"]
+        )
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header_rows = [next(reader) for _ in range(2)]
+
+        assert len(header_rows[1]) == 2
+        assert "Water recovery" in header_rows[1][0]
+        assert "pump" in header_rows[1][1]
+
+    def test_exact_keys_filters_by_exact_match(self, tmp_path):
+        """exact_keys filters keys by exact match only."""
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "water_recovery",
+            PsData("Water recovery", "test", [0.85]),
+        )
+        dm.add_data(
+            "dir_a",
+            ("costing", "pump"),
+            PsData("pump", "test", [100.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            ("data_a", "data_b"),
+            PsData("data_b", "test", [200.0]),
+        )
+
+        save_path = str(tmp_path / "exact.csv")
+        exporter = PsDataExporter(
+            dm, save_path, exact_keys=["water_recovery", ("costing", "pump")]
+        )
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header_rows = [next(reader) for _ in range(2)]
+
+        assert len(header_rows[1]) == 2
+        assert "Water recovery" in header_rows[1][0]
+        assert "pump" in header_rows[1][1]
+
+    def test_exact_keys_takes_precedence_over_export_keys(self, tmp_path):
+        """exact_keys takes precedence over export_keys."""
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "key_a",
+            PsData("Key A", "test", [1.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            ("costing", "pump"),
+            PsData("pump", "test", [100.0]),
+        )
+
+        save_path = str(tmp_path / "precedence.csv")
+        exporter = PsDataExporter(
+            dm, save_path, export_keys=["key_a"], exact_keys=[("costing", "pump")]
+        )
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header_rows = [next(reader) for _ in range(2)]
+
+        assert len(header_rows[1]) == 1
+        assert "pump" in header_rows[1][0]
+
+    def test_no_filter_exports_all_non_internal_keys(self, manual_single_dm, tmp_path):
+        """Without filter, all non-internal keys are exported."""
+        save_path = str(tmp_path / "no_filter.csv")
+        exporter = PsDataExporter(manual_single_dm, save_path)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert len(headers) == 2
+
+    def test_export_keys_via_manager_method(self, tmp_path):
+        """export_data_to_csv on PsDataManager accepts export_keys parameter."""
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            ("costing", "pump"),
+            PsData("pump", "test", [100.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            ("data_a", "data_b"),
+            PsData("data_b", "test", [200.0]),
+        )
+
+        save_path = str(tmp_path / "via_manager.csv")
+        dm.export_data_to_csv(save_path, export_keys=["costing"])
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header_rows = [next(reader) for _ in range(2)]
+
+        assert len(header_rows[1]) == 1
+        assert "pump" in header_rows[1][0]

From 769bf190b15804abb68a4a22ad0e1b037654547e Mon Sep 17 00:00:00 2001
From: avdudchenko <33663878+avdudchenko@users.noreply.github.com>
Date: Mon, 22 Jun 2026 19:35:21 -0400
Subject: [PATCH 3/5] adding export capability

---
 .../data_manager/ps_data_exporter.py          |  75 ++++++-
 src/psPlotKit/data_manager/ps_data_manager.py |   6 +-
 .../tests/test_ps_data_exporter.py            | 193 ++++++++++++++++++
 3 files changed, 269 insertions(+), 5 deletions(-)

diff --git a/src/psPlotKit/data_manager/ps_data_exporter.py b/src/psPlotKit/data_manager/ps_data_exporter.py
index 894119a..b572528 100644
--- a/src/psPlotKit/data_manager/ps_data_exporter.py
+++ b/src/psPlotKit/data_manager/ps_data_exporter.py
@@ -1,4 +1,5 @@
 import os
+import re
 import csv
 import numpy as np
 from psPlotKit.util.logger import define_logger
@@ -31,6 +32,9 @@ class PsDataExporter:
         List of keys to include in export with exact matching. The entire key
         (including all tuple elements) must match exactly. If specified, this
         takes precedence over export_keys.
+    skip_zero_data : bool, optional
+        If True, columns whose data are all zeros are excluded from export.
+        Defaults to False so existing behavior is preserved.
 
     Notes
     -----
@@ -41,7 +45,9 @@ class PsDataExporter:
       CSV file inside that folder.
     * Column headers are built from each :class:`PsData` object's
       ``data_label`` and ``mpl_units`` attributes (set by
-      :meth:`PsData.set_label`).
+      :meth:`PsData.set_label`).  LaTeX/math-mode markers are removed
+      from ``mpl_units`` so that headers contain plain text (e.g.
+      ``$/m^3`` instead of ``$\\$$/$m^3$``).
     * Internal keys (e.g., ``_zero_sentinel``) are automatically excluded
       from export.
     * For tuple data keys (e.g., ``("costing", "pump")``), multi-row headers
@@ -56,12 +62,14 @@ def __init__(
         first_key=None,
         export_keys=None,
         exact_keys=None,
+        skip_zero_data=False,
     ):
         self.ps_data_manager = ps_data_manager
         self.save_location = save_location
         self.first_key = first_key
         self.export_keys = export_keys
         self.exact_keys = exact_keys
+        self.skip_zero_data = skip_zero_data
 
     @staticmethod
     def _ensure_csv_extension(path):
@@ -77,6 +85,37 @@ def _strip_csv_extension(path):
             return path[:-4]
         return path
 
+    @staticmethod
+    def _clean_unit_for_export(units):
+        """Remove matplotlib/LaTeX math formatting from a unit string.
+
+        :class:`PsData.set_label` wraps unit tokens in math mode for
+        matplotlib rendering (e.g. ``$\\$$`` for a dollar sign,
+        ``$m^3$`` for an exponent).  CSV exports should contain plain
+        text such as ``$/m^3`` instead.
+
+        Parameters
+        ----------
+        units : str
+            The unit string as stored in ``PsData.mpl_units``.
+
+        Returns
+        -------
+        str
+            Unit string with LaTeX/math-mode markers removed.
+        """
+        if not units or units == "-":
+            return units
+
+        # Unwrap math-mode exponents first, e.g. $m^3$ -> m^3
+        units = re.sub(r"\$([^\$]*\^[^\$]*)\$", r"\1", units)
+        # USD/kUSD/MUSD math-mode dollar signs: $\$$ -> $, k$\$$ -> k$, etc.
+        units = re.sub(r"([kM]?)\$\s*\\\$\s*\$", r"\1$", units)
+        # Drop any remaining LaTeX escape backslashes
+        units = units.replace("\\", "")
+
+        return units
+
     def export(self):
         """Export the data manager contents to CSV.
 
@@ -133,12 +172,42 @@ def _group_data_by_directory(self):
             if not self._should_export_key(data_key):
                 continue
 
+            # Skip all-zero columns when requested
+            if self.skip_zero_data and self._is_zero_data(ps_data):
+                _logger.info(
+                    "Skipping '{}': all values are zero".format(data_key)
+                )
+                continue
+
             if dir_key not in grouped:
                 grouped[dir_key] = []
             grouped[dir_key].append((data_key, ps_data))
 
         return grouped
 
+    def _is_zero_data(self, ps_data):
+        """Return True if every value in *ps_data* is exactly zero.
+
+        Empty arrays are not considered all-zero.
+
+        Parameters
+        ----------
+        ps_data : PsData
+            The data object to inspect.
+
+        Returns
+        -------
+        bool
+            True if the data array is non-empty and all elements are 0.
+        """
+        data = ps_data.data
+        if data is None:
+            return False
+        arr = np.asarray(data)
+        if arr.size == 0:
+            return False
+        return np.all(arr == 0)
+
     def _should_export_key(self, data_key):
         """Check if a data key should be exported based on export_keys/exact_keys filters.
 
@@ -278,7 +347,7 @@ def _build_single_row_header(self, data_items):
         headers = []
         for _, ps_data in data_items:
             label = ps_data.data_label
-            units = getattr(ps_data, "mpl_units", "-")
+            units = self._clean_unit_for_export(getattr(ps_data, "mpl_units", "-"))
             if units and units != "-":
                 header = "{} ({})".format(label, units)
             else:
@@ -341,7 +410,7 @@ def _build_multi_row_header(self, data_items):
 
         for idx, (data_key, ps_data) in enumerate(data_items):
             label = ps_data.data_label
-            units = getattr(ps_data, "mpl_units", "-")
+            units = self._clean_unit_for_export(getattr(ps_data, "mpl_units", "-"))
 
             flat_key = flattened_keys[idx]
 
diff --git a/src/psPlotKit/data_manager/ps_data_manager.py b/src/psPlotKit/data_manager/ps_data_manager.py
index dc92d59..5e0cb0e 100644
--- a/src/psPlotKit/data_manager/ps_data_manager.py
+++ b/src/psPlotKit/data_manager/ps_data_manager.py
@@ -465,7 +465,7 @@ def get_data_keys(self, selected_keys):
             data_keys = self.data_keys[:]
         return data_keys
 
-    def export_data_to_csv(self, save_location, first_key=None, export_keys=None, exact_keys=None):
+    def export_data_to_csv(self, save_location, first_key=None, export_keys=None, exact_keys=None, skip_zero_data=False):
         """Export all loaded data to CSV file(s).
 
         Convenience wrapper around :class:`PsDataExporter`.  If this
@@ -487,6 +487,8 @@ def export_data_to_csv(self, save_location, first_key=None, export_keys=None, ex
         exact_keys : list, optional
             List of keys to include in export with exact matching. The entire
             key must match exactly. Takes precedence over export_keys.
+        skip_zero_data : bool, optional
+            If True, columns whose data are all zeros are excluded from export.
 
         Returns
         -------
@@ -495,7 +497,7 @@ def export_data_to_csv(self, save_location, first_key=None, export_keys=None, ex
         """
         from psPlotKit.data_manager.ps_data_exporter import PsDataExporter
 
-        exporter = PsDataExporter(self, save_location, first_key, export_keys, exact_keys)
+        exporter = PsDataExporter(self, save_location, first_key, export_keys, exact_keys, skip_zero_data)
         return exporter.export()
 
     def display_loaded_contents(self):
diff --git a/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py b/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
index 18d876d..bc704ae 100644
--- a/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
+++ b/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py
@@ -268,6 +268,199 @@ def test_dimensionless_units_no_parentheses(self, tmp_path):
         assert "(" not in headers[0]
 
 
+class TestUnitFormattingCleanup:
+    def test_latex_dollar_sign_exported_as_plain_dollar(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "LCOW",
+            PsData("LCOW", "test", [1.0, 2.0], import_units="USD/m**3"),
+        )
+        save_path = str(tmp_path / "lcow.csv")
+        exporter = PsDataExporter(dm, save_path)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header = next(reader)[0]
+
+        assert "LCOW" in header
+        assert "$/m^3" in header
+        assert "$\\$$" not in header
+        assert "${" not in header
+
+    def test_latex_cleanup_for_kusd_and_musd(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "capex_k",
+            PsData("capex_k", "test", [1.0], import_units="kUSD/m**3"),
+        )
+        dm.add_data(
+            "dir_a",
+            "capex_M",
+            PsData("capex_M", "test", [1.0], import_units="MUSD/m**3"),
+        )
+        save_path = str(tmp_path / "capex.csv")
+        exporter = PsDataExporter(dm, save_path)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert "k$/m^3" in headers[0]
+        assert "M$/m^3" in headers[1]
+
+    def test_plain_units_passthrough_unchanged(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "metric_1",
+            PsData("metric_1", "test", [1.0, 2.0, 3.0], import_units="m"),
+        )
+        save_path = str(tmp_path / "plain.csv")
+        exporter = PsDataExporter(dm, save_path)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header = next(reader)[0]
+
+        assert "metric_1 (m)" == header
+
+    def test_cleanup_multi_row_header_for_tuple_keys(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            ("costing", "LCOW"),
+            PsData("LCOW", "test", [1.0], import_units="USD/m**3"),
+        )
+        save_path = str(tmp_path / "multirow.csv")
+        exporter = PsDataExporter(dm, save_path)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            header_rows = [next(reader) for _ in range(2)]
+
+        assert header_rows[0] == ["costing"]
+        assert "LCOW" in header_rows[1][0]
+        assert "$/m^3" in header_rows[1][0]
+        assert "$\\$$" not in header_rows[1][0]
+
+
+class TestSkipZeroData:
+    def test_zero_column_skipped_when_enabled(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "nonzero",
+            PsData("nonzero", "test", [1.0, 2.0, 3.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            "all_zero",
+            PsData("all_zero", "test", [0.0, 0.0, 0.0]),
+        )
+        save_path = str(tmp_path / "skip_zero.csv")
+        exporter = PsDataExporter(dm, save_path, skip_zero_data=True)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert len(headers) == 1
+        assert "nonzero" in headers[0]
+        assert "all_zero" not in headers[0]
+
+    def test_zero_column_kept_when_disabled(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "nonzero",
+            PsData("nonzero", "test", [1.0, 2.0, 3.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            "all_zero",
+            PsData("all_zero", "test", [0.0, 0.0, 0.0]),
+        )
+        save_path = str(tmp_path / "keep_zero.csv")
+        exporter = PsDataExporter(dm, save_path, skip_zero_data=False)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert len(headers) == 2
+        assert "all_zero" in headers[1]
+
+    def test_single_zero_value_skipped(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "scalar_zero",
+            PsData("scalar_zero", "test", [0.0]),
+        )
+        save_path = str(tmp_path / "scalar_zero.csv")
+        exporter = PsDataExporter(dm, save_path, skip_zero_data=True)
+        written = exporter.export()
+
+        assert written == []
+        assert not os.path.exists(save_path)
+
+    def test_mixed_zero_and_nonzero_keeps_nonzero(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "mixed",
+            PsData("mixed", "test", [0.0, 1.0, 0.0]),
+        )
+        save_path = str(tmp_path / "mixed.csv")
+        exporter = PsDataExporter(dm, save_path, skip_zero_data=True)
+        exporter.export()
+
+        with open(save_path, "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert "mixed" in headers[0]
+
+    def test_skip_zero_data_via_manager_method(self, tmp_path):
+        dm = PsDataManager()
+        dm.PsDataImportInstances = []
+        dm.add_data(
+            "dir_a",
+            "nonzero",
+            PsData("nonzero", "test", [1.0, 2.0]),
+        )
+        dm.add_data(
+            "dir_a",
+            "all_zero",
+            PsData("all_zero", "test", [0.0, 0.0]),
+        )
+        save_path = str(tmp_path / "manager_skip_zero.csv")
+        written = dm.export_data_to_csv(save_path, skip_zero_data=True)
+
+        with open(written[0], "r") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+
+        assert len(headers) == 1
+        assert "nonzero" in headers[0]
+
+
 # ---------- PsDataManager.export_data_to_csv convenience method ----------
 
 

From cc0e63f1ada5320ef2236e916898d908b2256e22 Mon Sep 17 00:00:00 2001
From: avdudchenko <33663878+avdudchenko@users.noreply.github.com>
Date: Tue, 23 Jun 2026 18:58:31 -0400
Subject: [PATCH 4/5] fixing exports

---
 src/psPlotKit/data_manager/ps_data.py         | 78 +++++++++++++++----
 .../data_manager/ps_data_exporter.py          | 38 +++++----
 .../data_manager/tests/test_ps_data.py        | 53 +++++++++++++
 3 files changed, 141 insertions(+), 28 deletions(-)

diff --git a/src/psPlotKit/data_manager/ps_data.py b/src/psPlotKit/data_manager/ps_data.py
index c736485..12a1f07 100644
--- a/src/psPlotKit/data_manager/ps_data.py
+++ b/src/psPlotKit/data_manager/ps_data.py
@@ -169,13 +169,44 @@ def _define_custom_units(self, custom_units=None):
             custom_units = CustomUnits()
         self.custom_units = custom_units.get_units_dict()
 
+    def _raise_unit_error(
+        self, error, operation, requested_units=None, other=None
+    ):
+        """Build and raise an enhanced unit error preserving the original type.
+
+        The raised exception keeps the same Python type as *error* so existing
+        ``except ValueError`` / ``except LookupError`` blocks continue to work,
+        but the message now includes the data_key and data_type so users can
+        identify which data triggered the incompatible unit problem.
+        """
+        msg = "Unit error for data_key={!r}, data_type={!r}. Operation: {}. ".format(
+            self.data_key, self.data_type, operation
+        )
+        if requested_units is not None:
+            msg += "Requested units: {!r}. ".format(requested_units)
+        if other is not None:
+            if isinstance(other, PsData):
+                msg += "Other operand: data_key={!r}, data_type={!r}. ".format(
+                    other.data_key, other.data_type
+                )
+            else:
+                msg += "Other operand: {!r}. ".format(other)
+        msg += "Original error: {}: {}".format(type(error).__name__, error)
+        _logger.error(msg)
+        raise type(error)(msg) from error
+
     def _assign_units(self, manual_conversion=1):
         if self.data_is_numbers:
             self.data = self.data * manual_conversion
         qsunits = self._get_qs_unit()
-        self.data_with_units = qs.Quantity(self.data.copy(), qsunits)
-        self.raw_data_with_units = qs.Quantity(self.raw_data.copy(), qsunits)
-        self.data = self.data_with_units.magnitude
+        try:
+            self.data_with_units = qs.Quantity(self.data.copy(), qsunits)
+            self.raw_data_with_units = qs.Quantity(self.raw_data.copy(), qsunits)
+            self.data = self.data_with_units.magnitude
+        except (ValueError, LookupError) as e:
+            self._raise_unit_error(
+                e, "assign units", requested_units=self.sunits
+            )
         self.set_label()
 
     @property
@@ -280,15 +311,23 @@ def set_data(self, data):
         self.raw_data = data
 
     def to_units(self, new_units):
-        self.sunits = self._convert_string_unit(new_units)
+        converted_units = self._convert_string_unit(new_units)
+        old_sunits = self.sunits
+        self.sunits = converted_units
         qsunits = self._get_qs_unit()
-        if new_units == "degC" and str(self.data_with_units.units) == "1.0 K":
-            self.data_with_units = qs.Quantity(
-                self.data_with_units.magnitude[:] - 273.15, new_units
+        try:
+            if new_units == "degC" and str(self.data_with_units.units) == "1.0 K":
+                self.data_with_units = qs.Quantity(
+                    self.data_with_units.magnitude[:] - 273.15, new_units
+                )
+            else:
+                self.data_with_units = self.data_with_units.rescale(qsunits)
+            self.raw_data_with_units = self.raw_data_with_units.rescale(qsunits)
+        except (ValueError, LookupError) as e:
+            self.sunits = old_sunits
+            self._raise_unit_error(
+                e, "convert units", requested_units=new_units
             )
-        else:
-            self.data_with_units = self.data_with_units.rescale(qsunits)
-        self.raw_data_with_units = self.raw_data_with_units.rescale(qsunits)
         self.data = self.data_with_units.magnitude[:]
         self.raw_data = self.raw_data_with_units.magnitude
         self.set_label()
@@ -349,7 +388,12 @@ def _arithmetic_op(self, other, op, symbol):
                 "Arithmetic operations require a PsData object or numeric "
                 "scalar, got {}".format(type(other))
             )
-        result_quantity = op(self.data_with_units, other_val)
+        try:
+            result_quantity = op(self.data_with_units, other_val)
+        except (ValueError, LookupError) as e:
+            self._raise_unit_error(
+                e, "arithmetic '{}'".format(symbol), other=other
+            )
         result_key = "({} {} {})".format(self.data_key, symbol, other_label)
         return PsData(
             data_key=result_key,
@@ -371,7 +415,12 @@ def _r_arithmetic_op(self, other, op, symbol):
                 "Arithmetic operations require a PsData object or numeric "
                 "scalar, got {}".format(type(other))
             )
-        result_quantity = op(other_val, self.data_with_units)
+        try:
+            result_quantity = op(other_val, self.data_with_units)
+        except (ValueError, LookupError) as e:
+            self._raise_unit_error(
+                e, "reflected arithmetic '{}'".format(symbol), other=other
+            )
         result_key = "({} {} {})".format(other_label, symbol, self.data_key)
         return PsData(
             data_key=result_key,
@@ -426,7 +475,10 @@ def __rpow__(self, other):
         )
 
     def __neg__(self):
-        result_quantity = -1 * self.data_with_units
+        try:
+            result_quantity = -1 * self.data_with_units
+        except (ValueError, LookupError) as e:
+            self._raise_unit_error(e, "negation")
         result_key = "(-{})".format(self.data_key)
         return PsData(
             data_key=result_key,
diff --git a/src/psPlotKit/data_manager/ps_data_exporter.py b/src/psPlotKit/data_manager/ps_data_exporter.py
index b572528..3f74864 100644
--- a/src/psPlotKit/data_manager/ps_data_exporter.py
+++ b/src/psPlotKit/data_manager/ps_data_exporter.py
@@ -174,21 +174,18 @@ def _group_data_by_directory(self):
 
             # Skip all-zero columns when requested
             if self.skip_zero_data and self._is_zero_data(ps_data):
-                _logger.info(
-                    "Skipping '{}': all values are zero".format(data_key)
-                )
+                _logger.info("Skipping '{}': all values are zero".format(data_key))
                 continue
 
             if dir_key not in grouped:
                 grouped[dir_key] = []
             grouped[dir_key].append((data_key, ps_data))
-
         return grouped
 
     def _is_zero_data(self, ps_data):
         """Return True if every value in *ps_data* is exactly zero.
 
-        Empty arrays are not considered all-zero.
+        Empty arrays are considered all-zero.
 
         Parameters
         ----------
@@ -198,15 +195,22 @@ def _is_zero_data(self, ps_data):
         Returns
         -------
         bool
-            True if the data array is non-empty and all elements are 0.
+            True if the data array is empty or all elements are 0.
         """
         data = ps_data.data
         if data is None:
-            return False
+            return True
         arr = np.asarray(data)
+
         if arr.size == 0:
-            return False
-        return np.all(arr == 0)
+            return True
+        arr = np.nan_to_num(arr, nan=0)
+        if np.sum(np.abs(arr)) < np.finfo(float).eps * 10 or np.isnan(
+            np.sum(np.abs(arr))
+        ):
+            return True
+        return False
+        # return np.all(arr == 0
 
     def _should_export_key(self, data_key):
         """Check if a data key should be exported based on export_keys/exact_keys filters.
@@ -374,7 +378,7 @@ def _flatten_key(data_key):
         """
         flattened = []
         for elem in data_key:
-            if isinstance(elem, tuple):
+            if isinstance(elem, (tuple, list)):
                 flattened.extend(elem)
             else:
                 flattened.append(elem)
@@ -397,7 +401,7 @@ def _build_multi_row_header(self, data_items):
         # Flatten nested tuples to determine max depth
         flattened_keys = []
         for key, _ in data_items:
-            if isinstance(key, tuple):
+            if isinstance(key, (tuple, list)):
                 flattened_keys.append(self._flatten_key(key))
             else:
                 flattened_keys.append(key)
@@ -415,15 +419,19 @@ def _build_multi_row_header(self, data_items):
             flat_key = flattened_keys[idx]
 
             # If label equals data_key, use the last tuple element as label
-            if label == data_key and isinstance(data_key, tuple):
+            if label == data_key and isinstance(data_key, (tuple, list)):
                 label = flat_key[-1]
-
+            if isinstance(label, (tuple, list)):
+                if len(label) == 1:
+                    label = str(label[0])
+                else:
+                    label = ", ".join(label)
             if units and units != "-":
+
                 full_label = "{} ({})".format(label, units)
             else:
                 full_label = str(label)
-
-            if isinstance(flat_key, tuple):
+            if isinstance(flat_key, (tuple, list)):
                 key_len = len(flat_key)
                 # Calculate starting row (align shorter tuples to bottom)
                 start_row = max_depth - key_len
diff --git a/src/psPlotKit/data_manager/tests/test_ps_data.py b/src/psPlotKit/data_manager/tests/test_ps_data.py
index b5e91c9..7a1e493 100644
--- a/src/psPlotKit/data_manager/tests/test_ps_data.py
+++ b/src/psPlotKit/data_manager/tests/test_ps_data.py
@@ -382,3 +382,56 @@ def test_scalar_sub_and_div(self):
         a = PsData("a", "test", [10.0, 20.0, 30.0])
         result = (a - 5) / 2
         np.testing.assert_array_almost_equal(result.data, [2.5, 7.5, 12.5])
+
+
+# ---------- unit error handling ----------
+
+
+class TestUnitErrors:
+    def test_to_units_incompatible_raises_with_metadata(self, meter_ps):
+        with pytest.raises(ValueError) as excinfo:
+            meter_ps.to_units("kg")
+        msg = str(excinfo.value)
+        assert "length" in msg
+        assert "test" in msg
+        assert "kg" in msg
+        assert "m" in msg
+        assert excinfo.value.__cause__ is not None
+
+    def test_to_units_preserves_original_units_on_failure(self, meter_ps):
+        with pytest.raises(ValueError):
+            meter_ps.to_units("kg")
+        assert meter_ps.sunits == "m"
+        np.testing.assert_array_almost_equal(meter_ps.data, [10.0, 20.0, 30.0])
+
+    def test_constructor_units_incompatible_raises_with_metadata(self):
+        with pytest.raises(ValueError) as excinfo:
+            PsData("length", "test", [1.0, 2.0], import_units="m", units="kg")
+        msg = str(excinfo.value)
+        assert "length" in msg
+        assert "test" in msg
+        assert "kg" in msg
+
+    def test_assign_units_invalid_string_raises_with_metadata(self, meter_ps):
+        with pytest.raises(LookupError) as excinfo:
+            meter_ps.assign_units("totally_bogus_unit")
+        msg = str(excinfo.value)
+        assert "length" in msg
+        assert "test" in msg
+        assert "totally_bogus_unit" in msg
+
+    def test_invalid_unit_string_raises_with_metadata(self):
+        with pytest.raises(LookupError) as excinfo:
+            PsData("x", "test", [1.0], import_units="totally_bogus_unit")
+        msg = str(excinfo.value)
+        assert "x" in msg
+        assert "test" in msg
+        assert "totally_bogus_unit" in msg
+
+    def test_arithmetic_incompatible_units_raises_with_metadata(self, meter_ps, usd_ps):
+        with pytest.raises(ValueError) as excinfo:
+            meter_ps + usd_ps
+        msg = str(excinfo.value)
+        assert "length" in msg
+        assert "cost" in msg
+        assert "test" in msg

From 2cbaec08ed7c3d209827b4e0f3ab5f5b35db756e Mon Sep 17 00:00:00 2001
From: avdudchenko <33663878+avdudchenko@users.noreply.github.com>
Date: Tue, 23 Jun 2026 20:02:52 -0400
Subject: [PATCH 5/5] black formatting

---
 src/psPlotKit/data_manager/ps_data.py         | 16 ++++------------
 src/psPlotKit/data_manager/ps_data_manager.py | 13 +++++++++++--
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/psPlotKit/data_manager/ps_data.py b/src/psPlotKit/data_manager/ps_data.py
index 12a1f07..a627a9e 100644
--- a/src/psPlotKit/data_manager/ps_data.py
+++ b/src/psPlotKit/data_manager/ps_data.py
@@ -169,9 +169,7 @@ def _define_custom_units(self, custom_units=None):
             custom_units = CustomUnits()
         self.custom_units = custom_units.get_units_dict()
 
-    def _raise_unit_error(
-        self, error, operation, requested_units=None, other=None
-    ):
+    def _raise_unit_error(self, error, operation, requested_units=None, other=None):
         """Build and raise an enhanced unit error preserving the original type.
 
         The raised exception keeps the same Python type as *error* so existing
@@ -204,9 +202,7 @@ def _assign_units(self, manual_conversion=1):
             self.raw_data_with_units = qs.Quantity(self.raw_data.copy(), qsunits)
             self.data = self.data_with_units.magnitude
         except (ValueError, LookupError) as e:
-            self._raise_unit_error(
-                e, "assign units", requested_units=self.sunits
-            )
+            self._raise_unit_error(e, "assign units", requested_units=self.sunits)
         self.set_label()
 
     @property
@@ -325,9 +321,7 @@ def to_units(self, new_units):
             self.raw_data_with_units = self.raw_data_with_units.rescale(qsunits)
         except (ValueError, LookupError) as e:
             self.sunits = old_sunits
-            self._raise_unit_error(
-                e, "convert units", requested_units=new_units
-            )
+            self._raise_unit_error(e, "convert units", requested_units=new_units)
         self.data = self.data_with_units.magnitude[:]
         self.raw_data = self.raw_data_with_units.magnitude
         self.set_label()
@@ -391,9 +385,7 @@ def _arithmetic_op(self, other, op, symbol):
         try:
             result_quantity = op(self.data_with_units, other_val)
         except (ValueError, LookupError) as e:
-            self._raise_unit_error(
-                e, "arithmetic '{}'".format(symbol), other=other
-            )
+            self._raise_unit_error(e, "arithmetic '{}'".format(symbol), other=other)
         result_key = "({} {} {})".format(self.data_key, symbol, other_label)
         return PsData(
             data_key=result_key,
diff --git a/src/psPlotKit/data_manager/ps_data_manager.py b/src/psPlotKit/data_manager/ps_data_manager.py
index 5e0cb0e..886720d 100644
--- a/src/psPlotKit/data_manager/ps_data_manager.py
+++ b/src/psPlotKit/data_manager/ps_data_manager.py
@@ -465,7 +465,14 @@ def get_data_keys(self, selected_keys):
             data_keys = self.data_keys[:]
         return data_keys
 
-    def export_data_to_csv(self, save_location, first_key=None, export_keys=None, exact_keys=None, skip_zero_data=False):
+    def export_data_to_csv(
+        self,
+        save_location,
+        first_key=None,
+        export_keys=None,
+        exact_keys=None,
+        skip_zero_data=False,
+    ):
         """Export all loaded data to CSV file(s).
 
         Convenience wrapper around :class:`PsDataExporter`.  If this
@@ -497,7 +504,9 @@ def export_data_to_csv(self, save_location, first_key=None, export_keys=None, ex
         """
         from psPlotKit.data_manager.ps_data_exporter import PsDataExporter
 
-        exporter = PsDataExporter(self, save_location, first_key, export_keys, exact_keys, skip_zero_data)
+        exporter = PsDataExporter(
+            self, save_location, first_key, export_keys, exact_keys, skip_zero_data
+        )
         return exporter.export()
 
     def display_loaded_contents(self):