Phase Association code

pyoctoAssociator.py

@@ -0,0 +1,100 @@
+import os
+
+import numpy as np
+import pandas as pd
+import scipy.io
+
+import pyocto
+
+
+class PhaseAssociation:
+    """
+    Class for phase association using PyOcto.
+    """
+
+    def __init__(self, inventory, config):
+        """
+        Initialize the PhaseAssociation class.
+
+        Args:
+            inventory: ObsPy inventory object.
+            config (dict): Configuration parameters.
+        """
+        self._inv = inventory
+        self._config = config
+        self._setup_velocity_model()
+        self._setup_associator()
+
+    def _setup_velocity_model(self):
+        """
+        Set up the velocity model based on configuration.
+        """
+        velocity_config = self._config["velocity"]
+        if self._config["velocity"]["path"] is None:
+            self.velocity_model = pyocto.VelocityModel0D(
+                p_velocity=self._config["velocity"]["0D_P_vel"],
+                s_velocity=self._config["velocity"]["0D_S_vel"],
+                tolerance=self._config["params"]["tolerance"],
+                association_cutoff_distance=self._config["params"]["association_cutoff_distance"])
+        else:
+            if not os.path.exists(velocity_config["path_out"]):
+                self._create_1d_vel_model()
+            self.velocity_model = pyocto.VelocityModel1D(velocity_config["path_out"],
+                                                         tolerance=self._config["params"]["tolerance"])
+
+    def _mat2df(self):
+        # Load the .mat file
+        mat_contents = scipy.io.loadmat(self._config["velocity"]["path"])
+        data = mat_contents["d"][0][0]
+        depth = data[0].T[0]
+        depth = np.append(depth[::2], depth[-1])
+        vel_p = data[1].T[0]
+        vel_p = np.append(vel_p[0], vel_p[::2])
+        vel_s = data[2].T[0]
+        vel_s = np.append(vel_s[0], vel_s[::2])
+
+        # Create DataFrame from the lists
+        data = {
+            'depth': depth,
+            'vp': vel_p,
+            'vs': vel_s
+        }
+        layers = pd.DataFrame(data)
+
+        return layers
+
+    def _create_1d_vel_model(self):
+
+        layers = self._mat2df()
+        pyocto.VelocityModel1D.create_model(
+            layers,
+            self._config["velocity"]["grid_spacing"],
+            self._config["velocity"]["grid_hor_extent"],
+            self._config["velocity"]["grid_ver_extent"],
+            self._config["velocity"]["path_out"]
+        )
+
+    def _setup_associator(self):
+        """
+        Set up the associator based on configuration.
+        """
+        grid_config = self._config["grid"]
+        associator_params = self._config["params"]
+        self.associator = pyocto.OctoAssociator.from_area(
+            lat=(grid_config["lat_min"], grid_config["lat_max"]),
+            lon=(grid_config["lon_min"], grid_config["lon_max"]),
+            zlim=(grid_config["z_min"], grid_config["z_max"]),
+            time_before=associator_params["time_before"],
+            velocity_model=self.velocity_model,
+            n_picks=associator_params["n_picks"],
+            n_p_picks=associator_params["n_p_picks"],
+            n_s_picks=associator_params["n_s_picks"],
+            n_p_and_s_picks=associator_params["n_p_and_s_picks"],
+            pick_match_tolerance=associator_params["pick_match_tolerance"]
+        )
+
+    def associate_phases(self, picks):
+        stations = self.associator.inventory_to_df(self._inv)
+        events, assignments = self.associator.associate_seisbench(picks, stations)
+        events = self.associator.transform_events(events)
+        return events, assignments