diff --git a/lassie/apps/lassie.py b/lassie/apps/lassie.py index c00f7fbc..84a946f1 100644 --- a/lassie/apps/lassie.py +++ b/lassie/apps/lassie.py @@ -81,7 +81,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", ) @@ -95,14 +95,14 @@ 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") subparsers.add_parser( "clear-cache", - help="clear the cached travel times", + help="clear the cach directory", ) dump_schemas = subparsers.add_parser( @@ -178,7 +178,7 @@ async def extract() -> None: asyncio.run(extract()) - elif args.command == "station-corrections": + elif args.command == "corrections": rundir = Path(args.rundir) station_corrections = StationCorrections(rundir=rundir) if args.plot: diff --git a/lassie/models/detection.py b/lassie/models/detection.py index 541d73cf..c11c486c 100644 --- a/lassie/models/detection.py +++ b/lassie/models/detection.py @@ -20,7 +20,7 @@ from lassie.models.location import Location from lassie.models.station import Station, Stations from lassie.tracers import RayTracerArrival -from lassie.utils import PhaseDescription, time_to_path +from lassie.utils import PhaseDescription, Symbols, time_to_path if TYPE_CHECKING: from pyrocko.squirrel import Response, Squirrel @@ -449,6 +449,17 @@ def add(self, detection: EventDetection) -> None: marker.save_markers(detection.get_pyrocko_markers(), str(markers_file)) self.detections.append(detection) + logger.info( + "%s detection #%d %s: %.5f°, %.5f°, depth %.1f m, " + "border distance %.1f m, semblance %.3f", + Symbols.Target, + self.n_detections, + detection.time, + *detection.effective_lat_lon, + detection.depth, + detection.distance_border, + detection.semblance, + ) # This has to happen after the markers are saved detection.dump_append(self.rundir, self.n_detections - 1) diff --git a/lassie/models/location.py b/lassie/models/location.py index eebe15c0..be6efaa5 100644 --- a/lassie/models/location.py +++ b/lassie/models/location.py @@ -3,10 +3,11 @@ import hashlib import math import struct -from typing import TYPE_CHECKING, Iterable, Literal, Self, TypeVar +from typing import TYPE_CHECKING, Iterable, Literal, TypeVar from pydantic import BaseModel, PrivateAttr from pyrocko import orthodrome as od +from typing_extensions import Self if TYPE_CHECKING: from pathlib import Path diff --git a/lassie/search/base.py b/lassie/search/base.py index a63a78e2..69624c6f 100644 --- a/lassie/search/base.py +++ b/lassie/search/base.py @@ -22,7 +22,7 @@ from lassie.station_corrections import StationCorrections from lassie.tracers import CakeTracer, ConstantVelocityTracer, RayTracers from lassie.tracers.fast_marching import FastMarchingTracer -from lassie.utils import PhaseDescription, Symbols, alog_call, time_to_path +from lassie.utils import PhaseDescription, alog_call, time_to_path if TYPE_CHECKING: from pyrocko.trace import Trace @@ -422,11 +422,6 @@ async def search( node_idx = (await semblance.maxima_node_idx())[time_idx] source_node = octree[node_idx] - if not octree.is_node_in_bounds(source_node): - logger.info( - "source node is inside octree's absorbing boundary (%.1f m)", - source_node.distance_border, - ) source_location = source_node.as_location() detection = EventDetection( @@ -464,15 +459,6 @@ async def search( ) detections.append(detection) - logger.info( - "%s new detection %s: %.5fE, %.5fN, depth %.1f m, semblance %.3f", - Symbols.Target, - detection.time, - *detection.effective_lat_lon, - detection.depth, - detection.semblance, - ) - # detection.plot() # plot_octree_movie(octree, semblance, file=Path("/tmp/test.mp4")) diff --git a/lassie/search/squirrel.py b/lassie/search/squirrel.py index 629b8573..6bdaa8b7 100644 --- a/lassie/search/squirrel.py +++ b/lassie/search/squirrel.py @@ -196,7 +196,7 @@ async def scan_squirrel(self) -> None: * (batch.n - batch.i - 1) ) logger.info( - "remaining %s, estimated finish at %s", + "remaining time %s, estimated finish at %s", remaining_time, datetime.now() + remaining_time, # noqa: DTZ005 ) diff --git a/lassie/tracers/cake.py b/lassie/tracers/cake.py index 896e0bdb..4a7535b5 100644 --- a/lassie/tracers/cake.py +++ b/lassie/tracers/cake.py @@ -393,7 +393,7 @@ def get_travel_times(self, octree: Octree, stations: Stations) -> np.ndarray: cache_hits, cache_misses = self._node_lut.get_stats() cache_hit_rate = cache_hits / (cache_hits + cache_misses) - logger.info( + logger.debug( "node LUT cache fill level %.1f%%, cache hit rate %.1f%%", self.lut_fill_level() * 100, cache_hit_rate * 100, diff --git a/lassie/tracers/fast_marching/fast_marching.py b/lassie/tracers/fast_marching/fast_marching.py index f02aaac4..04a79ed1 100644 --- a/lassie/tracers/fast_marching/fast_marching.py +++ b/lassie/tracers/fast_marching/fast_marching.py @@ -8,7 +8,7 @@ from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, Self, Sequence +from typing import TYPE_CHECKING, Any, Literal, Sequence import numpy as np from lru import LRU @@ -16,6 +16,7 @@ from pyrocko.modelling import eikonal from rich.progress import Progress from scipy.interpolate import RegularGridInterpolator +from typing_extensions import Self from lassie.models.location import Location from lassie.models.station import Station, Stations @@ -163,7 +164,7 @@ def filename(self) -> str: # TODO: Add origin to hash to avoid collisions return f"{self.station.pretty_nsl}-{self.velocity_model_hash}.3dtt" - def get_traveltime_interpolator(self) -> RegularGridInterpolator: + def get_travel_time_interpolator(self) -> RegularGridInterpolator: if self._interpolator is None: self._interpolator = RegularGridInterpolator( (self._east_coords, self._north_coords, self._depth_coords), @@ -178,7 +179,7 @@ def interpolate_travel_time( location: Location, method: Literal["nearest", "linear", "cubic"] = "linear", ) -> float: - interpolator = self.get_traveltime_interpolator() + interpolator = self.get_travel_time_interpolator() offset = location.offset_to(self.center) return interpolator([offset], method=method).astype(float, copy=False)[0] @@ -187,11 +188,19 @@ def interpolate_nodes( nodes: Sequence[Node], method: Literal["nearest", "linear", "cubic"] = "linear", ) -> np.ndarray: - interpolator = self.get_traveltime_interpolator() + interpolator = self.get_travel_time_interpolator() coordinates = [node.as_location().offset_to(self.center) for node in nodes] return interpolator(coordinates, method=method).astype(float, copy=False) + def get_meshgrid(self) -> list[np.ndarray]: + return np.meshgrid( + self._east_coords, + self._north_coords, + self._depth_coords, + indexing="ij", + ) + def save(self, path: Path) -> Path: """Save travel times to a zip file. @@ -221,10 +230,10 @@ def save(self, path: Path) -> Path: @classmethod def load(cls, file: Path) -> Self: - """Load 3D travel times from a zip file. + """Load 3D travel times from a .3dtt file. Args: - file (Path): path to the zip file containing the travel times + file (Path): path to the .3dtt file containing the travel times Returns: Self: 3D travel times @@ -441,7 +450,7 @@ def get_travel_times( cache_hits, cache_misses = self._node_lut.get_stats() cache_hit_rate = cache_hits / (cache_hits + cache_misses) - logger.info( + logger.debug( "node LUT cache fill level %.1f%%, cache hit rate %.1f%%", self.lut_fill_level() * 100, cache_hit_rate * 100, @@ -456,7 +465,7 @@ def fill_lut(self, nodes: Sequence[Node]) -> None: with Progress() as progress: status = progress.add_task( - f"interpolating {self.phase} traveltimes for {n_nodes} nodes", + f"interpolating {self.phase} travel times for {n_nodes} nodes", total=self._cached_stations.n_stations, ) for station in self._cached_stations: diff --git a/lassie/tracers/fast_marching/velocity_models.py b/lassie/tracers/fast_marching/velocity_models.py index 0f2ccc17..ac261592 100644 --- a/lassie/tracers/fast_marching/velocity_models.py +++ b/lassie/tracers/fast_marching/velocity_models.py @@ -4,7 +4,7 @@ import re from hashlib import sha1 from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Any, Literal, Self, Union +from typing import TYPE_CHECKING, Annotated, Any, Literal, Union import numpy as np from pydantic import ( @@ -17,6 +17,7 @@ ) from pydantic.dataclasses import dataclass from scipy.interpolate import RegularGridInterpolator +from typing_extensions import Self from lassie.models.location import Location @@ -94,7 +95,10 @@ def hash(self) -> str: return self._hash def get_source_arrival_grid(self, station: Station) -> np.ndarray: - times = np.full_like(self._velocity_model, fill_value=-1.0) + if not self.is_inside(station): + raise ValueError("Station is outside of velocity model.") + + times = np.full_like(self.velocity_model, fill_value=-1.0) station_offset = station.offset_to(self.center) east_idx = np.argmin(np.abs(self._east_coords - station_offset[0])) @@ -116,6 +120,7 @@ def get_meshgrid(self) -> list[np.ndarray]: self._east_coords, self._north_coords, self._depth_coords, + indexing="ij", ) def resample( diff --git a/test/test_fast_marching.py b/test/test_fast_marching.py index 07bd7571..12f12bd6 100644 --- a/test/test_fast_marching.py +++ b/test/test_fast_marching.py @@ -24,7 +24,10 @@ def stations_inside( - model: VelocityModel3D, nstations: int = 20, seed: int = 0 + model: VelocityModel3D, + nstations: int = 20, + seed: int = 0, + depth: float | None = None, ) -> Stations: stations = [] rng = np.random.RandomState(seed) @@ -37,7 +40,8 @@ def stations_inside( elevation=model.center.elevation, north_shift=model.center.north_shift + rng.uniform(*model.north_bounds), east_shift=model.center.east_shift + rng.uniform(*model.east_bounds), - depth=model.center.depth + rng.uniform(*model.depth_bounds), + depth=model.center.depth + + (depth if depth is not None else rng.uniform(*model.depth_bounds)), ) stations.append(station) return Stations(stations=stations) @@ -174,3 +178,45 @@ def test_non_lin_loc_model( ) fig.colorbar(cmap) plt.show() + + +@pytest.mark.plot +@pytest.mark.asyncio +async def test_non_lin_loc_travel_times(data_dir: Path, octree: Octree) -> None: + import matplotlib.pyplot as plt + + header_file = data_dir / "FORGE_3D_5_large.P.mod.hdr" + + tracer = FastMarchingTracer( + phase="fm:P", + velocity_model=NonLinLocVelocityModel( + header_file=header_file, + grid_spacing=100.0, + ), + ) + model_3d = tracer.velocity_model.get_model(octree) + octree = octree_cover(model_3d) + stations = stations_inside(model_3d, depth=0.0) + await tracer.prepare(octree, stations) + + volume = tracer.get_travel_time_volume(stations.stations[0]) + + # 3d figure of velocity model + fig = plt.figure() + ax = fig.add_subplot(projection="3d") + coords = volume.get_meshgrid() + print(coords[0].shape) + + cmap = ax.scatter( + coords[0], + coords[1], + coords[2], + c=volume.travel_times.ravel(), + alpha=0.2, + ) + + station_offet = volume.station.offset_to(volume.center) + print(station_offet) + ax.scatter(*station_offet, s=100, c="r") + fig.colorbar(cmap) + plt.show()