updating how z units are propagated

Author: kujaku11

mtpy/core/mt.py

@@ -18,6 +18,8 @@
     Z,
     Tipper,
     COORDINATE_REFERENCE_FRAME_OPTIONS,
+    MT_TO_OHM_FACTOR,
+    IMPEDANCE_UNITS,
 )
 from mtpy.core.mt_location import MTLocation
 from mtpy.core.mt_dataframe import MTDataFrame
@@ -78,7 +80,7 @@ class MT(TF, MTLocation):
 
     """
 
-    def __init__(self, fn=None, **kwargs):
+    def __init__(self, fn=None, impedance_units="mt", **kwargs):
         tf_kwargs = {}
         for key in [
             "period",
@@ -117,6 +119,9 @@ def __init__(self, fn=None, **kwargs):
             self.station_metadata.transfer_function.sign_convention
         )
 
+        self._impedance_unit_factors = IMPEDANCE_UNITS
+        self.impedance_units = impedance_units
+
         for key, value in kwargs.items():
             setattr(self, key, value)
 
@@ -215,6 +220,23 @@ def coordinate_reference_frame(self, value):
 
         self.station_metadata.transfer_function.sign_convention = value
 
+    @property
+    def impedance_units(self):
+        """impedance units"""
+        return self._impedance_units
+
+    @impedance_units.setter
+    def impedance_units(self, value):
+        """impedance units setter options are [ mt | ohm ]"""
+        if not isinstance(value, str):
+            raise TypeError("Units input must be a string.")
+        if value.lower() not in self._impedance_unit_factors.keys():
+            raise ValueError(
+                f"{value} is not an acceptable unit for impedance."
+            )
+
+        self._impedance_units = value
+
     @property
     def rotation_angle(self):
         """Rotation angle in degrees from north. In the coordinate reference frame"""
@@ -291,6 +313,7 @@ def Z(self):
                 z_error=self.impedance_error.to_numpy(),
                 frequency=self.frequency,
                 z_model_error=self.impedance_model_error.to_numpy(),
+                units=self.impedance_units,
             )
         return Z()
 
@@ -301,6 +324,9 @@ def Z(self, z_object):
         recalculate phase tensor and invariants, which shouldn't change except
         for strike angle.
         """
+        # if a z object is given the underlying data is in mt units, even
+        # if the units are set to ohm.
+        z_object.units = "mt"
         if not isinstance(z_object.frequency, type(None)):
             if self.frequency.size != z_object.frequency.size:
                 self.frequency = z_object.frequency
@@ -638,7 +664,7 @@ def plot_depth_of_penetration(self, **kwargs):
 
         return PlotPenetrationDepth1D(self, **kwargs)
 
-    def to_dataframe(self, utm_crs=None, cols=None):
+    def to_dataframe(self, utm_crs=None, cols=None, impedance_units="mt"):
         """Create a dataframe from the transfer function for use with plotting
         and modeling.
         :param cols:
@@ -648,6 +674,8 @@ def to_dataframe(self, utm_crs=None, cols=None):
         :param eter utm_crs: The utm zone to project station to, could be a
             name, pyproj.CRS, EPSG number, or anything that pyproj.CRS can intake.
         :type eter utm_crs: string, int, :class:`pyproj.CRS`
+        :param impedance_units: ["mt" [mV/km/nT] | "ohm" [Ohms] ]
+        :type impedance_units: str
         """
         if utm_crs is not None:
             self.utm_crs = utm_crs
@@ -671,19 +699,22 @@ def to_dataframe(self, utm_crs=None, cols=None):
 
         mt_df.dataframe.loc[:, "period"] = self.period
         if self.has_impedance():
-            mt_df.from_z_object(self.Z)
+            z_object = self.Z
+            z_object.units = impedance_units
+            mt_df.from_z_object(z_object)
+            mt_df.from_z_object(z_object)
         if self.has_tipper():
             mt_df.from_t_object(self.Tipper)
 
         return mt_df
 
-    def from_dataframe(self, mt_df):
+    def from_dataframe(self, mt_df, impedance_units="mt"):
         """Fill transfer function attributes from a dataframe for a single station.
         :param mt_df:
         :param df: DESCRIPTION.
         :type df: TYPE
-        :return: DESCRIPTION.
-        :rtype: TYPE
+        :param impedance_units: ["mt" [mV/km/nT] | "ohm" [Ohms] ]
+        :type impedance_units: str
         """
 
         if not isinstance(mt_df, MTDataFrame):
@@ -717,11 +748,9 @@ def from_dataframe(self, mt_df):
 
         self.tf_id = self.station
 
-        # self._transfer_function = self._initialize_transfer_function(
-        #     mt_df.period
-        # )
-
-        self.Z = mt_df.to_z_object()
+        z_obj = mt_df.to_z_object()
+        z_obj.units = impedance_units
+        self.Z = z_obj
         self.Tipper = mt_df.to_t_object()
 
     def compute_model_z_errors(

mtpy/core/transfer_function/base.py

@@ -102,8 +102,11 @@ def __eq__(self, other):
 
         # loop over variables to make sure they are all the same.
         for var in list(self._dataset.data_vars):
-            if not (self._dataset[var] == other._dataset[var]).all().data:
-                return False
+            try:
+                if not (self._dataset[var] == other._dataset[var]).all().data:
+                    return False
+            except TypeError:
+                continue
         return True
 
     def __deepcopy__(self, memo):
@@ -164,9 +167,7 @@ def _initialize(
             tf_error = np.zeros_like(
                 tf_model_error, dtype=self._tf_dtypes["tf_error"]
             )
-            periods = self._validate_frequency(
-                periods, tf_model_error.shape[0]
-            )
+            periods = self._validate_frequency(periods, tf_model_error.shape[0])
 
         else:
             periods = self._validate_frequency(periods)

mtpy/core/transfer_function/z.py

@@ -77,6 +77,19 @@ def __init__(
 
         """
 
+        self._ohm_factor = MT_TO_OHM_FACTOR
+        self._unit_factors = IMPEDANCE_UNITS
+        self.units = units
+
+        # if units input is ohms, then we want to scale them to mt units that
+        # way the underlying data is consistent in [mV/km/nT]
+        if z is not None:
+            z = z * self._scale_factor
+        if z_error is not None:
+            z_error = z_error * self._scale_factor
+        if z_model_error is not None:
+            z_model_error = z_model_error * self._scale_factor
+
         super().__init__(
             tf=z,
             tf_error=z_error,
@@ -85,10 +98,6 @@ def __init__(
             _name="impedance",
         )
 
-        self._ohm_factor = MT_TO_OHM_FACTOR
-        self._unit_factors = IMPEDANCE_UNITS
-        self.units = units
-
     @property
     def units(self):
         """impedance units"""
@@ -122,7 +131,7 @@ def z(self):
 
     @z.setter
     def z(self, z):
-        """Set the attribute 'z'.
+        """Set the attribute 'z'. Should be in units of mt [mV/km/nT]
         :param z: Complex impedance tensor array.
         :type z: np.ndarray(nfrequency, 2, 2)
         """

tests/core/test_mt.py

@@ -100,9 +100,7 @@ def setUpClass(self):
             [[[35.26438968, 0.20257033], [0.20257033, 35.26438968]]]
         )
 
-        self.pt = np.array(
-            [[[1.00020002, -0.020002], [-0.020002, 1.00020002]]]
-        )
+        self.pt = np.array([[[1.00020002, -0.020002], [-0.020002, 1.00020002]]])
         self.pt_error = np.array(
             [[[0.01040308, 0.02020604], [0.02020604, 0.01040308]]]
         )
@@ -367,6 +365,25 @@ def test_from_dataframe_fail(self):
         self.assertRaises(TypeError, self.m1.from_dataframe, "a")
 
 
+class TestMT2DataFrameOhms(unittest.TestCase):
+    @classmethod
+    def setUpClass(self):
+        self.m1 = MT(TF_EDI_CGG)
+        self.m1.read()
+
+        self.mt_df = self.m1.to_dataframe(impedance_units="ohm")
+
+    def test_impedance_in_ohms(self):
+        z_obj = self.m1.Z
+        z_obj.units = "ohm"
+
+        self.assertEqual(z_obj, self.mt_df.to_z_object())
+
+    def test_impedance_not_equal(self):
+
+        self.assertNotEqual(self.m1.Z, self.mt_df.to_z_object())
+
+
 # =============================================================================
 # Run
 # =============================================================================

tests/core/transfer_function/test_z.py

@@ -62,6 +62,9 @@ def test_resistivity_model_error(self):
     def test_phase_model_error(self):
         self.assertEqual(None, self.z.phase_model_error)
 
+    def test_units(self):
+        self.assertEqual("mt", self.z.units)
+
 
 class TestZSetResPhase(unittest.TestCase):
     """