From c323795fd67f8a57cd15aad8cb56d0d4a9490a89 Mon Sep 17 00:00:00 2001
From: domfournier <dominiquef@mirageoscience.com>
Date: Mon, 14 Apr 2025 15:19:48 -0700
Subject: [PATCH] Add directive for saving LP models.

---
 simpeg/directives/__init__.py    |  1 +
 simpeg/directives/_save_geoh5.py | 67 +++++++++++++++++++++++++-------
 2 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/simpeg/directives/__init__.py b/simpeg/directives/__init__.py
index 737425f76f..4581e48d29 100644
--- a/simpeg/directives/__init__.py
+++ b/simpeg/directives/__init__.py
@@ -126,6 +126,7 @@
     BaseSaveGeoH5,
     SaveDataGeoH5,
     SaveLogFilesGeoH5,
+    SaveLPModelGroup,
     SaveModelGeoH5,
     SavePropertyGroup,
     SaveSensitivityGeoH5,
diff --git a/simpeg/directives/_save_geoh5.py b/simpeg/directives/_save_geoh5.py
index 89c4575442..6d085c7478 100644
--- a/simpeg/directives/_save_geoh5.py
+++ b/simpeg/directives/_save_geoh5.py
@@ -8,8 +8,9 @@
 from .directives import InversionDirective
 from simpeg.maps import IdentityMap
 
+from geoh5py.data import NumericData
 from geoh5py.groups.property_group import GroupTypeEnum
-from geoh5py.groups import PropertyGroup, UIJsonGroup
+from geoh5py.groups import UIJsonGroup
 from geoh5py.objects import ObjectBase
 from geoh5py.ui_json.utils import fetch_active_workspace
 
@@ -419,7 +420,7 @@ def save_log(self):
 
 class SavePropertyGroup(BaseSaveGeoH5):
     """
-    Save the model as a property group in the geoh5 file
+    Assign the data to a property group in the geoh5 file
     """
 
     def __init__(
@@ -446,21 +447,61 @@ def write(self, iteration: int, **_):
                     channel_name, base_name = self.get_names(
                         component, channel, iteration
                     )
-                    child = [
+                    children = [
                         child
                         for child in h5_object.children
-                        if channel_name in child.name
-                    ][0]
+                        if (
+                            channel_name in child.name
+                            and isinstance(child, NumericData)
+                        )
+                    ]
 
-                    if child is not None:
-                        properties.append(child)
+                    if children[0] is not None:
+                        properties += children
 
                 if len(properties) == 0:
                     return
 
-                PropertyGroup(
-                    parent=h5_object,
-                    name=base_name,
-                    properties=properties,
-                    property_group_type=self.group_type,
-                )
+                prop_group = h5_object.get_property_group(base_name)[0]
+
+                if prop_group is None:
+                    prop_group = h5_object.create_property_group(
+                        name=base_name,
+                        properties=properties,
+                        property_group_type=self.group_type,
+                    )
+                else:
+                    prop_group.add_properties(properties)
+
+
+class SaveLPModelGroup(SavePropertyGroup):
+    """
+    Save the model as a property group in the geoh5 file
+    """
+
+    def __init__(
+        self,
+        h5_object,
+        irls_directive,
+        group_type: GroupTypeEnum = GroupTypeEnum.MULTI,
+        **kwargs,
+    ):
+        self.group_type = group_type
+        self.irls_directive = irls_directive
+
+        super().__init__(h5_object, **kwargs)
+
+    def get_names(
+        self, component: str, channel: str, iteration: int
+    ) -> tuple[str, str]:
+        """
+        Format the data and property_group name.
+        """
+        channel_name, base_name = super().get_names(component, channel, iteration)
+
+        if self.irls_directive.metrics.irls_iteration_count == 0:
+            base_name = "L2 models"
+        else:
+            base_name = "LP models"
+
+        return channel_name, base_name