updating z with units

Author: kujaku11

mtpy/core/mt.py

@@ -14,7 +14,11 @@
 
 from mt_metadata.transfer_functions.core import TF
 
-from mtpy.core import Z, Tipper, COORDINATE_REFERENCE_FRAME_OPTIONS
+from mtpy.core import (
+    Z,
+    Tipper,
+    COORDINATE_REFERENCE_FRAME_OPTIONS,
+)
 from mtpy.core.mt_location import MTLocation
 from mtpy.core.mt_dataframe import MTDataFrame
 from mtpy.utils.estimate_tf_quality_factor import EMTFStats
@@ -1072,9 +1076,7 @@ def add_white_noise(self, value, inplace=True):
             ] = self._transfer_function.transfer_function.real * (
                 noise_real
             ) + (
-                1j
-                * self._transfer_function.transfer_function.imag
-                * noise_imag
+                1j * self._transfer_function.transfer_function.imag * noise_imag
             )
 
             self._transfer_function["transfer_function_error"] = (
@@ -1088,9 +1090,7 @@ def add_white_noise(self, value, inplace=True):
             ] = self._transfer_function.transfer_function.real * (
                 noise_real
             ) + (
-                1j
-                * self._transfer_function.transfer_function.imag
-                * noise_imag
+                1j * self._transfer_function.transfer_function.imag * noise_imag
             )
 
             self._transfer_function["transfer_function_error"] = (

mtpy/core/transfer_function/__init__.py

@@ -1,3 +1,8 @@
+import numpy as np
+
+MT_TO_OHM_FACTOR = 1.0 / np.pi * np.sqrt(5.0 / 8.0) * 10**3.5
+IMPEDANCE_UNITS = {"mt": 1, "ohm": MT_TO_OHM_FACTOR}
+
 from .z import Z
 from .tipper import Tipper
 from .pt import PhaseTensor

mtpy/core/transfer_function/z.py

@@ -18,6 +18,7 @@
 import numpy as np
 
 from .base import TFBase
+from . import MT_TO_OHM_FACTOR, IMPEDANCE_UNITS
 from .pt import PhaseTensor
 from .z_analysis import (
     ZInvariants,
@@ -58,6 +59,7 @@ def __init__(
         z_error=None,
         frequency=None,
         z_model_error=None,
+        units="mt",
     ):
         """Initialize an instance of the Z class.
         :param z_model_error:
@@ -70,6 +72,9 @@ def __init__(
         :param frequency: Array of frequencyuency values corresponding to impedance
             tensor elements, defaults to None.
         :type frequency: np.ndarray(n_frequency), optional
+        :param units: units for the impedance [ "mt" [mV/km/nT] | ohm [Ohms] ]
+        :type units: str
+
         """
 
         super().__init__(
@@ -80,9 +85,9 @@ def __init__(
             _name="impedance",
         )
 
-        self._ohm_factor = 1.0 / np.pi * np.sqrt(5.0 / 8.0) * 10**3.5
-        self._unit_factors = {"mt": 1, "ohm": self._ohm_factor}
-        self.units = "mt"
+        self._ohm_factor = MT_TO_OHM_FACTOR
+        self._unit_factors = IMPEDANCE_UNITS
+        self.units = units
 
     @property
     def units(self):