updating mtdata

Author: kujaku11

mtpy/core/mt_data.py

@@ -59,15 +59,9 @@ class MTData(OrderedDict, MTStations):
     """
 
     def __init__(self, mt_list=None, **kwargs):
-        if mt_list is not None:
-            for mt_obj in mt_list:
-                self.add_station(mt_obj, compute_relative_location=False)
 
-        MTStations.__init__(
-            self,
-            kwargs.pop("utm_epsg", None),
-            datum_epsg=kwargs.pop("datum_epsg", None),
-            **kwargs,
+        self._coordinate_reference_frame_options = (
+            COORDINATE_REFERENCE_FRAME_OPTIONS
         )
 
         self.z_model_error = ModelErrors(
@@ -83,6 +77,7 @@ def __init__(self, mt_list=None, **kwargs):
             mode="tipper",
         )
         self.data_rotation_angle = 0
+        self.coordinate_reference_frame = "ned"
 
         self.model_parameters = {}
 
@@ -101,8 +96,15 @@ def __init__(self, mt_list=None, **kwargs):
             "model_parameters",
         ]
 
-        self._coordinate_reference_frame_options = (
-            COORDINATE_REFERENCE_FRAME_OPTIONS
+        if mt_list is not None:
+            for mt_obj in mt_list:
+                self.add_station(mt_obj, compute_relative_location=False)
+
+        MTStations.__init__(
+            self,
+            kwargs.pop("utm_epsg", None),
+            datum_epsg=kwargs.pop("datum_epsg", None),
+            **kwargs,
         )
 
     def _validate_item(self, mt_obj):
@@ -218,7 +220,9 @@ def clone_empty(self):
     @property
     def coordinate_reference_frame(self):
         """coordinate reference frame ned or enu"""
-        return self._coordinate_reference_frame
+        return self._coordinate_reference_frame_options[
+            self._coordinate_reference_frame
+        ].upper()
 
     @coordinate_reference_frame.setter
     def coordinate_reference_frame(self, value):
@@ -670,9 +674,7 @@ def interpolate(
                 )
 
             else:
-                mt_data.add_station(
-                    new_mt_obj, compute_relative_location=False
-                )
+                mt_data.add_station(new_mt_obj, compute_relative_location=False)
 
         if not inplace:
             return mt_data
@@ -696,9 +698,7 @@ def rotate(self, rotation_angle, inplace=True):
             if not inplace:
                 rot_mt_obj = mt_obj.copy()
                 rot_mt_obj.rotation_angle = rotation_angle
-                mt_data.add_station(
-                    rot_mt_obj, compute_relative_location=False
-                )
+                mt_data.add_station(rot_mt_obj, compute_relative_location=False)
             else:
                 mt_obj.rotation_angle = rotation_angle
 
@@ -792,12 +792,8 @@ def compute_model_errors(
             self.t_model_error.floor = t_floor
 
         for mt_obj in self.values():
-            mt_obj.compute_model_z_errors(
-                **self.z_model_error.error_parameters
-            )
-            mt_obj.compute_model_t_errors(
-                **self.t_model_error.error_parameters
-            )
+            mt_obj.compute_model_z_errors(**self.z_model_error.error_parameters)
+            mt_obj.compute_model_t_errors(**self.t_model_error.error_parameters)
 
     def get_nearby_stations(self, station_key, radius, radius_units="m"):
         """Get stations close to a given station.

tests/core/test_mt_data.py

@@ -75,7 +75,20 @@ def test_clone_empty(self):
 
     def test_initialization_utm_epsg_no_mt_list(self):
         md = MTData(utm_epsg=self.utm_epsg)
-        self.assertEqual(self.md.utm_epsg, self.utm_epsg)
+        self.assertEqual(md.utm_epsg, self.utm_epsg)
+
+    def test_coordinate_reference_frame(self):
+        self.assertEqual("NED", self.md.coordinate_reference_frame)
+
+    def test_coordinate_reference_frame_set(self):
+        md = MTData(mt_list=self.mt_list_01, coordinate_reference_frame="enu")
+
+        with self.subTest("mtdata"):
+            self.assertEqual("ENU", md.coordinate_reference_frame)
+
+        for mt_obj in md.values():
+            with self.subTest(mt_obj.station):
+                self.assertEqual("ENU", mt_obj.coordinate_reference_frame)
 
     def test_initialization_datum_epsg_no_mt_list(self):
         md = MTData(datum_epsg=self.datum_epsg)