sgn-drift 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

{sgn_drift-0.1.0.dist-info → sgn_drift-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sgn-drift
-Version: 0.1.0
+Version: 0.1.1
 Summary: PSD Drift Data Quality Extensions for the SGN Framework
 Author-email: James Kennington <jmk7376@psu.edu>, Zach Yarbrough <jmk7376@psu.edu>, Joshua Black <jmk7376@psu.edu>
 License-Expression: MPL-2.0

{sgn_drift-0.1.0.dist-info → sgn_drift-0.1.1.dist-info}/RECORD

@@ -1,12 +1,14 @@
 sgndrift/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sgndrift/_version.py,sha256=5jwwVncvCiTnhOedfkzzxmxsggwmTBORdFL_4wq0ZeY,704
+sgndrift/_version.py,sha256=m8HxkqoKGw_wAJtc4ZokpJKNLXqp4zwnNhbnfDtro7w,704
 sgndrift/bin/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sgndrift/bin/estimate_drift.py,sha256=qVp-PacdNRUAEWOrHlvzP-whMRCB_zhteZZDE1Ld_9g,8822
+sgndrift/bin/estimate_psd.py,sha256=Nb-12CzE1s2uOvFzsB3O8z-2wm72BuEOhvOVfq74n1U,6728
 sgndrift/bin/plot_drift.py,sha256=EYIhluD0asFWoYmLfpu2vsmU5TDHtjmPjCNuY2tvYLI,4746
 sgndrift/bin/plot_drift_comparison.py,sha256=0VknMCS1-szBk5trTARUYvYWEXrid-FnyOf0yM1_lnw,6355
 sgndrift/bin/plot_drift_super.py,sha256=W-OgqfL8_UFDKxyNhNe3ahXteQllu3oHPyidYcJd6Kc,7934
 sgndrift/bin/plot_drift_super_comp.py,sha256=lcHqw4w0xKfTW1DX5pAVOGM3mAlxfc9Xsz12fjlQlTQ,12222
 sgndrift/bin/plot_drift_time.py,sha256=EdLkzUe4Exsi6OJDQYLmY8PH1wSHLoRv5q3ZwKyDS8g,5929
+sgndrift/bin/plot_psd.py,sha256=SRkqZrJqZhbn8rRVPu5vShNMGm2VVQi1Y5-MzBGEhPo,8635
 sgndrift/psd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sgndrift/psd/drift.py,sha256=3v0UlLc3yQLSJwLHJXQo5Hy1wRM7wV8hkMiiSGtpY8o,2413
 sgndrift/psd/estimators.py,sha256=UdI7SaUOEiQeXbBzWNfY3C0w0tS8NWMzJHzULJW30Ck,4280
@@ -15,8 +17,8 @@ sgndrift/sinks/drift_sink.py,sha256=PHspvH9LupSmjLWjavvsk2a3TfU5jYXUMN0TV9lTdjs,
 sgndrift/transforms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sgndrift/transforms/drift.py,sha256=bbjmrW1q8awkb-oBpeWo94M-pGxHYN88iOKrkbQHL2A,4727
 sgndrift/transforms/psd.py,sha256=6LABGSUpbOE4nKzHEWybHT3ronnuUgLuV5tI46txn34,5427
-sgn_drift-0.1.0.dist-info/METADATA,sha256=-5mp3zaZIktyCf2Ap6ljRQXEGXwd0QMIWaUFqQhtoRw,2828
-sgn_drift-0.1.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-sgn_drift-0.1.0.dist-info/entry_points.txt,sha256=emRR36dQDFCkTxwuWEMnN-TeNNTrnlJ-HFeJEw_37vY,362
-sgn_drift-0.1.0.dist-info/top_level.txt,sha256=xmogf6bL8sl_VbrQ6pPMXLdY967VtYYUtFzIg_U3Rq4,9
-sgn_drift-0.1.0.dist-info/RECORD,,
+sgn_drift-0.1.1.dist-info/METADATA,sha256=AaZAytDaFP4AMNxOwK4CA6qc393f6itdjUB0j9YQgz8,2828
+sgn_drift-0.1.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+sgn_drift-0.1.1.dist-info/entry_points.txt,sha256=3duQ4GjWXADFIL0L_sXEjzhZS2-Z9wasOuMyRtnVO3c,457
+sgn_drift-0.1.1.dist-info/top_level.txt,sha256=xmogf6bL8sl_VbrQ6pPMXLdY967VtYYUtFzIg_U3Rq4,9
+sgn_drift-0.1.1.dist-info/RECORD,,

{sgn_drift-0.1.0.dist-info → sgn_drift-0.1.1.dist-info}/entry_points.txt

@@ -2,6 +2,8 @@
 sgn-drift-estimate = sgndrift.bin.estimate_drift:main
 sgn-drift-plot = sgndrift.bin.plot_drift:main
 sgn-drift-plot-comp = sgndrift.bin.plot_drift_comparison:main
+sgn-drift-plot-psd = sgndrift.bin.plot_psd:main
 sgn-drift-plot-super = sgndrift.bin.plot_drift_super:main
 sgn-drift-plot-super-comp = sgndrift.bin.plot_drift_super_comp:main
 sgn-drift-plot-time = sgndrift.bin.plot_drift_time:main
+sgn-drift-psd = sgndrift.bin.estimate_psd:main

sgndrift/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.0'
-__version_tuple__ = version_tuple = (0, 1, 0)
+__version__ = version = '0.1.1'
+__version_tuple__ = version_tuple = (0, 1, 1)
 
 __commit_id__ = commit_id = None

sgndrift/bin/estimate_psd.py

@@ -0,0 +1,240 @@
+#!/usr/bin/env python3
+"""
+Robust PSD Estimation Tool.
+
+This script estimates the Power Spectral Density (PSD) of gravitational wave data.
+It generalizes legacy functionality by leveraging the SGN framework for:
+1. Robust data fetching (GWOSC).
+2. Modern PSD estimators (MGM, Recursive).
+3. Flexible output formatting via PSDSink.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from typing import List, Optional
+
+# SGN & Family
+from sgn.apps import Pipeline
+from sgn.base import SourceElement
+from sgndrift.transforms.psd import MGMPSD, RecursivePSD
+from sgnligo.sinks.psd_sink import PSDSink
+from sgnligo.sources.gwosc import GWOSCSource
+
+
+def build_pipeline(
+    start: float,
+    end: float,
+    detectors: List[str],
+    sample_rate: int,
+    fft_length: float,
+    overlap: float,
+    estimator_type: str,
+    output_fmt: str,
+    write_interval: Optional[float] = None,
+    source_element: Optional[SourceElement] = None,
+    verbose: bool = False,
+) -> Pipeline:
+    """Constructs the processing pipeline for the requested interval."""
+    pipe = Pipeline()
+
+    # 1. Source
+    if source_element:
+        source = source_element
+        source_name = source.name
+    else:
+        source_name = f"gwosc_{int(start)}"
+        source = GWOSCSource(
+            name=source_name,
+            start=start,
+            end=end,
+            detectors=detectors,
+            sample_rate=sample_rate,
+            cache_data=True,
+            verbose=verbose,
+        )
+    pipe.insert(source)
+
+    # 2. Per-Detector Processing
+    for ifo in detectors:
+        # A. Estimator
+        est_name = f"psd_{ifo}"
+        if estimator_type.lower() == "recursive":
+            est = RecursivePSD(
+                name=est_name,
+                sink_pad_names=("in",),
+                source_pad_names=("out",),
+                fft_length=fft_length,
+                overlap=overlap,
+                sample_rate=int(sample_rate),
+            )
+        else:
+            # Default to MGM
+            est = MGMPSD(
+                name=est_name,
+                sink_pad_names=("in",),
+                source_pad_names=("out",),
+                fft_length=fft_length,
+                overlap=overlap,
+                sample_rate=int(sample_rate),
+            )
+
+        # B. Sink
+        # Format filename: defaults to "psd_{detector}_{start}.xml.gz"
+        fname = output_fmt.format(
+            detector=ifo,
+            start="{gps}",
+            duration=int(end - start),
+            rate=int(sample_rate),
+        )
+
+        # PSDSink usually expects the IFO name as the sink pad
+        # to correctly label the XML element
+        sink = PSDSink(
+            name=f"sink_{ifo}",
+            sink_pad_names=(ifo,),
+            fname=fname,
+            write_interval=write_interval,
+            verbose=verbose,
+        )
+
+        pipe.insert(est, sink)
+
+        # C. Link
+        # Source[IFO] -> Est[in]
+        # Est[out] -> Sink[IFO]
+
+        # Check source pads
+        if ifo in source.srcs:
+            src_pad = source.srcs[ifo]
+        else:
+            if verbose:
+                print(
+                    f"Warning: Source pad '{ifo}' not found in source '{source.name}'. Skipping {ifo}."
+                )
+            continue
+
+        pipe.link({est.snks["in"]: src_pad, sink.snks[ifo]: est.srcs["out"]})
+
+    return pipe
+
+
+def estimate_psd(
+    start: float,
+    end: float,
+    detectors: List[str],
+    output_fmt: str = "psd_{detector}_{start}.xml.gz",
+    estimator: str = "mgm",
+    fft_length: float = 4.0,
+    overlap: float = 0.5,
+    sample_rate: float = 4096.0,
+    write_interval: Optional[float] = None,
+    verbose: bool = False,
+):
+    """
+    Main driver for PSD estimation.
+    Runs the pipeline over the user-specified [start, end) interval.
+    """
+    duration = end - start
+    if duration < (fft_length * 4):
+        print(f"Error: Duration ({duration}s) is too short for stable PSD estimation.")
+        return
+
+    if verbose:
+        print(f"Processing interval: {start} - {end} ({duration}s)")
+        print(f"Detectors: {detectors}")
+        print(f"Estimator: {estimator}")
+        if write_interval:
+            print(f"Write Interval: {write_interval}s")
+
+    try:
+        pipe = build_pipeline(
+            start=start,
+            end=end,
+            detectors=detectors,
+            sample_rate=int(sample_rate),
+            fft_length=fft_length,
+            overlap=overlap,
+            estimator_type=estimator,
+            output_fmt=output_fmt,
+            write_interval=write_interval,
+            verbose=verbose,
+        )
+        pipe.run()
+        if verbose:
+            print("Pipeline finished successfully.")
+
+    except Exception as e:
+        print(f"Pipeline failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Estimate GW PSDs.")
+    parser.add_argument(
+        "--detectors", type=str, required=True, help="Comma-separated list (e.g. H1,L1)"
+    )
+    parser.add_argument("--start", type=float, required=True, help="GPS Start Time")
+    parser.add_argument("--end", type=float, required=True, help="GPS End Time")
+
+    parser.add_argument(
+        "--output-fmt",
+        type=str,
+        default="psd_{detector}_{start}.xml.gz",
+        help="Output filename format. Available keys: {detector}, {start}, {duration}, {rate}",
+    )
+
+    parser.add_argument(
+        "--estimator",
+        type=str,
+        default="mgm",
+        choices=["mgm", "recursive"],
+        help="PSD Estimation algorithm",
+    )
+    parser.add_argument(
+        "--fft-length", type=float, default=4.0, help="FFT length in seconds"
+    )
+    parser.add_argument(
+        "--overlap", type=float, default=0.5, help="Overlap fraction (0-1)"
+    )
+    parser.add_argument(
+        "--sample-rate", type=float, default=4096.0, help="Sampling rate in Hz"
+    )
+    parser.add_argument(
+        "--write-interval",
+        type=float,
+        default=None,
+        help="Interval in seconds to write intermediate PSDs to disk.",
+    )
+
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable verbose logging"
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    dets = [d.strip() for d in args.detectors.split(",")]
+
+    estimate_psd(
+        start=args.start,
+        end=args.end,
+        detectors=dets,
+        output_fmt=args.output_fmt,
+        estimator=args.estimator,
+        fft_length=args.fft_length,
+        overlap=args.overlap,
+        sample_rate=args.sample_rate,
+        write_interval=args.write_interval,
+        verbose=args.verbose,
+    )
+
+
+if __name__ == "__main__":
+    main()

sgndrift/bin/plot_psd.py

@@ -0,0 +1,280 @@
+#!/usr/bin/env python3
+"""
+PSD Visualization Tool.
+
+Generates comparison plots of Power Spectral Densities (PSDs).
+
+Features:
+- Supports multiple input files (LIGO LW XML, HDF5, TXT).
+- Auto-detects Time Evolution (Drift) using internal metadata.
+- Smart Labeling (Relative Time).
+- Sequential Colormaps for Drift Plots.
+"""
+
+import argparse
+import sys
+import os
+import re
+from typing import List, Optional, Any
+
+import matplotlib.pyplot as plt
+import matplotlib.cm as cm
+import numpy as np
+from gwpy.frequencyseries import FrequencySeries
+
+
+def load_psd(filepath: str, ifo: Optional[str] = None) -> FrequencySeries:
+    """Loads a PSD FrequencySeries from a file using GWpy."""
+    if not os.path.exists(filepath):
+        print(f"Error: File not found at {filepath}")
+        sys.exit(1)
+
+    try:
+        # For XML files, explicit IFO helps GWpy disambiguate
+        if ifo:
+            return FrequencySeries.read(filepath, instrument=ifo)
+        else:
+            return FrequencySeries.read(filepath)
+    except Exception as e:
+        print(f"Error reading {filepath}: {e}")
+        print(f"Tip: If {filepath} is an XML file, try providing --ifos")
+        sys.exit(1)
+
+
+def make_plot(
+    file_list: List[str],
+    output_path: str,
+    ifos: Optional[List[str]] = None,
+    labels: Optional[List[str]] = None,
+    title: Optional[str] = None,
+    xlim: Optional[List[float]] = None,
+    ylim: Optional[List[float]] = None,
+    asd_mode: bool = False,
+    sort_time: bool = True,
+):
+    """Generates and saves the PSD comparison plot."""
+
+    # --- 1. Load Data & Extract Metadata ---
+    loaded_items = []
+
+    for i, fpath in enumerate(file_list):
+        # IFO hint from CLI args
+        ifo_hint = ifos[i] if ifos and i < len(ifos) else None
+
+        # Load Series
+        series = load_psd(fpath, ifo_hint)
+
+        # Extract GPS Time from Object Attributes
+        gps_time = None
+        if series.epoch is not None:
+            # gwpy epoch is usually an astropy Time object
+            try:
+                gps_time = series.epoch.gps
+            except AttributeError:
+                # Fallback if it's just a number
+                gps_time = float(series.epoch)
+
+        # Extract IFO Name (Prefer Channel Name -> Filename -> Hint)
+        detected_ifo = None
+        if series.channel:
+            # Typical format: "H1:DMT-PSD..."
+            cname = str(series.channel)
+            if ":" in cname:
+                detected_ifo = cname.split(":")[0]
+
+        if not detected_ifo:
+            # Fallback regex on filename
+            match = re.search(r"([HLV]1)", os.path.basename(fpath))
+            if match:
+                detected_ifo = match.group(1)
+
+        # Final IFO resolution
+        final_ifo = ifo_hint if ifo_hint else detected_ifo
+
+        loaded_items.append(
+            {
+                "series": series,
+                "gps": gps_time,
+                "ifo": final_ifo,
+                "path": fpath,
+                "explicit_label": labels[i] if labels and i < len(labels) else None,
+            }
+        )
+
+    # --- 2. Sort by Time ---
+    if sort_time:
+        # Sort by GPS, handling None (put at start or end as preferred)
+        loaded_items.sort(key=lambda x: x["gps"] if x["gps"] is not None else -1.0)
+
+    # --- 3. Determine Plot Mode (Drift vs Comparison) ---
+    unique_ifos = set(item["ifo"] for item in loaded_items if item["ifo"])
+    # If strictly one IFO is found across multiple files, assume Drift Mode
+    is_drift_mode = (len(unique_ifos) == 1) and (len(loaded_items) > 1)
+
+    # --- 4. Setup Plotting ---
+    plt.figure(figsize=(12, 8))
+
+    # Colors
+    num_plots = len(loaded_items)
+    if is_drift_mode:
+        # Sequential (Viridis) for time evolution
+        colors = cm.viridis(np.linspace(0, 1, num_plots))
+    else:
+        # Qualitative (Tab10) for distinct items
+        colors = cm.tab10(np.linspace(0, 1, min(num_plots, 10)))
+
+    # Base Time (Start of sequence)
+    t0 = loaded_items[0]["gps"] if loaded_items and loaded_items[0]["gps"] else None
+
+    for i, item in enumerate(loaded_items):
+        series = item["series"]
+
+        # ASD Transform
+        if asd_mode:
+            # GWpy propagates units automatically
+            series = series**0.5
+            series.name = "Amplitude Spectral Density"
+
+        # Construct Label
+        if item["explicit_label"]:
+            label = item["explicit_label"]
+        elif is_drift_mode and item["gps"] is not None and t0 is not None:
+            # Smart Relative Label: "H1 (+16s)"
+            dt = item["gps"] - t0
+            ifo_str = item["ifo"] if item["ifo"] else "PSD"
+            label = f"{ifo_str} (T+{dt:.1f}s)"
+        else:
+            # Standard Label: "H1 (Channel)" or Filename
+            base = item["ifo"] if item["ifo"] else os.path.basename(item["path"])
+            # Only append GPS to label if we aren't in smart drift mode
+            if item["gps"]:
+                base += f" @ {int(item['gps'])}"
+            label = base
+
+        # Plot
+        plt.loglog(
+            series.frequencies.value,
+            series.value,
+            label=label,
+            alpha=0.8,
+            linewidth=1.5,
+            color=colors[i % len(colors)],
+        )
+
+    # --- 5. Styling ---
+    plt.xlabel("Frequency [Hz]", fontsize=12)
+
+    if asd_mode:
+        plt.ylabel(r"ASD [$\mathrm{Strain} / \sqrt{\mathrm{Hz}}$]", fontsize=12)
+    else:
+        plt.ylabel(r"PSD [$\mathrm{Strain}^2 / \mathrm{Hz}$]", fontsize=12)
+
+    plt.grid(True, which="both", alpha=0.2, ls="-")
+
+    if xlim:
+        plt.xlim(xlim)
+    else:
+        plt.xlim(10, 2048)
+
+    if ylim:
+        plt.ylim(ylim)
+
+    if title:
+        plt.title(title, fontsize=14)
+    else:
+        mode = "Drift Evolution" if is_drift_mode else "Noise Comparison"
+        if t0 and is_drift_mode:
+            mode += f" (Start GPS: {int(t0)})"
+        plt.title(mode, fontsize=14)
+
+    # Legend
+    legend = plt.legend(loc="upper right", framealpha=0.9, fontsize=10)
+    if is_drift_mode:
+        legend.set_title("Time Evolution")
+
+    try:
+        plt.savefig(output_path, dpi=150, bbox_inches="tight")
+        print(f"Plot saved to: {output_path}")
+    except Exception as e:
+        print(f"Error saving plot: {e}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Visualize and Compare GW PSDs.")
+    parser.add_argument(
+        "input_files", nargs="+", type=str, help="Path(s) to input PSD file(s)"
+    )
+    parser.add_argument(
+        "-o", "--output", type=str, default="psd_plot.png", help="Output filename"
+    )
+
+    # Plotting Options
+    parser.add_argument(
+        "--labels",
+        type=str,
+        default=None,
+        help="Comma-separated labels (overrides auto-detection)",
+    )
+    parser.add_argument(
+        "--ifos",
+        type=str,
+        default=None,
+        help="Comma-separated IFO names (hint for XML parsing)",
+    )
+    parser.add_argument(
+        "--asd", action="store_true", help="Plot Amplitude Spectral Density (sqrt(PSD))"
+    )
+    parser.add_argument(
+        "--no-sort", action="store_true", help="Disable automatic sorting by GPS time"
+    )
+
+    # Ranges
+    parser.add_argument("--fmin", type=float, default=10.0, help="Min Frequency")
+    parser.add_argument("--fmax", type=float, default=2048.0, help="Max Frequency")
+    parser.add_argument("--ymin", type=float, default=None, help="Min Y-axis")
+    parser.add_argument("--ymax", type=float, default=None, help="Max Y-axis")
+
+    parser.add_argument("--title", type=str, default=None, help="Custom plot title")
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # Process comma-separated lists
+    labels = [l.strip() for l in args.labels.split(",")] if args.labels else None
+    ifos = [i.strip() for i in args.ifos.split(",")] if args.ifos else None
+
+    # Handle shell globbing (Python glob if shell didn't expand)
+    expanded_files = []
+    import glob
+
+    for f in args.input_files:
+        if "*" in f:
+            expanded_files.extend(glob.glob(f))
+        else:
+            expanded_files.append(f)
+
+    # Remove duplicates and sort for consistency
+    expanded_files = sorted(list(set(expanded_files)))
+
+    if not expanded_files:
+        print("No input files found.")
+        sys.exit(1)
+
+    make_plot(
+        expanded_files,
+        output_path=args.output,
+        ifos=ifos,
+        labels=labels,
+        title=args.title,
+        xlim=[args.fmin, args.fmax],
+        ylim=[args.ymin, args.ymax] if args.ymin and args.ymax else None,
+        asd_mode=args.asd,
+        sort_time=not args.no_sort,
+    )
+
+
+if __name__ == "__main__":
+    main()