diff --git a/pyproject.toml b/pyproject.toml index 2a2491b0..2752ed72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ dependencies = [ "pyevtk>=1.6", "psutil>=5.9", "aiofiles>=23.0", + "typer >=0.12.3", ] classifiers = [ diff --git a/src/qseek/apps/qseek.py b/src/qseek/apps/qseek.py index 55561abb..b3acbd93 100644 --- a/src/qseek/apps/qseek.py +++ b/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 diff --git a/src/qseek/exporters/__init__.py b/src/qseek/exporters/__init__.py new file mode 100644 index 00000000..3e842644 --- /dev/null +++ b/src/qseek/exporters/__init__.py @@ -0,0 +1,2 @@ +from qseek.exporters.simple import Simple # noqa +from qseek.exporters.velest import Velest # noqa diff --git a/src/qseek/exporters/base.py b/src/qseek/exporters/base.py new file mode 100644 index 00000000..c26e412c --- /dev/null +++ b/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__()) diff --git a/src/qseek/exporters/simple.py b/src/qseek/exporters/simple.py new file mode 100644 index 00000000..63dde913 --- /dev/null +++ b/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 diff --git a/src/qseek/exporters/velest.py b/src/qseek/exporters/velest.py new file mode 100644 index 00000000..7286650a --- /dev/null +++ b/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 diff --git a/src/qseek/images/phase_net.py b/src/qseek/images/phase_net.py index cdd47c32..a5e2bf55 100644 --- a/src/qseek/images/phase_net.py +++ b/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`." diff --git a/src/qseek/models/detection.py b/src/qseek/models/detection.py index a070c15b..acf0e288 100644 --- a/src/qseek/models/detection.py +++ b/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 diff --git a/src/qseek/models/detection_uncertainty.py b/src/qseek/models/detection_uncertainty.py index 5b2b386c..6fcd4431 100644 --- a/src/qseek/models/detection_uncertainty.py +++ b/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))