Development (#2)

* fixes * adding plots * adding plots * eq rate plot * minor changes * adding fast-marching * fmm: implementation * wip: fast-marching implementation * cake: loading .nd files // adding depth constrain * wip: fast-marching implementation * cake: fixing new structure * fast-marching: adding NonLinLoc support * wip: fast-marching implementation * wip: fast-marching implementation * wip: fast-marching implementation * wip: fast-marching implementation * wip: fast-marching implementation * wip: fast-marching implementation * wip: fixing tests * wip: finishing up fast-marching * wip: finishing up fast-marching * fix typing * fixing tests * refactoring waveform provider * finishing up * update README --------- Co-authored-by: Marius Isken <marius.isken@gfz-potsdam.de>
pyrocko · Sep 13, 2023 · e26d5d2 · e26d5d2
1 parent 1094db9
commit e26d5d2
Show file tree

Hide file tree

Showing 43 changed files with 2,623 additions and 599 deletions.
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -23,4 +23,4 @@ jobs:
           pip install .[dev]
       - name: Test with pytest
         run: |
-          pytest
+          pytest -m "not plot"
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -19,6 +19,6 @@ repos:
         # language_version: python3.9
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     # Ruff version.
-    rev: "v0.0.275"
+    rev: "v0.0.287"
     hooks:
       - id: ruff
diff --git a/README.md b/README.md
@@ -8,21 +8,28 @@
 [![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)](https://pre-commit.com/)
 <!-- [![PyPI](https://img.shields.io/pypi/v/lassie)](https://pypi.org/project/lassie/) -->
 
-Lassie is an earthquake detector based on stacking and migration method. It combines neural network phase picks with an iterative octree localisation approach.
+Lassie is an earthquake detection and localisation framework based on stacking and migration method. It combines neural network phase picks with an iterative octree localisation approach for accurate localisation of seismic events.
 
-Key features are of the tools are:
+Key features are of the earthquake detection and localisation framework are:
 
-* Phase detection using SeisBench
-* Efficient and accurate Octree localisation approach
-* Extraction of event features
+* Earthquake phase detection using machine-learning pickers from [SeisBench](https://github.com/seisbench/seisbench)
+* Octree localisation approach for efficient and accurate search
+* Different velocity models:
+  * Constant velocity
+  * 1D Layered velocity model
+  * 3D fast-marching velocity model (NonLinLoc compatible)
+* Extraction of earthquake event features:
     * Local magnitudes
     * Ground motion attributes
-* Determination of station corrections
+* Automatic extraction of modelled and picked travel times
+* Calculation and application of station corrections / station delay times
+
+Lassie is built on top of [Pyrocko](https://pyrocko.org).
 
 ## Installation
 
 ```sh
-git clone https://github.com/miili/lassie-v2
+git clone https://github.com/pyrocko/lassie-v2
 cd lassie-v2
 pip3 install .
 ```
@@ -32,12 +39,12 @@ pip3 install .
 Initialize a new project in a fresh directory.
 
 ```sh
-lassie new project-dir/
+lassie init my-project/
 ```
 
-Edit the `search.json`
+Edit the `my-project.json`
 
-Start the detection
+Start the earthquake detection with
 
 ```sh
 lassie run search.json
@@ -52,14 +59,13 @@ The simplest and recommended way of installing from source:
 Local development through pip.
 
 ```sh
-cd lightguide
+cd lassie-v2
 pip3 install .[dev]
 ```
 
 The project utilizes pre-commit for clean commits, install the hooks via:
 
 ```sh
-pip install pre-commit
 pre-commit install
 ```
 
@@ -73,4 +79,4 @@ Please cite lassie as:
 
 Contribution and merge requests by the community are welcome!
 
-Lassie-v@ was written by Marius Paul Isken and is licensed under the GNU GENERAL PUBLIC LICENSE v3.
+Lassie-v2 was written by Marius Paul Isken and is licensed under the GNU GENERAL PUBLIC LICENSE v3.
diff --git a/docs/index.md b/docs/index.md
@@ -0,0 +1,17 @@
+# Welcome to MkDocs
+
+For full documentation visit [mkdocs.org](https://www.mkdocs.org).
+
+## Commands
+
+* `mkdocs new [dir-name]` - Create a new project.
+* `mkdocs serve` - Start the live-reloading docs server.
+* `mkdocs build` - Build the documentation site.
+* `mkdocs -h` - Print help message and exit.
+
+## Project layout
+
+    mkdocs.yml    # The configuration file.
+    docs/
+        index.md  # The documentation homepage.
+        ...       # Other markdown pages, images and other files.
diff --git a/lassie/apps/lassie.py b/lassie/apps/lassie.py
@@ -3,22 +3,18 @@
 import argparse
 import asyncio
 import logging
-from datetime import datetime
+import shutil
 from pathlib import Path
 
 import nest_asyncio
 from pkg_resources import get_distribution
 
 from lassie.console import console
-from lassie.images import ImageFunctions
-from lassie.images.phase_net import PhaseNet
 from lassie.models import Stations
-from lassie.search import SquirrelSearch
+from lassie.search import Search
 from lassie.server import WebServer
 from lassie.station_corrections import StationCorrections
-from lassie.tracers import RayTracers
-from lassie.tracers.cake import CakeTracer
-from lassie.utils import ANSI, setup_rich_logging
+from lassie.utils import CACHE_DIR, setup_rich_logging
 
 nest_asyncio.apply()
 
@@ -66,6 +62,14 @@ def main() -> None:
     )
     continue_run.add_argument("rundir", type=Path, help="existing runding to continue")
 
+    init_project = subparsers.add_parser(
+        "init",
+        help="initialize a new Lassie project",
+    )
+    init_project.add_argument(
+        "folder", type=Path, help="folder to initialize project in"
+    )
+
     features = subparsers.add_parser(
         "feature-extraction",
         help="extract features from an existing run",
@@ -74,7 +78,7 @@ def main() -> None:
     features.add_argument("rundir", type=Path, help="path of existing run")
 
     station_corrections = subparsers.add_parser(
-        "station-corrections",
+        "corrections",
         help="analyse station corrections from existing run",
         description="analyze and plot station corrections from a finished run",
     )
@@ -88,27 +92,26 @@ def main() -> None:
 
     serve = subparsers.add_parser(
         "serve",
-        help="serve results from an existing run",
+        help="start webserver and serve results from an existing run",
         description="start a webserver and serve detections and results from a run",
     )
     serve.add_argument("rundir", type=Path, help="rundir to serve")
 
-    new = subparsers.add_parser(
-        "new",
-        help="initialize a new project",
+    subparsers.add_parser(
+        "clear-cache",
+        help="clear the cach directory",
     )
-    new.add_argument("folder", type=Path, help="folder to initialize project in")
 
     dump_schemas = subparsers.add_parser(
         "dump-schemas",
-        help="dump models to json-schema (development)",
+        help="dump data models to json-schema (development)",
     )
     dump_schemas.add_argument("folder", type=Path, help="folder to dump schemas to")
 
     args = parser.parse_args()
     setup_rich_logging(level=logging.INFO - args.verbose * 10)
 
-    if args.command == "new":
+    if args.command == "init":
         folder: Path = args.folder
         if folder.exists():
             raise FileExistsError(f"Folder {folder} already exists")
@@ -117,79 +120,74 @@ def main() -> None:
         pyrocko_stations = folder / "pyrocko-stations.yaml"
         pyrocko_stations.touch()
 
-        config = SquirrelSearch(
-            ray_tracers=RayTracers(root=[CakeTracer()]),
-            image_functions=ImageFunctions(
-                root=[PhaseNet(phase_map={"P": "cake:P", "S": "cake:S"})]
-            ),
-            stations=Stations(pyrocko_station_yamls=[pyrocko_stations]),
-            waveform_data=[Path("/data/")],
-            time_span=(
-                datetime.fromisoformat("2023-04-11T00:00:00+00:00"),
-                datetime.fromisoformat("2023-04-18T00:00:00+00:00"),
-            ),
+        config = Search(
+            stations=Stations(
+                pyrocko_station_yamls=[pyrocko_stations.relative_to(folder)]
+            )
         )
 
-        config_file = folder / "config.json"
+        config_file = folder / f"{folder.name}.json"
         config_file.write_text(config.model_dump_json(by_alias=False, indent=2))
+
         logger.info("initialized new project in folder %s", folder)
-        logger.info(
-            "start detecting with:\n\t%slassie run config.json%s", ANSI.Bold, ANSI.Reset
-        )
+        logger.info("start detection with: lassie run %s", config_file.name)
 
     elif args.command == "run":
-        search = SquirrelSearch.from_config(args.config)
-        search.init_rundir(force=args.force)
+        search = Search.from_config(args.config)
 
         webserver = WebServer(search)
 
         async def _run() -> None:
             http = asyncio.create_task(webserver.start())
-            await search.scan_squirrel()
+            await search.start(force_rundir=args.force)
             await http
 
         asyncio.run(_run())
 
     elif args.command == "continue":
-        search = SquirrelSearch.load_rundir(args.rundir)
-        if search.progress.time_progress:
-            console.rule(f"Continuing search from {search.progress.time_progress}")
+        search = Search.load_rundir(args.rundir)
+        if search._progress.time_progress:
+            console.rule(f"Continuing search from {search._progress.time_progress}")
         else:
             console.rule("Starting search from scratch")
 
         webserver = WebServer(search)
 
         async def _run() -> None:
             http = asyncio.create_task(webserver.start())
-            await search.scan_squirrel()
+            await search.start()
             await http
 
         asyncio.run(_run())
 
     elif args.command == "feature-extraction":
-        search = SquirrelSearch.load_rundir(args.rundir)
+        search = Search.load_rundir(args.rundir)
 
         async def extract() -> None:
             for detection in search._detections.detections:
                 await search.add_features(detection)
 
         asyncio.run(extract())
 
-    elif args.command == "station-corrections":
+    elif args.command == "corrections":
         rundir = Path(args.rundir)
         station_corrections = StationCorrections(rundir=rundir)
         if args.plot:
             station_corrections.save_plots(rundir / "station_corrections")
         station_corrections.save_csv(filename=rundir / "station_corrections_stats.csv")
 
     elif args.command == "serve":
-        search = SquirrelSearch.load_rundir(args.rundir)
+        search = Search.load_rundir(args.rundir)
         webserver = WebServer(search)
 
         loop = asyncio.get_event_loop()
         loop.create_task(webserver.start())
         loop.run_forever()
 
+    elif args.command == "clear-cache":
+        logger.info("clearing cache directory %s", CACHE_DIR)
+        shutil.rmtree(CACHE_DIR)
+
     elif args.command == "dump-schemas":
         from lassie.models.detection import EventDetections
 
@@ -198,7 +196,7 @@ async def extract() -> None:
 
         file = args.folder / "search.schema.json"
         print(f"writing JSON schemas to {args.folder}")
-        file.write_text(SquirrelSearch.model_json_schema(indent=2))
+        file.write_text(Search.model_json_schema(indent=2))
 
         file = args.folder / "detections.schema.json"
         file.write_text(EventDetections.model_json_schema(indent=2))

diff --git a/lassie/features/local_magnitude.py b/lassie/features/local_magnitude.py
@@ -254,4 +254,4 @@ async def add_features(self, squirrel: Squirrel, event: EventDetection) -> None:
         print(event.time, local_magnitude)
         event.magnitude = local_magnitude.median
         event.magnitude_type = "local"
-        event.feature.add_feature(local_magnitude)
+        event.features.add_feature(local_magnitude)
diff --git a/lassie/images/__init__.py b/lassie/images/__init__.py
@@ -2,7 +2,8 @@
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Annotated, Iterator, Union
+from itertools import chain
+from typing import TYPE_CHECKING, Annotated, Any, Iterator, Union
 
 from pydantic import Field, RootModel
 
@@ -34,7 +35,13 @@
 
 
 class ImageFunctions(RootModel):
-    root: list[ImageFunctionType] = []
+    root: list[ImageFunctionType] = [PhaseNet()]
+
+    def model_post_init(self, __context: Any) -> None:
+        # Check if phases are provided twice
+        phases = self.get_phases()
+        if len(set(phases)) != len(phases):
+            raise ValueError("A phase was provided twice")
 
     async def process_traces(self, traces: list[Trace]) -> WaveformImages:
         images = []
@@ -44,6 +51,14 @@ async def process_traces(self, traces: list[Trace]) -> WaveformImages:
 
         return WaveformImages(root=images)
 
+    def get_phases(self) -> tuple[str, ...]:
+        """Get all phases that are available in the image functions.
+
+        Returns:
+            tuple[str, ...]: All available phases.
+        """
+        return tuple(chain.from_iterable(image.get_provided_phases() for image in self))
+
     def get_blinding(self) -> timedelta:
         return max(image.blinding for image in self)
 

diff --git a/lassie/images/base.py b/lassie/images/base.py
@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from lassie.models.phase_arrival import PhaseArrival
 from lassie.models.station import Stations
@@ -35,7 +35,7 @@ def blinding(self) -> timedelta:
         """Blinding duration for the image function. Added to padded waveforms."""
         raise NotImplementedError("must be implemented by subclass")
 
-    def get_available_phases(self) -> tuple[str]:
+    def get_provided_phases(self) -> tuple[str, ...]:
         ...
 
 
@@ -45,7 +45,7 @@ class WaveformImage:
     phase: PhaseDescription
     weight: float
     traces: list[Trace]
-    stations: Stations = Stations.construct()
+    stations: Stations = Field(default_factory=lambda: Stations.model_construct())
 
     @property
     def sampling_rate(self) -> float: