GEOPY-1965: Convert all EM params to BaseData

simpeg_drivers/__init__.py

@@ -91,27 +91,36 @@ def assets_path() -> Path:
     ),
     "fem": (
         "simpeg_drivers.electromagnetics.frequency_domain.driver",
-        {"inversion": "FrequencyDomainElectromagneticsDriver"},
+        {
+            "forward": "FrequenceyDomainElectromagneticsForwardDriver",
+            "inversion": "FrequencyDomainElectromagneticsInversionDriver",
+        },
     ),
     "joint cross gradient": (
         "simpeg_drivers.joint.joint_cross_gradient.driver",
         {"inversion": "JointCrossGradientDriver"},
     ),
     "tdem": (
         "simpeg_drivers.electromagnetics.time_domain.driver",
-        {"inversion": "TimeDomainElectromagneticsDriver"},
+        {
+            "forward": "TimeDomainElectromagneticsForwardDriver",
+            "inversion": "TimeDomainElectromagneticsInversionDriver",
+        },
     ),
     "magnetotellurics": (
         "simpeg_drivers.natural_sources.magnetotellurics.driver",
-        {"inversion": "MagnetotelluricsDriver"},
+        {
+            "forward": "MagnetotelluricsForwardDriver",
+            "inversion": "MagnetotelluricsInversionDriver",
+        },
     ),
     "tipper": (
         "simpeg_drivers.natural_sources.tipper.driver",
-        {"inversion": "TipperDriver"},
+        {"forward": "TipperForwardDriver", "inversion": "TipperInversionDriver"},
     ),
     "gravity": (
         "simpeg_drivers.potential_fields.gravity.driver",
-        {"inversion": "GravityInversionDriver", "forward": "GravityForwardDriver"},
+        {"forward": "GravityForwardDriver", "inversion": "GravityInversionDriver"},
     ),
     "magnetic scalar": (
         "simpeg_drivers.potential_fields.magnetic_scalar.driver",

simpeg_drivers/components/data.py

@@ -89,7 +89,6 @@ def __init__(self, workspace: Workspace, params: InversionBaseParams):
         self.indices: np.ndarray | None = None
         self.vector: bool | None = None
         self.n_blocks: int | None = None
-        self.components: list[str] | None = None
         self.observed: dict[str, np.ndarray] = {}
         self.predicted: dict[str, np.ndarray] = {}
         self.uncertainties: dict[str, np.ndarray] = {}
@@ -106,8 +105,10 @@ def _initialize(self) -> None:
         """Extract data from the workspace using params data."""
         self.vector = True if self.params.inversion_type == "magnetic vector" else False
         self.n_blocks = 3 if self.params.inversion_type == "magnetic vector" else 1
-        self.components, self.observed, self.uncertainties = self.get_data()
-        self.has_tensor = InversionData.check_tensor(self.components)
+        self.components = self.params.active_components
+        self.observed = self.params.data
+        self.uncertainties = self.params.uncertainties
+        self.has_tensor = InversionData.check_tensor(self.params.components)
         self.locations = super().get_locations(self.params.data_object)
 
         if "2d" in self.params.inversion_type:
@@ -280,7 +281,7 @@ def normalize(
         """
         d = deepcopy(data)
         for chan in getattr(self.params.data_object, "channels", [None]):
-            for comp in self.components:
+            for comp in self.params.active_components:
                 if isinstance(d[comp], dict):
                     if d[comp][chan] is not None:
                         d[comp][chan] *= self.normalizations[chan][comp]
@@ -298,7 +299,7 @@ def get_normalizations(self):
         normalizations = {}
         for chan in getattr(self.params.data_object, "channels", [None]):
             normalizations[chan] = {}
-            for comp in self.components:
+            for comp in self.params.active_components:
                 normalizations[chan][comp] = np.ones(self.mask.sum())
                 if comp in ["potential", "chargeability"]:
                     normalizations[chan][comp] = 1
@@ -488,7 +489,7 @@ def survey(self):
     @property
     def n_data(self):
         n_data = 0
-        for comp in self.components:
+        for comp in self.params.active_components:
             if isinstance(self.observed[comp], dict):
                 for channel in self.observed[comp]:
                     n_data += len(self.observed[comp][channel])

simpeg_drivers/components/locations.py

@@ -121,6 +121,8 @@ def get_locations(self, entity: ObjectBase) -> np.ndarray:
         return locations
 
     def _filter(self, a, mask):
+        if a is None:
+            return None
         for k, v in a.items():
             if not isinstance(v, np.ndarray):
                 a.update({k: self._filter(v, mask)})

simpeg_drivers/electromagnetics/frequency_domain/__init__.py

@@ -9,7 +9,10 @@
 # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
 
 
-from .params import FrequencyDomainElectromagneticsParams
+from .params import (
+    FrequencyDomainElectromagneticsForwardParams,
+    FrequencyDomainElectromagneticsInversionParams,
+)
 
 # pylint: disable=unused-import
 # flake8: noqa

simpeg_drivers/electromagnetics/frequency_domain/driver.py

@@ -14,12 +14,20 @@
 from simpeg_drivers.driver import InversionDriver
 
 from .constants import validations
-from .params import FrequencyDomainElectromagneticsParams
+from .params import (
+    FrequencyDomainElectromagneticsForwardParams,
+    FrequencyDomainElectromagneticsInversionParams,
+)
 
 
-class FrequencyDomainElectromagneticsDriver(InversionDriver):
-    _params_class = FrequencyDomainElectromagneticsParams
+class FrequencyDomainElectromagneticsForwardDriver(InversionDriver):
+    _params_class = FrequencyDomainElectromagneticsForwardParams
     _validations = validations
 
-    def __init__(self, params: FrequencyDomainElectromagneticsParams):
+    def __init__(self, params: FrequencyDomainElectromagneticsForwardParams):
         super().__init__(params)
+
+
+class FrequencyDomainElectromagneticsInversionDriver(InversionDriver):
+    _params_class = FrequencyDomainElectromagneticsInversionParams
+    _validations = validations

simpeg_drivers/electromagnetics/frequency_domain/params.py

@@ -11,68 +11,61 @@
 
 from __future__ import annotations
 
-from copy import deepcopy
-from uuid import UUID
+from pathlib import Path
+from typing import ClassVar, TypeAlias
+
+from geoh5py.groups import PropertyGroup
+from geoh5py.objects import (
+    AirborneFEMReceivers,
+    LargeLoopGroundFEMReceivers,
+    MovingLoopGroundFEMReceivers,
+)
+
+from simpeg_drivers import assets_path
+from simpeg_drivers.params import BaseForwardData, BaseInversionData, EMDataMixin
 
-from simpeg_drivers.params import InversionBaseParams
 
-from .constants import (
-    default_ui_json,
-    forward_defaults,
-    inversion_defaults,
-    validations,
+Receivers: TypeAlias = (
+    MovingLoopGroundFEMReceivers | LargeLoopGroundFEMReceivers | AirborneFEMReceivers
 )
 
 
-class FrequencyDomainElectromagneticsParams(InversionBaseParams):
+class FrequencyDomainElectromagneticsForwardParams(EMDataMixin, BaseForwardData):
     """
-    Parameter class for Frequency-domain Electromagnetic (FEM) -> conductivity inversion.
+    Parameter class for Frequency-domain Electromagnetic (FEM) simulation.
+
+    :param z_real_channel_bool: Real impedance channel boolean.
+    :param z_imag_channel_bool: Imaginary impedance channel boolean.
+    :param model_type: Specify whether the models are provided in resistivity or conductivity.
     """
 
-    _physical_property = "conductivity"
-
-    def __init__(self, input_file=None, forward_only=False, **kwargs):
-        self._default_ui_json = deepcopy(default_ui_json)
-        self._forward_defaults = deepcopy(forward_defaults)
-        self._inversion_defaults = deepcopy(inversion_defaults)
-        self._inversion_type = "fem"
-        self._validations = validations
-        self._tx_offsets = None
-        self._z_real_channel_bool = None
-        self._z_real_channel = None
-        self._z_real_uncertainty = None
-        self._z_imag_channel_bool = None
-        self._z_imag_channel = None
-        self._z_imag_uncertainty = None
-        self._model_type = "Conductivity (S/m)"
-
-        super().__init__(input_file=input_file, forward_only=forward_only, **kwargs)
-
-    def data_channel(self, component: str):
-        """Return uuid of data channel."""
-        return getattr(self, "_".join([component, "channel"]), None)
-
-    def uncertainty_channel(self, component: str):
-        """Return uuid of uncertainty channel."""
-        return getattr(self, "_".join([component, "uncertainty"]), None)
-
-    def property_group_data(self, property_group: UUID):
-        """
-        Return dictionary of channel/data.
-
-        :param property_group: Property group uid
-        """
-        channels = self.data_object.channels
-        if self.forward_only:
-            out = {k: None for k in channels}
-        else:
-            group = self.data_object.find_or_create_property_group(
-                name=property_group.name
-            )
-            properties = [self.geoh5.get_entity(p)[0].values for p in group.properties]
-            out = {f: properties[i] for i, f in enumerate(channels)}
-
-        return out
+    name: ClassVar[str] = "Frequency Domain Electromagnetics Forward"
+    title: ClassVar[str] = "Frequency-domain EM (FEM) Forward"
+    default_ui_json: ClassVar[Path] = assets_path() / "uijson/fem_forward.ui.json"
+
+    inversion_type: str = "fem"
+    physical_property: str = "conductivity"
+
+    data_object: Receivers
+    z_real_channel_bool: bool
+    z_imag_channel_bool: bool
+    model_type: str = "Conductivity (S/m)"
+
+    @property
+    def tx_offsets(self):
+        """Return transmitter offsets from frequency metadata"""
+
+        try:
+            offset_data = self.data_object.metadata["EM Dataset"][
+                "Frequency configurations"
+            ]
+            tx_offsets = {k["Frequency"]: k["Offset"] for k in offset_data}
+
+        except KeyError as exception:
+            msg = "Metadata must contain 'Frequency configurations' dictionary with 'Offset' data."
+            raise KeyError(msg) from exception
+
+        return tx_offsets
 
     @property
     def unit_conversion(self):
@@ -82,83 +75,52 @@ def unit_conversion(self):
         }
         return conversion[self.data_object.unit]
 
-    def data(self, component: str):
-        """Returns array of data for chosen data component."""
-        property_group = self.data_channel(component)
-        return self.property_group_data(property_group)
 
-    def uncertainty(self, component: str) -> float:
-        """Returns uncertainty for chosen data component."""
-        uid = self.uncertainty_channel(component)
-        return self.property_group_data(uid)
-
-    @property
-    def model_type(self):
-        """Model units."""
-        return self._model_type
-
-    @model_type.setter
-    def model_type(self, val):
-        self.setter_validator("model_type", val)
-
-    @property
-    def tx_offsets(self):
-        if self._tx_offsets is None and self.data_object is not None:
-            try:
-                offset_data = self.data_object.metadata["EM Dataset"][
-                    "Frequency configurations"
-                ]
-                self._tx_offsets = {k["Frequency"]: k["Offset"] for k in offset_data}
-            except KeyError as exception:
-                msg = "Metadata must contain 'Frequency configurations' dictionary with 'Offset' data."
-                raise KeyError(msg) from exception
-
-        return self._tx_offsets
+class FrequencyDomainElectromagneticsInversionParams(EMDataMixin, BaseInversionData):
+    """
+    Parameter class for Frequency-domain Electromagnetic (FEM) -> conductivity inversion.
 
-    @property
-    def z_real_channel_bool(self):
-        return self._z_real_channel_bool
+    :param z_real_channel: Real impedance channel.
+    :param z_real_uncertainty: Real impedance uncertainty channel.
+    :param z_imag_channel: Imaginary impedance channel.
+    :param z_imag_uncertainty: Imaginary impedance uncertainty channel.
+    :param model_type: Specify whether the models are provided in resistivity or conductivity.
+    """
 
-    @z_real_channel_bool.setter
-    def z_real_channel_bool(self, val):
-        self.setter_validator("z_real_channel_bool", val)
+    name: ClassVar[str] = "Frequency Domain Electromagnetics Inversion"
+    title: ClassVar[str] = "Frequency-domain EM (FEM) Inversion"
+    default_ui_json: ClassVar[Path] = assets_path() / "uijson/fem_inversion.ui.json"
 
-    @property
-    def z_real_channel(self):
-        return self._z_real_channel
+    inversion_type: str = "fem"
+    physical_property: str = "conductivity"
 
-    @z_real_channel.setter
-    def z_real_channel(self, val):
-        self.setter_validator("z_real_channel", val, fun=self._uuid_promoter)
+    data_object: Receivers
+    z_real_channel: PropertyGroup | None = None
+    z_real_uncertainty: PropertyGroup | None = None
+    z_imag_channel: PropertyGroup | None = None
+    z_imag_uncertainty: PropertyGroup | None = None
+    model_type: str = "Conductivity (S/m)"
 
     @property
-    def z_real_uncertainty(self):
-        return self._z_real_uncertainty
+    def tx_offsets(self):
+        """Return transmitter offsets from frequency metadata"""
 
-    @z_real_uncertainty.setter
-    def z_real_uncertainty(self, val):
-        self.setter_validator("z_real_uncertainty", val, fun=self._uuid_promoter)
+        try:
+            offset_data = self.data_object.metadata["EM Dataset"][
+                "Frequency configurations"
+            ]
+            tx_offsets = {k["Frequency"]: k["Offset"] for k in offset_data}
 
-    @property
-    def z_imag_channel_bool(self):
-        return self._z_imag_channel_bool
+        except KeyError as exception:
+            msg = "Metadata must contain 'Frequency configurations' dictionary with 'Offset' data."
+            raise KeyError(msg) from exception
 
-    @z_imag_channel_bool.setter
-    def z_imag_channel_bool(self, val):
-        self.setter_validator("z_imag_channel_bool", val)
+        return tx_offsets
 
     @property
-    def z_imag_channel(self):
-        return self._z_imag_channel
-
-    @z_imag_channel.setter
-    def z_imag_channel(self, val):
-        self.setter_validator("z_imag_channel", val, fun=self._uuid_promoter)
-
-    @property
-    def z_imag_uncertainty(self):
-        return self._z_imag_uncertainty
-
-    @z_imag_uncertainty.setter
-    def z_imag_uncertainty(self, val):
-        self.setter_validator("z_imag_uncertainty", val, fun=self._uuid_promoter)
+    def unit_conversion(self):
+        """Return time unit conversion factor."""
+        conversion = {
+            "Hertz (Hz)": 1.0,
+        }
+        return conversion[self.data_object.unit]

simpeg_drivers/electromagnetics/time_domain/__init__.py

@@ -9,7 +9,10 @@
 # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
 
 
-from .params import TimeDomainElectromagneticsParams
+from .params import (
+    TimeDomainElectromagneticsForwardParams,
+    TimeDomainElectromagneticsInversionParams,
+)
 
 # pylint: disable=unused-import
 # flake8: noqa