adding exporters

pyproject.toml

@@ -47,6 +47,7 @@ dependencies = [
     "pyevtk>=1.6",
     "psutil>=5.9",
     "aiofiles>=23.0",
+    "typer >=0.12.3",
 ]
 
 classifiers = [

src/qseek/apps/qseek.py

@@ -170,6 +170,41 @@
 )
 
 
+export = subparsers.add_parser(
+    "export",
+    help="Export detections to different output formats",
+    description="Export detections to different output formats."
+    " Get an overview with `qseek export list`",
+)
+
+export.add_argument(
+    "format",
+    type=str,
+    help="Name of export module, or `list` to list available modules",
+)
+
+export.add_argument(
+    "rundir",
+    type=Path,
+    help="path to existing qseek rundir",
+    nargs="?",
+)
+
+export.add_argument(
+    "export_dir",
+    type=Path,
+    help="path to export directory",
+    nargs="?",
+)
+
+export.add_argument(
+    "--force",
+    action="store_true",
+    default=False,
+    help="overwrite existing output directory",
+)
+
+
 subparsers.add_parser(
     "clear-cache",
     help="clear the cach directory",
@@ -374,6 +409,51 @@ async def start() -> None:
             logger.info("clearing cache directory %s", CACHE_DIR)
             shutil.rmtree(CACHE_DIR)
 
+        case "export":
+            from qseek.exporters.base import Exporter
+
+            def show_table():
+                table = Table(box=box.SIMPLE, header_style=None)
+                table.add_column("Exporter")
+                table.add_column("Description")
+                for exporter in Exporter.get_subclasses():
+                    table.add_row(exporter.__name__.lower(), exporter.__doc__)
+                console.print(table)
+
+            if args.format == "list":
+                show_table()
+                parser.exit()
+
+            if not args.rundir:
+                parser.error("rundir is required for export")
+
+            if args.export_dir is None:
+                parser.error("export directory is required")
+
+            if args.export_dir.exists():
+                if not args.force:
+                    parser.error(f"export directory {args.export_dir} already exists")
+                shutil.rmtree(args.export_dir)
+
+            for exporter in Exporter.get_subclasses():
+                if exporter.__name__.lower() == args.format.lower():
+                    exporter_instance = exporter()
+                    asyncio.run(
+                        exporter_instance.export(
+                            rundir=args.rundir,
+                            outdir=args.export_dir,
+                        )
+                    )
+                    break
+            else:
+                available_exporters = ", ".join(
+                    exporter.__name__ for exporter in Exporter.get_subclasses()
+                )
+                parser.error(
+                    f"unknown exporter: {args.format}"
+                    f"choose fom: {available_exporters}"
+                )
+
         case "modules":
             from qseek.corrections.base import TravelTimeCorrections
             from qseek.features.base import FeatureExtractor

src/qseek/exporters/__init__.py

@@ -0,0 +1,2 @@
+from qseek.exporters.simple import Simple  # noqa
+from qseek.exporters.velest import Velest  # noqa

src/qseek/exporters/base.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import BaseModel
+
+
+class Exporter(BaseModel):
+    async def export(self, rundir: Path, outdir: Path) -> Path:
+        raise NotImplementedError
+
+    @classmethod
+    def get_subclasses(cls) -> tuple[type[Exporter], ...]:
+        """Get the subclasses of this class.
+
+        Returns:
+            list[type]: The subclasses of this class.
+        """
+        return tuple(cls.__subclasses__())

src/qseek/exporters/simple.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from rich import progress
+
+from qseek.exporters.base import Exporter
+from qseek.search import Search
+from qseek.utils import time_to_path
+
+logger = logging.getLogger(__name__)
+
+
+class Simple(Exporter):
+    """Export simple travel times in CSV format (E. Biondi, 2023)."""
+
+    async def export(self, rundir: Path, outdir: Path) -> Path:
+        logger.info("Export simple travel times in CSV format.")
+
+        search = Search.load_rundir(rundir)
+        catalog = search.catalog
+
+        traveltime_dir = outdir / "traveltimes"
+        outdir.mkdir(parents=True)
+        traveltime_dir.mkdir()
+
+        event_file = outdir / "events.csv"
+        self.search.stations.export_csv(outdir / "stations.csv")
+        await catalog.export_csv(event_file)
+
+        for ev in progress.track(
+            catalog,
+            description="Exporting travel times",
+            total=catalog.n_events,
+        ):
+            traveltime_file = traveltime_dir / f"{time_to_path(ev.time)}.csv"
+            with traveltime_file.open("w") as file:
+                file.write(f"# event_id: {ev.uid}\n")
+                file.write(f"# event_time: {ev.time}\n")
+                file.write(f"# event_lat: {ev.lat}\n")
+                file.write(f"# event_lon: {ev.lon}\n")
+                file.write(f"# event_depth: {ev.effective_depth}\n")
+                file.write(f"# event_semblance: {ev.semblance}\n")
+                file.write("# traveltime observations:\n")
+                file.write(
+                    "lat,lon,elevation,network,station,location,phase,confidence,traveltime\n"
+                )
+
+                for receiver in ev.receivers:
+                    for phase, arrival in receiver.phase_arrivals.items():
+                        if arrival.observed is None:
+                            continue
+                        traveltime = arrival.observed.time - ev.time
+                        file.write(
+                            f"{receiver.lat},{receiver.lon},{receiver.effective_elevation},{receiver.network},{receiver.station},{receiver.location},{phase},{arrival.observed.detection_value},{traveltime.total_seconds()}\n",
+                        )
+
+        return outdir

src/qseek/exporters/velest.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import rich
+from rich.prompt import FloatPrompt
+
+from qseek.exporters.base import Exporter
+from qseek.search import Search
+
+logger = logging.getLogger(__name__)
+
+
+class Velest(Exporter):
+    """Crate a VELEST project folder for 1D velocity model estimation."""
+
+    min_pick_semblance: float = 0.3
+    n_picks: dict[str, int] = {}
+    n_events: int = 0
+
+    async def export(self, rundir: Path, outdir: Path) -> Path:
+        rich.print("Exporting qseek search to VELEST project folder")
+        min_pick_semblance = FloatPrompt.ask("Minimum pick confidence", default=0.3)
+
+        self.min_pick_semblance = min_pick_semblance
+
+        outdir.mkdir()
+        search = Search.load_rundir(rundir)
+        catalog = search.catalog  # noqa
+
+        export_info = outdir / "export_info.json"
+        export_info.write_text(self.model_dump_json(indent=2))
+
+        return outdir

src/qseek/images/phase_net.py

@@ -137,7 +137,7 @@ class PhaseNet(ImageFunction):
 
     image: Literal["PhaseNet"] = "PhaseNet"
     model: ModelName = Field(
-        default="ethz",
+        default="original",
         description="SeisBench pre-trained PhaseNet model to use. "
         "Choose from `ethz`, `geofon`, `instance`, `iquique`, `lendb`, `neic`, `obs`,"
         " `original`, `scedc`, `stead`."

src/qseek/models/detection.py

@@ -695,13 +695,20 @@ def get_csv_dict(self) -> dict[str, Any]:
             "time": self.time,
             "lat": round(self.effective_lat, 6),
             "lon": round(self.effective_lon, 6),
-            "depth_ellipsoid": round(self.effective_depth, 2),
+            "depth": round(self.effective_depth, 2),
             "east_shift": round(self.east_shift, 2),
             "north_shift": round(self.north_shift, 2),
             "distance_border": round(self.distance_border, 2),
             "semblance": self.semblance,
             "n_stations": self.n_stations,
         }
+        if self.uncertainty:
+            csv_line.update(
+                {
+                    "uncertainty_horizontal": self.uncertainty.horizontal,
+                    "uncertainty_vertical": self.uncertainty.depth,
+                }
+            )
         for magnitude in self.magnitudes:
             csv_line.update(magnitude.csv_row())
         return csv_line

src/qseek/models/detection_uncertainty.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, computed_field
 from typing_extensions import Self
 
 if TYPE_CHECKING:
@@ -63,3 +63,15 @@ def from_event(
             north=(float(min_offsets[1]), float(max_offsets[1])),
             depth=(float(min_offsets[2]), float(max_offsets[2])),
         )
+
+    @computed_field
+    def total(self) -> float:
+        """Calculate the total uncertainty in [m]."""
+        return float(
+            np.sqrt(sum(self.east) ** 2 + sum(self.north) ** 2 + sum(self.depth) ** 2)
+        )
+
+    @computed_field
+    def horizontal(self) -> float:
+        """Calculate the horizontal uncertainty in [m]."""
+        return float(np.sqrt(sum(self.east) ** 2 + sum(self.north) ** 2))