groundtrack 0.1.0__py3-none-any.whl

groundtrack/__init__.py

@@ -0,0 +1,23 @@
+"""
+groundtrack
+
+Tools for discovering seismic signals from atmospheric re-entry events
+using orbital ground tracks and seismic station data.
+"""
+
+__version__ = "0.1.0"
+
+# Allow user to import below functions directly from the package
+from .download import download_boxes 
+from .stations import load_and_filter_stations
+from .tiling import track_to_box_windows, box_windows_to_download_requests
+from .processing import process_boxes, process_box, process_stream
+from .pipeline import run_pipeline
+
+from .plotting import (
+    plot_track_and_boxes,
+    plot_stations,
+    plot_station_comparison,
+    plot_waveform_comparison,
+    plot_all_waveforms,
+)

groundtrack/download.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Iterable, Sequence
+
+from obspy import UTCDateTime
+from obspy.clients.fdsn import Client
+from obspy.clients.fdsn.header import FDSNNoDataException
+from obspy.clients.fdsn.mass_downloader import (
+    MassDownloader,
+    RectangularDomain,
+    Restrictions,
+)
+
+from .stations import (
+    parse_stationxml_files,
+    deduplicate_stations,
+    filter_stations_by_track_distance,
+)
+
+from .processing import (
+    process_box,
+    DEFAULT_PRE_FILT_LOW,
+    DEFAULT_WATER_LEVEL,
+    DEFAULT_OUTPUT,
+    DEFAULT_TAPER_PCT,
+    DEFAULT_FREQMIN,
+    DEFAULT_FREQMAX,
+    DEFAULT_CORNERS,
+    DEFAULT_ZEROPHASE,
+)
+
+
+def _normalize_providers(providers):
+    """
+    Make sure providers are always in a clean list format.
+
+    Why:
+    ObsPy expects provider names as strings. This avoids weird behavior if
+    something else gets passed in accidentally.
+    """
+    if providers is None:
+        return []
+    return [str(p) for p in providers]
+
+
+def _count_files(path: Path, pattern: str) -> int:
+    """
+    Count files recursively underneath a folder.
+
+    Why:
+    We use this for quick bookkeeping, like:
+        - checking whether a box already has downloaded content
+        - recording how many files ended up on disk
+    """
+    return sum(1 for p in path.rglob(pattern) if p.is_file())
+
+
+def _provider_station_query(
+    client: Client,
+    req: dict,
+    channel_priorities: Sequence[str],
+    level: str = "station",
+):
+    """
+    Query one provider for station inventory inside one box.
+
+    Why:
+    This is the coarse discovery step. We are asking:
+        "What stations exist in this box and this time window?"
+    """
+    return client.get_stations(
+        starttime=UTCDateTime(req["t_start_utc"]),
+        endtime=UTCDateTime(req["t_end_utc"]),
+        minlatitude=req["lat_min"],
+        maxlatitude=req["lat_max"],
+        minlongitude=req["lon_min"],
+        maxlongitude=req["lon_max"],
+        channel=",".join(channel_priorities),
+        level=level,
+    )
+
+
+def _make_mseed_storage(approved_stations: set[tuple[str, str]], wav_dir: Path):
+    """
+    Build a callable for MassDownloader's mseed_storage parameter.
+
+    MassDownloader calls this once per channel/time-interval it considers
+    downloading. Returning a file path means "download here"; returning
+    None means "skip this station entirely."
+
+    We use this to enforce the corridor distance filter: only stations
+    that passed Phase 1 filtering are in approved_stations.
+    """
+    def storage(network, station, location, channel, starttime, endtime):
+        if (network, station) not in approved_stations:
+            return None
+        loc = location if location else ""
+        return str(wav_dir / f"{network}.{station}.{loc}.{channel}.mseed")
+
+    return storage
+
+
+def download_boxes(
+    download_requests: Iterable[dict],
+    track_points,
+    output_base: str | Path,
+    event_name: str,
+    corridor_km: float = 100.0,
+    providers: Sequence[str] | None = ("EARTHSCOPE",),
+    channel_priorities: Sequence[str] = ("HHZ", "BHZ"),
+    location_priorities: Sequence[str] = ("", "00", "10", "20"),
+    overwrite_existing: bool = False,
+    verbose: bool = True,
+
+   # ------------------------------------------------------------------
+    # Optional post-download processing
+    # ------------------------------------------------------------------
+    # When apply_processing is False (default), raw data is downloaded and
+    # nothing else happens -- this preserves the original behavior. When it
+    # is True, each box is processed immediately after its waveforms land
+    # on disk, using the parameters below. All defaults match process_boxes()
+    # so the one-shot and two-step workflows produce identical output.
+    apply_processing: bool = False,
+    # Pre-processing (before response removal)
+    demean: bool = True,
+    detrend_linear: bool = True,
+    taper_pct: float = DEFAULT_TAPER_PCT,
+    # Response removal
+    output: str = DEFAULT_OUTPUT,
+    pre_filt_low: tuple[float, float] | None = DEFAULT_PRE_FILT_LOW,
+    water_level: float = DEFAULT_WATER_LEVEL,
+    # Bandpass (after response removal)
+    apply_bandpass: bool = True,
+    freqmin: float = DEFAULT_FREQMIN,
+    freqmax: float = DEFAULT_FREQMAX,
+    corners: int = DEFAULT_CORNERS,
+    zerophase: bool = DEFAULT_ZEROPHASE
+) -> dict:
+    """
+    Main box-based download pipeline.
+
+    High-level logic:
+    -----------------
+    For each box:
+        1. Query candidate station inventory inside the box
+        2. Parse and deduplicate stations
+        3. Filter stations by distance to the propagated track
+        4. Download waveforms only for stations that survive the filter
+        5. Save a manifest describing what happened
+
+    This is the box -> candidate inventory -> distance filter -> waveform
+    download architecture we discussed.
+    """
+    requests = list(download_requests)
+
+    run_folder = Path(output_base) / event_name
+    boxes_root = run_folder / "boxes"
+    boxes_root.mkdir(parents=True, exist_ok=True)
+
+    provider_names = _normalize_providers(providers)
+
+    # Initialize provider clients once up front.
+    clients = {}
+    for provider_name in provider_names:
+        try:
+            clients[provider_name] = Client(provider_name)
+        except Exception as e:
+            if verbose:
+                print(f"Could not initialize provider {provider_name}: {repr(e)}")
+
+    mdl = MassDownloader(providers=provider_names)  # One instance used for every box
+
+    results = []
+    n_ok = 0
+    n_skip = 0
+    n_fail = 0
+
+    for k, req in enumerate(requests, start=1):
+        box_id = req["box_id"]
+
+        box_dir = boxes_root / box_id
+        inv_dir = box_dir / "stations"
+        wav_dir = box_dir / "waveforms"
+        log_dir = box_dir / "logs"
+
+        inv_dir.mkdir(parents=True, exist_ok=True)
+        wav_dir.mkdir(parents=True, exist_ok=True)
+        log_dir.mkdir(parents=True, exist_ok=True)
+
+        existing_mseed = _count_files(wav_dir, "*.mseed")
+        existing_xml = _count_files(inv_dir, "*.xml")
+
+        if not overwrite_existing and existing_mseed > 0 and existing_xml > 0:
+            n_skip += 1
+            results.append(
+                {
+                    "box_id": box_id,
+                    "status": "skipped_existing",
+                    "existing_mseed_files": existing_mseed,
+                    "existing_stationxml_files": existing_xml,
+                }
+            )
+
+            if verbose:
+                print(f"[{k}/{len(requests)}] skipping {box_id} (already has data)")
+
+            continue
+
+        if verbose:
+            print(f"[{k}/{len(requests)}] processing {box_id} ...")
+
+        # ------------------------------------------------------------
+        # Stage 1: gather candidate inventory files from providers
+        # ------------------------------------------------------------
+        provider_stationxml_files = []
+
+        for provider_name, client in clients.items():
+            try:
+                inv = _provider_station_query(
+                    client=client,
+                    req=req,
+                    channel_priorities=channel_priorities,
+                    level="station",
+                )
+
+                xml_path = inv_dir / f"{provider_name}_stations.xml"
+                inv.write(str(xml_path), format="STATIONXML")
+                provider_stationxml_files.append(xml_path)
+
+                if verbose:
+                    print(f"    inventory from {provider_name}: saved {xml_path.name}")
+
+            except FDSNNoDataException:
+                if verbose:
+                    print(f"    inventory from {provider_name}: no data")
+            except Exception as e:
+                if verbose:
+                    print(f"    inventory from {provider_name}: FAILED -> {repr(e)}")
+
+        # Parse raw candidate stations from the StationXML files we just saved.
+        station_rows = parse_stationxml_files(provider_stationxml_files)
+        unique_stations = deduplicate_stations(station_rows)
+
+        # ------------------------------------------------------------
+        # Stage 2: keep only stations within the physical corridor
+        # ------------------------------------------------------------
+        kept_stations = filter_stations_by_track_distance(
+            unique_stations,
+            track_points=track_points,
+            corridor_km=corridor_km,
+        )
+
+        if verbose:
+            print(
+                f"    candidate stations: {len(unique_stations)} | "
+                f"kept within {corridor_km} km: {len(kept_stations)}"
+            )
+
+        # Save a simple filtered station summary for debugging / inspection.
+        filtered_summary_path = log_dir / "filtered_stations.json"
+        with open(filtered_summary_path, "w", encoding="utf-8") as f:
+            json.dump(kept_stations, f, indent=2, default=str)
+
+
+        # ------------------------------------------------------------
+        # Stage 3: bulk waveform download via MassDownloader
+        # ------------------------------------------------------------
+        approved = {(s["network"], s["station"]) for s in kept_stations}
+
+        if approved:
+            domain = RectangularDomain(
+                minlatitude=req["lat_min"],
+                maxlatitude=req["lat_max"],
+                minlongitude=req["lon_min"],
+                maxlongitude=req["lon_max"],
+            )
+
+            restrictions = Restrictions(
+                starttime=UTCDateTime(req["t_start_utc"]),
+                endtime=UTCDateTime(req["t_end_utc"]),
+                network=",".join(sorted({s["network"] for s in kept_stations})),
+                station=",".join(sorted({s["station"] for s in kept_stations})),
+                channel_priorities=list(channel_priorities),
+                location_priorities=list(location_priorities),
+                reject_channels_with_gaps=True,
+                minimum_length=0.9,
+            )
+
+            try:
+                mdl.download(
+                    domain,
+                    restrictions,
+                    mseed_storage=_make_mseed_storage(approved, wav_dir),
+                    stationxml_storage=str(inv_dir / "{network}.{station}.xml"),
+                )
+            except Exception as e:
+                if verbose:
+                    print(f"    MassDownloader error: {repr(e)}")
+
+        new_mseed = _count_files(wav_dir, "*.mseed")
+        new_xml = _count_files(inv_dir, "*.xml")
+
+        box_result = {
+            "box_id": box_id,
+            "status": "ok",
+            "candidate_station_count": len(unique_stations),
+            "filtered_station_count": len(kept_stations),
+            "mseed_files": new_mseed,
+            "stationxml_files": new_xml,
+        }
+
+        # ------------------------------------------------------------
+        # Optional: process this box's waveforms right after download
+        # ------------------------------------------------------------
+        if apply_processing and new_mseed > 0:
+            if verbose:
+                print(f"    processing {box_id} ...")
+            proc_result = process_box(
+                box_dir,
+                demean=demean,
+                detrend_linear=detrend_linear,
+                taper_pct=taper_pct,
+                output=output,
+                pre_filt_low=pre_filt_low,
+                water_level=water_level,
+                apply_bandpass=apply_bandpass,
+                freqmin=freqmin,
+                freqmax=freqmax,
+                corners=corners,
+                zerophase=zerophase,
+                overwrite_existing=overwrite_existing,
+                verbose=verbose,
+            )
+            box_result["processing"] = {
+                "traces_in": proc_result["traces_in"],
+                "traces_out": proc_result["traces_out"],
+                "skipped_existing": proc_result["skipped_existing"],
+                "n_errors": len(proc_result["errors"]),
+            }        
+
+        n_ok += 1
+        results.append(box_result)
+
+    manifest = {
+        "event_name": event_name,
+        "run_folder": str(run_folder),
+        "boxes_root": str(boxes_root),
+        "total_requests": len(requests),
+        "providers": provider_names,
+        "channel_priorities": list(channel_priorities),
+        "location_priorities": list(location_priorities),
+        "corridor_km": corridor_km,
+        "ok": n_ok,
+        "skipped_existing": n_skip,
+        "failed": n_fail,
+        "processing": {
+            "applied": apply_processing,
+            "params": {
+                "demean": demean,
+                "detrend_linear": detrend_linear,
+                "taper_pct": taper_pct,
+                "output": output,
+                "pre_filt_low": list(pre_filt_low) if pre_filt_low is not None else None,
+                "water_level": water_level,
+                "apply_bandpass": apply_bandpass,
+                "freqmin": freqmin,
+                "freqmax": freqmax,
+                "corners": corners,
+                "zerophase": zerophase,
+            } if apply_processing else None,
+        },
+        "results": results,
+    }
+
+    manifest_path = run_folder / "download_manifest.json"
+    with open(manifest_path, "w", encoding="utf-8") as f:
+        json.dump(manifest, f, indent=2, default=str)
+
+    if verbose:
+        print("\n=== Download Summary ===")
+        print("OK:", n_ok, "| Skipped:", n_skip, "| Failed:", n_fail)
+        print("Manifest saved to:", manifest_path)
+
+    return manifest

groundtrack/geodesy.py

@@ -0,0 +1,140 @@
+import numpy as np
+
+from .types import TrackSegment
+
+from obspy.geodetics import (
+    locations2degrees,
+    degrees2kilometers,
+)
+
+
+def wrap_lon_deg(lon: float) -> float:
+    return (lon + 180.0) % 360.0 - 180.0
+
+
+_AMBIGUOUS_BAND_KM = 10.0
+
+def min_distance_km_to_track(
+    sta_lat: float,
+    sta_lon: float,
+    track_points,
+    corridor_km: float | None = None,
+) -> float:
+    """
+    Minimum great-circle distance in km from a station to the ground track.
+
+    For stations clearly inside or outside the corridor, returns the
+    point-sampled distance (fast, vectorized). For stations in the ambiguous
+    band near the corridor boundary, computes the true cross-track distance
+    to the one or two arc segments neighboring the nearest track point.
+
+    Args:
+        sta_lat:      Station latitude in degrees.
+        sta_lon:      Station longitude in degrees.
+        track_points: List of TrackPoint objects.
+        corridor_km:  The distance threshold being tested. Required to
+                      determine whether a station is in the ambiguous band.
+                      If None, always returns the point-sampled distance.
+    """
+    sta_lon = wrap_lon_deg(float(sta_lon))
+
+    lats = np.array([p.lat for p in track_points], dtype=float)
+    lons = np.array([wrap_lon_deg(p.lon) for p in track_points], dtype=float)
+
+    # Fast vectorized point-to-point pass
+    deg_to_points = locations2degrees(sta_lat, sta_lon, lats, lons)
+    min_km = float(degrees2kilometers(np.min(deg_to_points)))
+
+    # If no corridor given, or station is clearly inside/outside the
+    # ambiguous band, return the point-sampled answer immediately
+    if corridor_km is None:
+        return min_km
+    if min_km <= corridor_km - _AMBIGUOUS_BAND_KM:
+        return min_km
+    if min_km > corridor_km + _AMBIGUOUS_BAND_KM:
+        return min_km
+
+    # --- Ambiguous band: compute cross-track distance for the one or two
+    # segments immediately neighboring the nearest track point ---
+    nearest_idx = int(np.argmin(deg_to_points))
+    candidate_indices = set()
+    if nearest_idx > 0:
+        candidate_indices.add(nearest_idx - 1)
+    if nearest_idx < len(track_points) - 1:
+        candidate_indices.add(nearest_idx)
+
+    sta_lat_r = np.deg2rad(sta_lat)
+    sta_lon_r = np.deg2rad(sta_lon)
+
+    for i in candidate_indices:
+        seg = _compute_segment(lats, lons, i)
+
+        d_AS = np.deg2rad(deg_to_points[i])
+        if d_AS == 0.0:
+            continue
+
+        bearing_AS = _bearing_rad(seg.lat_a_rad, seg.lon_a_rad, sta_lat_r, sta_lon_r)
+
+        sin_dxt = np.sin(d_AS) * np.sin(bearing_AS - seg.bearing_ab_rad)
+        dxt_rad = np.arcsin(np.clip(sin_dxt, -1.0, 1.0))
+
+        cos_dxt = np.cos(dxt_rad)
+        dat_rad = np.arccos(np.clip(np.cos(d_AS) / max(cos_dxt, 1e-10), -1.0, 1.0))
+
+        # Only use cross-track distance if the perpendicular foot
+        # falls within the segment, not beyond either endpoint
+        if 0.0 <= dat_rad <= seg.length_rad:
+            min_km = min(min_km, float(degrees2kilometers(np.rad2deg(abs(dxt_rad)))))
+
+    return min_km
+
+
+def lon_bounds_dateline_safe(lons_wrapped):
+    """
+    Return (lon_min, lon_max), possibly with lon_min > lon_max to indicate
+    dateline crossing. Chooses the shortest longitude interval containing
+    all points.
+    """
+    lons = np.array([wrap_lon_deg(x) for x in lons_wrapped], dtype=float)
+    s = np.sort(lons)
+
+    gaps = np.diff(np.r_[s, s[0] + 360.0])
+    k = int(np.argmax(gaps))
+
+    lon_min = s[(k + 1) % len(s)]
+    lon_max = s[k]
+
+    return wrap_lon_deg(lon_min), wrap_lon_deg(lon_max)
+
+
+def _compute_segment(lats_deg, lons_deg, i: int) -> TrackSegment:
+    """
+    Compute a single TrackSegment on the fly from degree arrays.
+    Used when no precomputed segment list is available.
+    """
+    lat_a = np.deg2rad(lats_deg[i])
+    lon_a = np.deg2rad(lons_deg[i])
+    lat_b = np.deg2rad(lats_deg[i + 1])
+    lon_b = np.deg2rad(lons_deg[i + 1])
+
+    return TrackSegment(
+        lat_a_rad=lat_a,
+        lon_a_rad=lon_a,
+        lat_b_rad=lat_b,
+        lon_b_rad=lon_b,
+        bearing_ab_rad=_bearing_rad(lat_a, lon_a, lat_b, lon_b),
+        length_rad=np.deg2rad(float(locations2degrees(
+            lats_deg[i], lons_deg[i],
+            lats_deg[i + 1], lons_deg[i + 1],
+        ))),
+    )
+
+
+def _bearing_rad(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
+    """
+    Initial bearing in radians from point 1 to point 2 (all inputs in radians).
+    """
+    dlon = lon2 - lon1
+    x = np.sin(dlon) * np.cos(lat2)
+    y = np.cos(lat1) * np.sin(lat2) - np.sin(lat1) * np.cos(lat2) * np.cos(dlon)
+    return float(np.arctan2(x, y))

groundtrack/io.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+
+def ensure_dir(path: str | Path) -> Path:
+    """
+    Make sure a directory exists and return it as a Path object.
+
+    Why this exists:
+    ----------------
+    A lot of this pipeline depends on folders being present before we write
+    cache files, manifests, or downloaded data. This keeps that logic in one
+    place instead of repeating mkdir everywhere.
+    """
+    path = Path(path)
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def get_tle_cache_dir(base_cache_dir: str | Path) -> Path:
+    """
+    Return the TLE cache directory and create it if needed.
+
+    Example:
+        cache/
+            tle/
+
+    Why:
+    ----
+    TLE caching is a must so we do not keep hitting the Space-Track API and
+    risk running into request limits.
+    """
+    base_cache_dir = ensure_dir(base_cache_dir)
+    return ensure_dir(base_cache_dir / "tle")
+
+
+def get_tle_cache_path(base_cache_dir: str | Path, cache_key: str) -> Path:
+    """
+    Return the full file path for one cached TLE file.
+    """
+    tle_cache_dir = get_tle_cache_dir(base_cache_dir)
+    return tle_cache_dir / cache_key
+
+
+def get_run_folder(output_base: str | Path, event_name: str) -> Path:
+    """
+    Return the main output folder for one event/run.
+
+    Example:
+        output_base/
+            shenzhou15_20240402T084000Z_to_20240402T090000Z/
+
+    Why:
+    ----
+    This keeps event outputs grouped cleanly and makes the demo notebook
+    easier to read.
+    """
+    output_base = ensure_dir(output_base)
+    return ensure_dir(output_base / event_name)
+
+
+def write_json(path: str | Path, obj) -> Path:
+    """
+    Write a JSON artifact to disk.
+
+    Why:
+    ----
+    We use JSON for simple reproducibility artifacts like:
+        - selected TLE metadata
+        - propagation metadata
+        - request plans
+        - manifests
+    """
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(obj, f, indent=2, default=str)
+
+    return path