fixes

lassie/search.py

@@ -3,11 +3,11 @@
 import asyncio
 import contextlib
 import logging
+import os
 from collections import deque
 from copy import deepcopy
 from datetime import datetime, timedelta, timezone
 from itertools import chain
-import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Deque, Literal
 
@@ -98,19 +98,22 @@ def processed_percent(self) -> float:
             return 0.0
         return self.batch_count / self.batch_count_total * 100.0
 
+    def reset_start_time(self) -> None:
+        self._search_start = datetime_now()
+
     def add_processed_batch(
         self,
         batch: WaveformBatch,
         duration: timedelta,
-        log: bool = False,
+        show_log: bool = False,
     ) -> None:
         self.batch_count = batch.i_batch
         self.batch_count_total = batch.n_batches
         self.batch_time = batch.end_time
         self._batch_processing_times.append(duration)
         self.processing_rate_bytes = batch.cumulative_bytes / duration.total_seconds()
         self.processing_rate_time = batch.duration / duration.total_seconds()
-        if log:
+        if show_log:
             self.log()
 
     def log(self) -> None:
@@ -368,6 +371,7 @@ async def start(self, force_rundir: bool = False) -> None:
 
         logger.info("starting search...")
         stats = self._stats
+        stats.reset_start_time()
 
         processing_start = datetime_now()
 
@@ -411,7 +415,7 @@ async def start(self, force_rundir: bool = False) -> None:
             stats.add_processed_batch(
                 batch,
                 duration=datetime_now() - batch_processing_start,
-                log=True,
+                show_log=True,
             )
 
             self.set_progress(batch.end_time)
@@ -535,7 +539,9 @@ async def calculate_semblance(
         if cache_mask is not None:
             weights[cache_mask] = 0.0
 
-        threads = parent.n_threads_argmax or max(1, os.cpu_count() - 4)  # Hold threads back for I/O
+        threads = parent.n_threads_argmax or max(
+            1, os.cpu_count() - 4
+        )  # Hold threads back for I/O
         semblance_data, offsets = await asyncio.to_thread(
             parstack.parstack,
             arrays=image.get_trace_data(),

lassie/utils.py

@@ -7,7 +7,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Annotated, Awaitable, Callable, ParamSpec, TypeVar
 
-from pydantic import BaseModel, constr
+from pydantic import BaseModel, ByteSize, constr
 from pyrocko.util import UnavailableDecimation
 from rich.logging import RichHandler
 
@@ -99,11 +99,7 @@ async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
 
 def human_readable_bytes(size: int | float) -> str:
     """Return a human readable string representation of bytes"""
-    for unit in ["B", "KiB", "MiB", "GiB", "TiB"]:
-        if size < 1024.0:
-            return f"{size:.2f} {unit}"
-        size = size / 1024.0
-    return f"{size:.2f} PiB"
+    return ByteSize(size).human_readable()
 
 
 def datetime_now() -> datetime: