wip

Author: Marius Isken

src/qseek/search.py

@@ -24,7 +24,6 @@
 
 from qseek.cache_lru import CACHES
 from qseek.corrections.corrections import StationCorrectionType, corrections_from_path
-from qseek.distance_weights import DistanceWeights
 from qseek.features import FeatureExtractorType
 from qseek.images.images import ImageFunctions, WaveformImages
 from qseek.magnitudes import EventMagnitudeCalculatorType
@@ -38,14 +37,15 @@
 from qseek.pre_processing.module import Downsample, PreProcessing
 from qseek.reduce import DelaySumReduce
 from qseek.signals import Signal
+from qseek.station_weights import StationWeights
 from qseek.stats import RuntimeStats, Stats
 from qseek.tracers.tracers import RayTracer, RayTracers
 from qseek.utils import (
     BackgroundTasks,
     CpuCount,
     PhaseDescription,
+    _get_cpu_count,
     datetime_now,
-    get_cpu_count,
     get_total_memory,
     human_readable_bytes,
     time_to_path,
@@ -265,10 +265,11 @@ class Search(BaseModel):
         ),
         description="List of ray tracers for travel time calculation.",
     )
-    distance_weights: DistanceWeights | None = Field(
-        default_factory=DistanceWeights,
+    station_weights: StationWeights | None = Field(
+        default_factory=StationWeights,
         validation_alias=AliasChoices("spatial_weights", "distance_weights"),
-        description="Spatial weights for distance weighting.",
+        description="Station weighting based on station density and "
+        "source-station distance.",
     )
     station_corrections: StationCorrectionType | None = Field(
         default=None,
@@ -357,7 +358,7 @@ class Search(BaseModel):
     _rundir: Path = PrivateAttr()
 
     _compute_semaphore: asyncio.Semaphore = PrivateAttr(
-        asyncio.Semaphore(max(1, get_cpu_count() - 4))
+        asyncio.Semaphore(max(1, _get_cpu_count() - 4))
     )
 
     # Signals
@@ -493,8 +494,8 @@ async def prepare(self) -> None:
         await self.pre_processing.prepare()
         await self.image_functions.prepare()
 
-        if self.distance_weights:
-            self.distance_weights.prepare(self.stations, self.octree)
+        if self.station_weights:
+            self.station_weights.prepare(self.stations, self.octree)
 
         if self.station_corrections:
             await self.station_corrections.prepare(
@@ -552,7 +553,7 @@ async def start(self, force_rundir: bool = False) -> None:
             ray_tracers=self.ray_tracers,
             window_padding=window_padding,
             station_corrections=self.station_corrections,
-            distance_weights=self.distance_weights,
+            distance_weights=self.station_weights,
             detection_threshold=self.detection_threshold,
             pick_confidence_threshold=self.pick_confidence_threshold,
             node_interpolation=self.node_interpolation,
@@ -698,7 +699,7 @@ def __init__(
         ray_tracers: RayTracers,
         window_padding: timedelta,
         station_corrections: StationCorrectionType | None = None,
-        distance_weights: DistanceWeights | None = None,
+        distance_weights: StationWeights | None = None,
         detection_threshold: float | Literal["MAD"] = "MAD",
         detection_blinding: timedelta = timedelta(seconds=1.0),
         pick_confidence_threshold: float = 0.3,

src/qseek/distance_weights.py src/qseek/station_weights.pysrc/qseek/distance_weights.py renamed to src/qseek/station_weights.py

@@ -210,7 +210,7 @@ def weights_gaussian(
     return weights
 
 
-class DistanceWeights(BaseModel):
+class StationWeights(BaseModel):
     plateau_weight: PositiveFloat = Field(
         default=4.0,
         description="The cumulative station weight required to define the"
@@ -256,7 +256,7 @@ def prepare(self, stations: StationInventory, octree: Octree) -> None:
         logger.info("calculating overall station weights")
         station_weights = _station_weights(self._interstation_distances)
         for sta, weight in zip(stations, station_weights, strict=True):
-            sta.set_apparent_weight(weight)
+            sta.set_apparent_weight(float(weight))
 
         self.fill_lut(nodes=octree.nodes)
 

src/qseek/tracers/fast_marching.py

@@ -22,7 +22,7 @@
 from qseek.stats import get_progress
 from qseek.tracers.base import ModelledArrival, RayTracer
 from qseek.tracers.utils import LayeredEarthModel1D, surface_distances_reference
-from qseek.utils import Range, alog_call, datetime_now, get_cpu_count
+from qseek.utils import Range, _get_cpu_count, alog_call, datetime_now
 
 if TYPE_CHECKING:
     from qseek.models.location import Location
@@ -321,7 +321,7 @@ async def _calculate_travel_times(self) -> None:
         if self._layered_model is None:
             raise ValueError("layered model must be set for FastMarchingTracer")
 
-        nthreads = self.nthreads if self.nthreads > 0 else get_cpu_count() * 2
+        nthreads = self.nthreads if self.nthreads > 0 else _get_cpu_count() * 2
         executor = ThreadPoolExecutor(
             max_workers=nthreads,
             thread_name_prefix="qseek-fmm",

src/qseek_inversion/inversion.py

@@ -65,8 +65,8 @@ async def prepare(self) -> None:
         await self.event_selection.prepare(rundir=self.import_rundir)
         # await search.prepare()
         self._search = search
-        if search.distance_weights:
-            search.distance_weights.prepare(
+        if search.station_weights:
+            search.station_weights.prepare(
                 self.event_selection.stations,
                 search.octree,
             )
@@ -84,7 +84,7 @@ async def prepare(self) -> None:
             detection_blinding=search.detection_blinding,
             pick_confidence_threshold=search.pick_confidence_threshold,
             station_corrections=search.station_corrections,
-            distance_weights=search.distance_weights,
+            distance_weights=search.station_weights,
             ignore_boundary=search.ignore_boundary,
             ignore_boundary_width=search.ignore_boundary_width,
             node_interpolation=False,