bugfixes

.pre-commit-config.yaml

@@ -2,23 +2,24 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v3.2.0
+    rev: v4.5.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: check-yaml
       - id: check-added-large-files
+      - id: mixed-line-ending
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     # Ruff version.
-    rev: "v0.0.291"
+    rev: "v0.1.1"
     hooks:
       - id: ruff
   - repo: https://github.com/psf/black
-    rev: 23.9.1
+    rev: 23.10.0
     hooks:
       - id: black
         # It is recommended to specify the latest version of Python
         # supported by your project here, or alternatively use
         # pre-commit's default_language_version, see
         # https://pre-commit.com/#top_level-default_language_version
-        # language_version: python3.9
+        # language_version: python3.10

lassie/models/detection.py

@@ -366,8 +366,7 @@ def as_pyrocko_event(self) -> Event:
             lon=self.lon,
             east_shift=self.east_shift,
             north_shift=self.north_shift,
-            depth=self.depth,
-            elevation=self.elevation,
+            depth=self.effective_depth,
             magnitude=self.magnitude or self.semblance,
             magnitude_type=self.magnitude_type,
         )
@@ -535,7 +534,7 @@ def save_csv(self, file: Path, jitter_location: float = 0.0) -> None:
                 detection = detection.jitter_location(jitter_location)
             lat, lon = detection.effective_lat_lon
             lines.append(
-                f"{lat:.5f}, {lon:.5f}, {-detection.effective_elevation:.1f},"
+                f"{lat:.5f}, {lon:.5f}, {detection.effective_depth:.1f},"
                 f" {detection.semblance}, {detection.time}, {detection.distance_border}"
             )
         file.write_text("\n".join(lines))

lassie/models/location.py

@@ -5,7 +5,7 @@
 import struct
 from typing import TYPE_CHECKING, Iterable, Literal, TypeVar
 
-from pydantic import BaseModel, PrivateAttr
+from pydantic import BaseModel, Field, PrivateAttr
 from pyrocko import orthodrome as od
 from typing_extensions import Self
 
@@ -18,10 +18,22 @@
 class Location(BaseModel):
     lat: float
     lon: float
-    east_shift: float = 0.0
-    north_shift: float = 0.0
-    elevation: float = 0.0
-    depth: float = 0.0
+    east_shift: float = Field(
+        default=0.0,
+        description="east shift towards geographical reference in meters.",
+    )
+    north_shift: float = Field(
+        default=0.0,
+        description="north shift towards geographical reference in meters.",
+    )
+    elevation: float = Field(
+        default=0.0,
+        description="elevation in meters.",
+    )
+    depth: float = Field(
+        default=0.0,
+        description="depth in meters, positive is down.",
+    )
 
     _cached_lat_lon: tuple[float, float] | None = PrivateAttr(None)
 

lassie/search.py

@@ -253,7 +253,7 @@ async def start(self, force_rundir: bool = False) -> None:
                 self._detections.add(detection)
                 await self._new_detection.emit(detection)
 
-            if batch.i_batch % 50 == 0:
+            if self._detections.n_detections % 50 == 0:
                 self._detections.dump_detections(jitter_location=self.octree.size_limit)
 
             processing_time = datetime_now() - batch_processing_start
@@ -283,7 +283,9 @@ async def start(self, force_rundir: bool = False) -> None:
             batch_processing_start = datetime_now()
             self.set_progress(batch.end_time)
 
+        self._detections.dump_detections(jitter_location=self.octree.size_limit)
         logger.info("finished search in %s", datetime_now() - processing_start)
+        logger.info("found %d detections", self._detections.n_detections)
 
     async def add_features(self, event: EventDetection) -> None:
         try:
@@ -326,6 +328,7 @@ def from_config(
         return model
 
     def __del__(self) -> None:
+        # FIXME: Replace with signal overserver?
         if hasattr(self, "_detections"):
             with contextlib.suppress(Exception):
                 self._detections.dump_detections(jitter_location=self.octree.size_limit)
@@ -379,10 +382,10 @@ async def calculate_semblance(
 
         traveltimes_bad = np.isnan(traveltimes)
         traveltimes[traveltimes_bad] = 0.0
-        station_contribution = (~traveltimes_bad).sum(axis=1, dtype=float)
+        station_contribution = (~traveltimes_bad).sum(axis=1, dtype=np.float32)
 
         shifts = np.round(-traveltimes / image.delta_t).astype(np.int32)
-        weights = np.full_like(shifts, fill_value=image.weight, dtype=float)
+        weights = np.full_like(shifts, fill_value=image.weight, dtype=np.float32)
 
         # Normalize by number of station contribution
         with np.errstate(divide="ignore", invalid="ignore"):

lassie/tracers/cake.py

@@ -2,6 +2,7 @@
 
 import logging
 import re
+import struct
 import zipfile
 from datetime import datetime, timedelta
 from functools import cached_property
@@ -91,21 +92,21 @@ class CakeArrival(ModelledArrival):
 
 class EarthModel(BaseModel):
     filename: FilePath | None = Field(
-        DEFAULT_VELOCITY_MODEL_FILE,
+        default=DEFAULT_VELOCITY_MODEL_FILE,
         description="Path to velocity model.",
     )
     format: Literal["nd", "hyposat"] = Field(
-        "nd",
+        default="nd",
         description="Format of the velocity model. nd or hyposat is supported.",
     )
     crust2_profile: constr(to_upper=True) | tuple[float, float] = Field(
-        "",
+        default="",
         description="Crust2 profile name or a tuple of (lat, lon) coordinates.",
     )
 
     raw_file_data: str | None = Field(
-        None,
-        description="Raw .nd file data.",
+        default=None,
+        description="Raw `.nd` file data.",
     )
     _layered_model: LayeredModel = PrivateAttr()
 
@@ -173,6 +174,27 @@ def id(self) -> str:
         return re.sub(r"[\,\s\;]", "", self.definition)
 
 
+def surface_distances(nodes: Sequence[Node], stations: Stations) -> np.ndarray:
+    """Returns the surface distance from all nodes to all stations.
+
+    Args:
+        nodes (Sequence[Node]): Nodes to calculate distance from.
+        stations (Stations): Stations to calculate distance to.
+
+    Returns:
+        np.ndarray: Distances in shape (n-nodes, n-stations).
+    """
+    node_coords = get_node_coordinates(nodes, system="geographic")
+    n_nodes = node_coords.shape[0]
+
+    node_coords = np.repeat(node_coords, stations.n_stations, axis=0)
+    sta_coords = np.vstack(n_nodes * [stations.get_coordinates(system="geographic")])
+
+    return od.distance_accurate50m_numpy(
+        node_coords[:, 0], node_coords[:, 1], sta_coords[:, 0], sta_coords[:, 1]
+    ).reshape(-1, stations.n_stations)
+
+
 class TravelTimeTree(BaseModel):
     earthmodel: EarthModel
     timing: Timing
@@ -234,11 +256,11 @@ def is_suited(
         time_tolerance: float,
         spatial_tolerance: float,
     ) -> bool:
-        def check_bounds(self, requested) -> bool:
+        def check_bounds(self, requested: tuple[float, float]) -> bool:
             return self[0] <= requested[0] and self[1] >= requested[1]
 
         return (
-            str(self.earthmodel.layered_model) == str(earthmodel.layered_model)
+            str(self.earthmodel) == str(earthmodel.hash)
             and self.timing == timing
             and check_bounds(self.distance_bounds, distance_bounds)
             and check_bounds(self.source_depth_bounds, source_depth_bounds)
@@ -249,7 +271,18 @@ def check_bounds(self, requested) -> bool:
 
     @property
     def filename(self) -> Path:
-        return Path(f"{self.timing.id}-{self.earthmodel.hash}.sptree")
+        hash = sha1(self.earthmodel.hash.encode())
+        hash.update(
+            struct.pack(
+                "dddddddd",
+                *self.distance_bounds,
+                *self.source_depth_bounds,
+                *self.receiver_depth_bounds,
+                self.time_tolerance,
+                self.spatial_tolerance,
+            )
+        )
+        return Path(f"{self.timing.id}-{hash.hexdigest()}.sptree")
 
     @classmethod
     def new(cls, **data) -> Self:
@@ -302,7 +335,6 @@ def load(cls, file: Path) -> Self:
         with zipfile.ZipFile(file, "r") as archive:
             path = zipfile.Path(archive)
             model_file = path / "model.json"
-            print(model_file.read_text())
             model = cls.model_validate_json(model_file.read_text())
         model._file = file
         return model
@@ -345,38 +377,34 @@ def init_lut(self, octree: Octree, stations: Stations) -> None:
         for node, traveltimes in zip(octree, station_traveltimes, strict=True):
             self._node_lut[node.hash()] = traveltimes.astype(np.float32)
 
+    def lut_fill_level(self) -> float:
+        """Return the fill level of the LUT as a float between 0.0 and 1.0"""
+        return len(self._node_lut) / self._node_lut.get_size()
+
     def fill_lut(self, nodes: Sequence[Node]) -> None:
         logger.debug("filling traveltimes LUT for %d nodes", len(nodes))
         stations = self._cached_stations
 
-        node_coords = get_node_coordinates(nodes, system="geographic")
-        sta_coords = stations.get_coordinates(system="geographic")
-
-        sta_coords = np.array(od.geodetic_to_ecef(*sta_coords.T)).T
-        node_coords = np.array(od.geodetic_to_ecef(*node_coords.T)).T
-
-        receiver_distances = np.linalg.norm(
-            sta_coords - node_coords[:, np.newaxis], axis=2
-        )
-
         traveltimes = self._interpolate_travel_times(
-            receiver_distances,
+            surface_distances(nodes, stations),
             np.array([sta.effective_depth for sta in stations]),
-            np.array([node.depth for node in nodes]),
+            np.array([node.as_location().effective_depth for node in nodes]),
         )
 
         for node, times in zip(nodes, traveltimes, strict=True):
             self._node_lut[node.hash()] = times.astype(np.float32)
 
-    def lut_fill_level(self) -> float:
-        """Return the fill level of the LUT as a float between 0.0 and 1.0"""
-        return len(self._node_lut) / self._node_lut.get_size()
-
     def get_travel_times(self, octree: Octree, stations: Stations) -> np.ndarray:
-        station_indices = np.fromiter(
-            (self._cached_station_indeces[sta.pretty_nsl] for sta in stations),
-            dtype=int,
-        )
+        try:
+            station_indices = np.fromiter(
+                (self._cached_station_indeces[sta.pretty_nsl] for sta in stations),
+                dtype=int,
+            )
+        except KeyError as exc:
+            raise ValueError(
+                "stations not found in cached stations, "
+                "was the LUT initialized with `TravelTimeTree.init_lut`?"
+            ) from exc
 
         stations_traveltimes = []
         fill_nodes = []
@@ -407,9 +435,11 @@ def interpolate_travel_times(
         octree: Octree,
         stations: Stations,
     ) -> np.ndarray:
-        receiver_distances = octree.distances_stations(stations)
+        receiver_distances = surface_distances(octree, stations)
         receiver_depths = np.array([sta.effective_depth for sta in stations])
-        source_depths = np.array([node.depth for node in octree])
+        source_depths = np.array(
+            [node.as_location().effective_depth for node in octree]
+        )
 
         return self._interpolate_travel_times(
             receiver_distances, receiver_depths, source_depths
@@ -452,7 +482,7 @@ def get_travel_time(self, source: Location, receiver: Location) -> float:
         coordinates = [
             receiver.effective_depth,
             source.effective_depth,
-            receiver.distance_to(source),
+            receiver.surface_distance_to(source),
         ]
         try:
             traveltime = self._get_sptree().interpolate(coordinates) or np.nan
@@ -521,16 +551,15 @@ async def prepare(self, octree: Octree, stations: Stations) -> None:
         ]
         logger.debug("loaded %d cached travel time trees", len(cached_trees))
 
-        distances = octree.distances_stations(stations)
-        source_depths = np.asarray(octree.depth_bounds)
+        distances = surface_distances(octree, stations)
+        source_depths = np.asarray(octree.depth_bounds) - octree.reference.elevation
         receiver_depths = np.fromiter((sta.effective_depth for sta in stations), float)
 
-        receiver_depths_bounds = (receiver_depths.min(), receiver_depths.max())
-        source_depth_bounds = (source_depths.min(), source_depths.max())
         distance_bounds = (distances.min(), distances.max())
+        source_depth_bounds = (source_depths.min(), source_depths.max())
+        receiver_depths_bounds = (receiver_depths.min(), receiver_depths.max())
         # FIXME: Time tolerance is too hardcoded. Is 5x a good value?
         time_tolerance = octree.smallest_node_size() / (self.get_vmin() * 5.0)
-
         # if self.trim_earth_model_depth:
         #     self.earthmodel.trim(-source_depth_bounds[1])