fixing bugs

Author: kujaku11

mtpy/core/mt.py

@@ -1142,6 +1142,8 @@ def edit_curve(self, method="default", tolerance=0.05):
 
         """
 
+        # bring up a gui of some sort.
+
     def to_occam1d(self, data_filename=None, mode="det"):
         """Write an Occam1DData data file.
         :param data_filename: Path to write file, if None returns Occam1DData

mtpy/imaging/plot_stations.py

@@ -31,6 +31,10 @@
 class PlotStations(PlotBase):
     """Plot station locations in map view.
 
+    Can set `x_limits` and `y_limits` to zoom in on a specific area.
+    or set `image_extent` to plot an image in the background.
+    or set `pad` to set the padding around the stations.
+
     Uses contextily to get the basemap.
     See https://contextily.readthedocs.io/en/latest/index.html for more
     information about options.
@@ -115,8 +119,32 @@ def plot(self):
         # make a figure instance
         self.fig = plt.figure(self.fig_num, self.fig_size, dpi=self.fig_dpi)
 
+        if self.pad is None:
+            self.pad = self._get_pad()
+
+        if self.image_extent:
+            x_limits = (self.image_extent[0], self.image_extent[2])
+            y_limits = (self.image_extent[1], self.image_extent[3])
+        elif self.x_limits is not None:
+            x_limits = self.x_limits
+            if self.y_limits is None:
+                y_limits = self._get_ylimits(self.gdf.geometry.y)
+            else:
+                y_limits = self.y_limits
+        elif self.y_limits is not None:
+            y_limits = self.y_limits
+            if self.x_limits is None:
+                x_limits = self._get_xlimits(self.gdf.geometry.x)
+            else:
+                x_limits = self.x_limits
+        else:
+            x_limits = self._get_xlimits(self.gdf.geometry.x)
+            y_limits = self._get_ylimits(self.gdf.geometry.y)
+
         # add and axes
         self.ax = self.fig.add_subplot(1, 1, 1, aspect="equal")
+        self.ax.set_xlim(x_limits)
+        self.ax.set_ylim(y_limits)
 
         # --> plot the background image if desired-----------------------
         if self.image_file is not None:
@@ -164,8 +192,6 @@ def plot(self):
                         f"Could not add base map because {error}"
                     )
 
-        if self.pad is None:
-            self.pad = self._get_pad()
         # set axis properties
         if self.plot_cx:
             self.ax.set_ylabel("latitude (deg)", fontdict=self.font_dict)
@@ -174,8 +200,6 @@ def plot(self):
             self.ax.set_xlabel("relative east (m)", fontdict=self.font_dict)
             self.ax.set_ylabel("relative north (m)", fontdict=self.font_dict)
         self.ax.grid(alpha=0.35, color=(0.35, 0.35, 0.35), lw=0.35)
-        self.ax.set_xlim(self._get_xlimits(self.gdf.geometry.x))
-        self.ax.set_ylim(self._get_ylimits(self.gdf.geometry.y))
         self.ax.set_axisbelow(True)
 
         plt.show()

mtpy/modeling/simpeg/data_2d.py

@@ -13,6 +13,8 @@
 from simpeg.electromagnetics import natural_source as nsem
 from simpeg import data
 
+from mtpy.imaging.mtplot_tools import plot_phase, plot_resistivity
+
 import matplotlib.pyplot as plt
 
 
@@ -305,9 +307,15 @@ def plot_response(self, **kwargs):
         te_data = self.te_data.dobs.reshape(
             (self.n_frequencies, 2, self.n_stations)
         )
+        te_data_errors = self.te_data.standard_deviation.reshape(
+            (self.n_frequencies, 2, self.n_stations)
+        )
         tm_data = self.tm_data.dobs.reshape(
             (self.n_frequencies, 2, self.n_stations)
         )
+        tm_data_errors = self.tm_data.standard_deviation.reshape(
+            (self.n_frequencies, 2, self.n_stations)
+        )
 
         if not self.invert_impedance:
             ax_xy_res = fig.add_subplot(2, 2, 1)
@@ -358,26 +366,58 @@ def plot_response(self, **kwargs):
                 2, 2, 4, sharex=ax_xy_res, sharey=ax_xy_phase
             )
             for ii in range(self.n_stations):
-                ax_xy_res.loglog(
+                plot_resistivity(
+                    ax_xy_res,
                     1.0 / self.frequencies,
-                    np.abs(te_data[:, 0, ii]),
+                    te_data[:, 0, ii],
                     color=(0.5, 0.5, ii / self.n_stations),
+                    label=self.dataframe.station.unique()[ii],
+                    error=te_data_errors[:, 0, ii],
                 )
-                ax_xy_phase.loglog(
+                plot_phase(
+                    ax_xy_phase,
                     1.0 / self.frequencies,
-                    np.abs(te_data[:, 1, ii]),
+                    te_data[:, 1, ii],
                     color=(0.25, 0.25, ii / self.n_stations),
+                    label=self.dataframe.station.unique()[ii],
+                    error=te_data_errors[:, 1, ii],
                 )
-                ax_yx_res.loglog(
+                plot_resistivity(
+                    ax_yx_res,
                     1.0 / self.frequencies,
-                    np.abs(tm_data[:, 0, ii]),
+                    tm_data[:, 0, ii],
                     color=(0.5, ii / self.n_stations, 0.75),
+                    label=self.dataframe.station.unique()[ii],
+                    error=tm_data_errors[:, 0, ii],
                 )
-                ax_yx_phase.loglog(
+                plot_phase(
+                    ax_yx_phase,
                     1.0 / self.frequencies,
-                    np.abs(tm_data[:, 1, ii]),
+                    tm_data[:, 1, ii],
                     color=(0.25, ii / self.n_stations, 0.75),
+                    label=self.dataframe.station.unique()[ii],
+                    error=tm_data_errors[:, 1, ii],
                 )
+                # ax_xy_res.loglog(
+                #     1.0 / self.frequencies,
+                #     np.abs(te_data[:, 0, ii]),
+                #     color=(0.5, 0.5, ii / self.n_stations),
+                # )
+                # ax_xy_phase.loglog(
+                #     1.0 / self.frequencies,
+                #     np.abs(te_data[:, 1, ii]),
+                #     color=(0.25, 0.25, ii / self.n_stations),
+                # )
+                # ax_yx_res.loglog(
+                #     1.0 / self.frequencies,
+                #     np.abs(tm_data[:, 0, ii]),
+                #     color=(0.5, ii / self.n_stations, 0.75),
+                # )
+                # ax_yx_phase.loglog(
+                #     1.0 / self.frequencies,
+                #     np.abs(tm_data[:, 1, ii]),
+                #     color=(0.25, ii / self.n_stations, 0.75),
+                # )
 
             ax_xy_phase.set_xlabel("Period (s)")
             ax_yx_phase.set_xlabel("Period (s)")

mtpy/modeling/simpeg/make_2d_mesh.py

@@ -9,6 +9,8 @@
 # =============================================================================
 
 import numpy as np
+import pandas as pd
+from numpy.typing import NDArray
 
 from discretize import TensorMesh
 from discretize import TreeMesh
@@ -25,7 +27,12 @@
 
 
 class StructuredMesh:
-    def __init__(self, station_locations, frequencies, **kwargs):
+    def __init__(
+        self,
+        station_locations: pd.DataFrame,
+        frequencies: NDArray | list[float],
+        **kwargs
+    ) -> None:
         self.station_locations = station_locations
         self.frequencies = frequencies
         self.topography = None
@@ -48,29 +55,29 @@ def __init__(self, station_locations, frequencies, **kwargs):
             setattr(self, key, value)
 
     @property
-    def frequency_max(self):
+    def frequency_max(self) -> float:
         return self.frequencies.max()
 
     @property
-    def frequency_min(self):
+    def frequency_min(self) -> float:
         return self.frequencies.min()
 
     @property
-    def z1_layer_thickness(self):
+    def z1_layer_thickness(self) -> float:
         return np.round(
             skin_depth(self.frequency_max, self.sigma_background)
             / self.z_factor_max
         )
 
     @property
-    def z_bottom(self):
+    def z_bottom(self) -> float:
         return (
             skin_depth(self.sigma_background, self.frequency_min)
             * self.z_factor_max
         )
 
     @property
-    def z_mesh_down(self):
+    def z_mesh_down(self) -> NDArray[np.float64]:
         for nz_down in range(self.n_max):
             z_mesh_down = (
                 self.z1_layer_thickness
@@ -81,7 +88,7 @@ def z_mesh_down(self):
         return z_mesh_down
 
     @property
-    def z_mesh_up(self):
+    def z_mesh_up(self) -> NDArray[np.float64]:
         z_mesh_up = [
             (
                 self.z1_layer_thickness,
@@ -91,38 +98,38 @@ def z_mesh_up(self):
         ]
         return dis_utils.unpack_widths(z_mesh_up)
 
-    def _make_z_mesh(self):
+    def _make_z_mesh(self) -> NDArray:
         """
         create vertical mesh
         """
 
         return np.r_[self.z_mesh_down, self.z_mesh_up]
 
     @property
-    def dx(self):
+    def dx(self) -> float:
         d_station = np.diff(self.station_locations[:, 0]).min()
         return np.round(d_station / self.x_spacing_factor)
 
     @property
-    def station_total_length(self):
+    def station_total_length(self) -> float:
         return (
             self.station_locations[:, 0].max()
             - self.station_locations[:, 0].min()
         )
 
     @property
-    def n_station_x_cells(self):
+    def n_station_x_cells(self) -> int:
         return int(self.station_total_length / self.dx)
 
     @property
-    def x_padding_cells(self):
+    def x_padding_cells(self) -> float:
         return (
             skin_depth(self.sigma_background, self.frequency_min)
             * self.x_factor_max
         )
 
     @property
-    def n_x_padding(self):
+    def n_x_padding(self) -> int:
         for npadx in range(self.n_max):
             x_pad = dis_utils.unpack_widths(
                 [
@@ -137,7 +144,7 @@ def n_x_padding(self):
                 break
         return npadx
 
-    def _make_x_mesh(self):
+    def _make_x_mesh(self) -> list[tuple[float, int, float]]:
         """
         make horizontal mesh