Add code snippets

python/phaseassoc/catalogConverter.py

@@ -0,0 +1,120 @@
+import pandas as pd
+
+from obspy.core.event import Catalog
+from obspy.core.event import Event
+from obspy.core.event import Pick
+from obspy.core.event import Origin
+from obspy.core.event import OriginQuality
+
+from obspy import UTCDateTime
+from obspy.core.event.base import WaveformStreamID, Comment
+
+
+class CatalogConverter:
+    """
+    Class for converting SeisBench and pyOcto detection results to ObsPy Catalog, which can be saved as QuakeML.
+    """
+    def __init__(self, config, picks, catalog_df, assignments_df, name_of_algorithm):
+        self._catalog_df = catalog_df
+        self._assignments_df = assignments_df
+        self._picks = picks
+        self._config = config
+        self._name_of_algorithm = name_of_algorithm
+        self.catalog = None
+
+    def _convert_pick(self, p):
+        """
+        Function converts picks from SeisBench to ObsPy format
+        :param p: SeisBench pick
+        :return: ObsPy pick
+        """
+        pick = Pick()
+        pick.time = UTCDateTime(p.peak_time.datetime)
+        pick.waveform_id = WaveformStreamID(network_code=p.trace_id.split(".")[0],
+                                            station_code=p.trace_id.split(".")[1],
+                                            channel_code=p.trace_id.split(".")[2])
+        if p.phase == 'P':
+            pick.phase_hint = self._config["P_hint"]
+        elif p.phase == 'S':
+            pick.phase_hint = self._config["S_hint"]
+        pick.evaluation_mode = 'automatic'
+        pick.evaluation_status = 'preliminary'
+        return pick
+
+    def _convert_origin(self, origin_sb, list_of_picks_sb):
+        origin = Origin()
+        origin.time = UTCDateTime(pd.to_datetime(origin_sb.time, unit='s').to_pydatetime())
+        origin.latitude = origin_sb.latitude  # float
+        origin.longitude = origin_sb.longitude  # float
+        origin.depth = origin_sb.depth  # float in kilometers (SWIP5 origin version) down the see level
+        origin.depth_type = 'operator assigned'
+        # TODO: make sure that region is not necessary
+        # origin.region = self._config["region"]
+        origin.evaluation_mode = "automatic"
+        origin.evaluation_status = 'preliminary'
+        origin.comments.append(Comment(text=f"Localized by: {self._name_of_algorithm}", force_resource_id=False))
+        origin.quality = OriginQuality(used_phase_count=len(list_of_picks_sb))
+        return origin
+
+    def _convert_event(self, origin_sb, list_of_picks_sb):
+        """
+        Function convert GaMMa detection to ObsPy Event
+        :param origin_sb:
+        :param list_of_picks_sb:
+        :return:
+        """
+        event = Event()
+        for p in list_of_picks_sb:
+            pick = self._convert_pick(p)
+            event.picks.append(pick)
+        origin = self._convert_origin(origin_sb, list_of_picks_sb)
+        event.origins.append(origin)
+        return event
+
+    @staticmethod
+    def _append_pick_trace_id(pick, stream):
+        """
+        Function assigns channel to pick - it is useful for work with SWIP
+        :param pick:
+        :param stream:
+        :return:
+        """
+        channel = stream[0].stats.channel
+        if pick.phase == "P":
+            pick.trace_id = pick.trace_id + channel[:-1] + "Z"
+        if pick.phase == "S":
+            pick.trace_id = pick.trace_id + channel[:-1] + "E"
+        return pick
+
+    def catalog2obspy(self):
+        """
+        Function convert GaMMa catalog and SeisBench picks
+        :return: ObsPy Catalog object
+        """
+        # TODO: make sure that resource id is necessary
+        #cat = Catalog(resource_id=self._config["resource_id"])
+        cat = Catalog()
+        for j, row in self._catalog_df.iterrows():
+            event = self._catalog_df.iloc[j]
+            event_picks = [self._picks[i] for i in
+                           self._assignments_df[self._assignments_df["event_idx"] ==
+                                                event["idx"]]["pick_idx"]]
+            event_obspy = self._convert_event(event, event_picks)
+            cat.append(event_obspy)
+        self.catalog = cat
+
+    def save_catalog_to_file(self, file_path):
+        """
+        Save ObsPy catalog to a file.
+
+        Args:
+            file_path (str): The file path where the catalog will be saved.
+
+        Returns:
+            None
+        """
+        try:
+            self.catalog.write(file_path, format="QUAKEML")
+            print(f"Catalog saved successfully to {file_path}")
+        except Exception as e:
+            print(f"Error occurred while saving catalog: {e}")