updates to occam and ws3din

v

Author: alkirkby

mtpy/modeling/occam2d/data.py

@@ -102,8 +102,8 @@ def __init__(self, dataframe=None, center_point=None, **kwargs):
         self.occam_dict = {
             "1": "res_xy",
             "2": "phase_xy",
-            "3": "tzx_real",
-            "4": "tzx_imag",
+            "3": "t_zx_real",
+            "4": "t_zx_imag",
             "5": "res_yx",
             "6": "phase_yx",
             "9": "res_xy",
@@ -113,7 +113,8 @@ def __init__(self, dataframe=None, center_point=None, **kwargs):
         self.df_dict = {
             "1": "res_xy",
             "2": "phase_xy",
-            "3": "tzx",
+            "3": "t_zx",
+            "4": "t_zy",
             "5": "res_yx",
             "6": "phase_yx",
         }
@@ -437,26 +438,48 @@ def read_data_file(self, data_fn=None):
 
         # format dataframe
         df = pd.DataFrame(entries)
-        df["tzx"] = df.tzx_real + 1j * df.tzx_imag
-        df["tzx_model_error"] = df.tzx_real_model_error
+        if 't_zy_real' in list(df):
+            df['t_zy_real'] = df['t_zy_real'].fillna(0)
+            df['t_zy_imag'] = df['t_zy_imag'].fillna(0)
+            df['t_zy'] = df['t_zy_real'] + 1j*df['t_zy_imag']
         df["period"] = 1.0 / df.frequency
-        df = df.drop(
-            columns=[
-                "tzx_real",
-                "tzx_imag",
-                "tzx_real_model_error",
-                "tzx_imag_model_error",
-                "frequency",
-            ],
-            axis=1,
-        )
-
-        df = df.groupby(["station", "period"]).agg("first")
+        # df = df.drop(
+        #     columns=[
+        #         "tzx_real",
+        #         "tzx_imag",
+        #         "tzx_real_model_error",
+        #         "tzx_imag_model_error",
+        #         "frequency",
+        #     ],
+        #     axis=1,
+        # )
+
+        # df = df.groupby(["station", "period"]).agg("first")
         df = df.sort_values("profile_offset").reset_index()
         self.dataframe = df
+        
+        # use workaround function to group
+        self._group_df()        
 
         self.model_mode = self._get_model_mode_from_data(res_log)
 
+    def _group_df(self):
+        for station in np.unique(self.dataframe['station']):
+            for period in np.unique(self.dataframe['period']):
+                filt = np.all([self.dataframe['station']==station,
+                               self.dataframe['period']==period],axis=0)
+                if np.any(filt):
+                    for key in ['res_xy','res_yx','phase_xy','phase_yx']:
+                        self.dataframe[key][filt] = np.unique(self.dataframe[key][filt])[0]
+
+                    tx_real_vals = np.unique(np.real(self.dataframe['t_zy'][filt]))
+                    tx_imag_vals = np.unique(np.imag(self.dataframe['t_zy'][filt]))
+                    if np.any(tx_real_vals != 0):
+                        self.dataframe['t_zy'][filt] = tx_real_vals[tx_real_vals != 0][0] +\
+                            1j * tx_imag_vals[tx_imag_vals != 0][0]
+
+                    self.dataframe.drop(self.dataframe[filt].index[1:],inplace=True)
+
     def _get_model_mode_from_data(self, res_log):
         """Get inversion mode from the data.
         :param res_log: DESCRIPTION.
@@ -477,9 +500,9 @@ def _get_model_mode_from_data(self, res_log):
                         inv_list.append(5)
                     else:
                         inv_list.append(10)
-                elif comp == "tzx":
-                    inv_list.append(3)
-                    inv_list.append(4)
+                # elif comp == "tzx":
+                #     inv_list.append(3)
+                #     inv_list.append(4)
                 else:
                     inv_list.append(int(inv_mode))
 
@@ -526,7 +549,8 @@ def _get_data_block(self):
                             value = value.imag
                             error_value = fdf[f"{comp}_model_error"].values[0]
                         elif comp_number in [6]:
-                            value = value + 180
+                            if -180 <= value <= -90:
+                                value = value + 180
                             error_value = fdf[f"{comp}_model_error"].values[0]
                         else:
                             value = value

mtpy/modeling/structured_mesh_3d.py

@@ -2755,17 +2755,18 @@ def to_ws3dinv_intial(self, initial_fn, res_list=None):
         l1 = 0
         layers = []
         for zz in range(self.nodes_z.shape[0] - 1):
-            if not (
-                write_res_model[:, :, zz] == write_res_model[:, :, zz + 1]
-            ).all():
+            # if not (
+            #     write_res_model[:, :, zz] == write_res_model[:, :, zz + 1]
+            # ).all():
                 layers.append((l1, zz))
                 l1 = zz + 1
         # need to add on the bottom layers
         layers.append((l1, self.nodes_z.shape[0] - 1))
 
         # write out the layers from resmodel
         for ll in layers:
-            lines.append(f"{ll[0] + 1} {ll[1] + 1}\n")
+            if nr > 0:
+                lines.append(f"{ll[0] + 1} {ll[1] + 1}\n")
             for nn in range(self.nodes_north.shape[0]):
                 for ee in range(self.nodes_east.shape[0]):
                     if nr > 0:

mtpy/modeling/ws3dinv/data.py

@@ -205,10 +205,10 @@ def get_n_stations(self):
         if self.dataframe is not None:
             return (
                 self.dataframe.loc[
-                    (self.dataframe.zxx != 0)
-                    | (self.dataframe.zxy != 0)
-                    | (self.dataframe.zyx != 0)
-                    | (self.dataframe.zyy != 0),
+                    (self.dataframe.z_xx != 0)
+                    | (self.dataframe.z_xy != 0)
+                    | (self.dataframe.z_yx != 0)
+                    | (self.dataframe.z_yy != 0),
                     "station",
                 ]
                 .unique()
@@ -256,7 +256,7 @@ def write_data_file(self, **kwargs):
         # -----Write data file--------------------------------------------------
         n_stations = self.get_n_stations()
         n_periods = self.period.size
-        station_locations = self.dataframe.station_locations.copy()
+        station_locations = self.station_locations.copy()
 
         lines = []
         data_lines = []
@@ -266,56 +266,62 @@ def write_data_file(self, **kwargs):
 
         # write N-S locations
         lines.append("Station_Location: N-S \n")
-        for ii in range(n_stations / self.n_z + 1):
-            for ll in range(self.n_z):
-                index = ii * self.n_z + ll
-                try:
-                    lines.append(
-                        f"{station_locations.model_north[index]:+.4e} "
-                    )
-                except IndexError:
-                    pass
-            lines.append("\n")
+        for ii in range(n_stations):# / self.n_z + 1
+            # for ll in range(self.n_z):
+            #     index = ii * self.n_z + ll
+            try:
+                lines.append(
+                    f"{station_locations.model_north[ii]:+.4e} "
+                )
+            except IndexError:
+                pass
+            if ii % 8 == 7:
+                lines.append("\n")
+        lines.append("\n")
 
         # write E-W locations
         lines.append("Station_Location: E-W \n")
-        for ii in range(n_stations / self.n_z + 1):
-            for ll in range(self.n_z):
-                index = ii * self.n_z + ll
-                try:
-                    lines.append(
-                        f"{station_locations.model_east[index]:+.4e} "
-                    )
-                except IndexError:
-                    pass
-            lines.append("\n")
+        for ii in range(n_stations):# / self.n_z + 1
+            # for ll in range(self.n_z):
+            #     index = ii * self.n_z + ll
+            try:
+                lines.append(
+                    f"{station_locations.model_east[ii]:+.4e} "
+                )
+            except IndexError:
+                pass
+
+            if ii % 8 == 7:
+                lines.append("\n")
+        lines.append("\n")
 
         # write impedance tensor components
         for ii, p1 in enumerate(self.period):
-            pdf = self.get_period_df(p1)
+            # pdf = self.get_period_df(p1)
+            pdf = self.dataframe[self.dataframe['period']==p1]
             data_lines.append(f"DATA_Period: {p1:3.6f}\n")
             error_lines.append(f"ERROR_Period: {p1:3.6f}\n")
             error_map_lines.append(f"ERMAP_Period: {p1:3.6f}\n")
             for row in pdf.itertuples():
                 data_lines.append(
-                    f"{row.zxx.real * zconv:+.4e} "
-                    f"{row.zxx.imag * zconv:+.4e} "
-                    f"{row.zxy.real * zconv:+.4e} "
-                    f"{row.zxy.imag * zconv:+.4e} "
-                    f"{row.zyx.real * zconv:+.4e} "
-                    f"{row.zyx.imag * zconv:+.4e} "
-                    f"{row.zyy.real * zconv:+.4e} "
-                    f"{row.zyy.imag * zconv:+.4e}\n"
+                    f"{row.z_xx.real * zconv:+.4e} "
+                    f"{row.z_xx.imag * zconv:+.4e} "
+                    f"{row.z_xy.real * zconv:+.4e} "
+                    f"{row.z_xy.imag * zconv:+.4e} "
+                    f"{row.z_yx.real * zconv:+.4e} "
+                    f"{row.z_yx.imag * zconv:+.4e} "
+                    f"{row.z_yy.real * zconv:+.4e} "
+                    f"{row.z_yy.imag * zconv:+.4e}\n"
                 )
                 error_lines.append(
-                    f"{row.zxx.real * self.z_error:+.4e} "
-                    f"{row.zxx.imag * self.z_error:+.4e} "
-                    f"{row.zxy.real * self.z_error:+.4e} "
-                    f"{row.zxy.imag * self.z_error:+.4e} "
-                    f"{row.zyx.real * self.z_error:+.4e} "
-                    f"{row.zyx.imag * self.z_error:+.4e} "
-                    f"{row.zyy.real * self.z_error:+.4e} "
-                    f"{row.zyy.imag * self.z_error:+.4e} "
+                    f"{row.z_xx.real * self.z_error:+.4e} "
+                    f"{row.z_xx.imag * self.z_error:+.4e} "
+                    f"{row.z_xy.real * self.z_error:+.4e} "
+                    f"{row.z_xy.imag * self.z_error:+.4e} "
+                    f"{row.z_yx.real * self.z_error:+.4e} "
+                    f"{row.z_yx.imag * self.z_error:+.4e} "
+                    f"{row.z_yy.real * self.z_error:+.4e} "
+                    f"{row.z_yy.imag * self.z_error:+.4e} "
                 )
 
                 error_map_lines.append(
@@ -326,11 +332,13 @@ def write_data_file(self, **kwargs):
                     f"{self.z_error_map[2]:+.4e} "
                     f"{self.z_error_map[2]:+.4e} "
                     f"{self.z_error_map[3]:+.4e} "
-                    f"{self.z_error_map[3]:+.4e} "
+                    f"{self.z_error_map[3]:+.4e}\n"
                 )
 
         with open(self.data_filename, "w") as fid:
-            fid.write("".join(lines, data_lines, error_lines, error_map_lines))
+            for dlist in [lines, data_lines, error_lines, error_map_lines]:
+                for line in dlist:
+                    fid.write(line)
 
         self.logger.info(f"Wrote WS3DINV file to {self.data_filename}")
         return self.data_filename