deepextractor 0.1.0__py3-none-any.whl

deepextractor/__init__.py

@@ -0,0 +1,28 @@
+"""
+DeepExtractor: Deep learning for gravitational-wave glitch reconstruction.
+
+Quick usage::
+
+    import deepextractor
+    reconstructed = deepextractor.reconstruct(noisy_strain)
+
+    # Or with explicit model control:
+    model = deepextractor.DeepExtractorModel(checkpoint="DeepExtractor_257")
+    signal = model.reconstruct(noisy_strain)
+    background = model.background(noisy_strain)
+
+Paper: https://arxiv.org/abs/2501.18423
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("deepextractor")
+except PackageNotFoundError:
+    # Running from source without pip install -e .
+    __version__ = "0.0.0.dev"
+
+from deepextractor.model import DeepExtractorModel
+from deepextractor.api import extract, reconstruct
+
+__all__ = ["__version__", "DeepExtractorModel", "reconstruct", "extract"]

deepextractor/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

deepextractor/api.py

@@ -0,0 +1,60 @@
+"""
+Top-level convenience functions for DeepExtractor inference.
+
+For one-shot use. For repeated inference on many signals, instantiate
+:class:`DeepExtractorModel` directly to amortise the model load cost.
+"""
+
+import numpy as np
+
+from deepextractor.model import DeepExtractorModel
+from deepextractor.utils.checkpoints import CHECKPOINT_BILBY
+
+
+def reconstruct(
+    noisy_input: np.ndarray,
+    checkpoint: str = "DeepExtractor_257",
+    checkpoint_filename: str = CHECKPOINT_BILBY,
+    checkpoint_dir: str | None = None,
+    device: str | None = None,
+    scaler_path: str | None = None,
+) -> np.ndarray:
+    """
+    Extract the transient signal from a noisy gravitational-wave strain.
+
+    Loads a DeepExtractor model, runs inference, and returns the reconstructed
+    signal. For repeated calls, prefer instantiating :class:`DeepExtractorModel`
+    directly to avoid reloading weights on each call.
+
+    Parameters
+    ----------
+    noisy_input : np.ndarray
+        1-D array of shape ``(T,)`` or 2-D batch of shape ``(N, T)``.
+    checkpoint : str
+        Model name. Default ``"DeepExtractor_257"``.
+    checkpoint_filename : str
+        Checkpoint filename. Defaults to the bilby-noise checkpoint.
+    checkpoint_dir : str | None
+        Local checkpoint directory. Falls back to HuggingFace Hub if None.
+    device : str | None
+        Torch device string. Auto-detected if None.
+    scaler_path : str | None
+        Path to scaler .pkl. Uses bundled asset if None.
+
+    Returns
+    -------
+    np.ndarray
+        Reconstructed signal, same shape as ``noisy_input``.
+    """
+    model = DeepExtractorModel(
+        checkpoint=checkpoint,
+        checkpoint_filename=checkpoint_filename,
+        checkpoint_dir=checkpoint_dir,
+        device=device,
+        scaler_path=scaler_path,
+    )
+    return model.reconstruct(noisy_input)
+
+
+# `extract` and `reconstruct` are synonyms at the API level.
+extract = reconstruct

deepextractor/generation/__init__.py

@@ -0,0 +1 @@
+"""Synthetic glitch signal generators and data generation scripts."""

deepextractor/generation/generate_spectrograms.py

@@ -0,0 +1,158 @@
+"""
+Convert time-domain .npy arrays to STFT spectrograms (magnitude + phase).
+
+Also provides a utility to concatenate chunked spectrogram files.
+
+Usage::
+
+    deepextractor-specgen --input-dir data/pycbc_noise/time_domain/ --output-dir data/pycbc_noise/spectrogram_domain/
+
+"""
+
+import argparse
+import os
+
+import numpy as np
+import torch
+
+
+# Default STFT parameters (257x257 output shape)
+DEFAULT_N_FFT = 256 * 2
+DEFAULT_WIN_LENGTH = DEFAULT_N_FFT // 8
+DEFAULT_HOP_LENGTH = DEFAULT_WIN_LENGTH // 2
+
+
+def apply_stft_and_save(
+    array_path, save_path, n_fft, hop_length, win_length, window, chunk_size=5000
+):
+    """Apply STFT to a .npy array in chunks and save the result."""
+    array = np.load(array_path)
+    print(f"Loaded {array_path}, shape: {array.shape}")
+
+    total_chunks = array.shape[0] // chunk_size
+    stft_list = []
+
+    for i in range(0, array.shape[0], chunk_size):
+        chunk = array[i : i + chunk_size]
+        tensor = torch.tensor(chunk, dtype=torch.float32)
+        stft_result = torch.stft(
+            tensor,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=window,
+            return_complex=True,
+        )
+        magnitude = torch.abs(stft_result)
+        phase = torch.angle(stft_result)
+        stft_mag_phase = torch.stack([magnitude, phase], dim=1)
+        stft_list.append(stft_mag_phase)
+
+        del tensor, stft_result, magnitude, phase
+        torch.cuda.empty_cache()
+
+        print(f"Processed chunk {i // chunk_size + 1}/{max(total_chunks, 1)}")
+
+    stft_final = torch.cat(stft_list, dim=0)
+    stft_numpy = stft_final.cpu().numpy()
+    np.save(save_path, stft_numpy)
+    print(f"STFT saved to {save_path}.npy, final shape: {stft_numpy.shape}")
+
+    del array, stft_list, stft_final, stft_numpy
+    torch.cuda.empty_cache()
+
+
+def load_and_concatenate_chunks(data_dir, base_filename, total_chunks):
+    """Load and concatenate chunked numpy arrays saved as ``{base}_chunk_{i}.npy``."""
+    stft_list = []
+    for i in range(total_chunks):
+        chunk_filename = f"{base_filename}_chunk_{i}.npy"
+        chunk_path = os.path.join(data_dir, chunk_filename)
+        if os.path.exists(chunk_path):
+            print(f"Loading {chunk_filename}...")
+            stft_list.append(np.load(chunk_path))
+        else:
+            print(f"Chunk {chunk_filename} not found. Skipping.")
+    print("Concatenating chunks...")
+    return np.concatenate(stft_list, axis=0)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert time-domain .npy arrays to STFT spectrogram arrays",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--input-dir",
+        type=str,
+        required=True,
+        help="Directory containing the time-domain .npy files.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        required=True,
+        help="Directory to save the spectrogram .npy files.",
+    )
+    parser.add_argument("--n-fft", type=int, default=DEFAULT_N_FFT)
+    parser.add_argument("--win-length", type=int, default=DEFAULT_WIN_LENGTH)
+    parser.add_argument("--hop-length", type=int, default=DEFAULT_HOP_LENGTH)
+    parser.add_argument("--chunk-size", type=int, default=5000)
+    parser.add_argument(
+        "--combine-chunks",
+        action="store_true",
+        help="Combine pre-existing chunk files instead of generating new spectrograms.",
+    )
+    parser.add_argument(
+        "--chunks-glitch-train", type=int, default=16,
+        help="Number of chunks for glitch_train (used with --combine-chunks).",
+    )
+    parser.add_argument(
+        "--chunks-background-train", type=int, default=16,
+        help="Number of chunks for background_train (used with --combine-chunks).",
+    )
+    parser.add_argument(
+        "--chunks-glitch-val", type=int, default=2,
+        help="Number of chunks for glitch_val (used with --combine-chunks).",
+    )
+    parser.add_argument(
+        "--chunks-background-val", type=int, default=2,
+        help="Number of chunks for background_val (used with --combine-chunks).",
+    )
+    args = parser.parse_args()
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    window = torch.hann_window(args.win_length)
+
+    if args.combine_chunks:
+        for base, n_chunks in [
+            ("glitch_train_scaled_mag_phase", args.chunks_glitch_train),
+            ("background_train_scaled_mag_phase", args.chunks_background_train),
+            ("glitch_val_scaled_mag_phase", args.chunks_glitch_val),
+            ("background_val_scaled_mag_phase", args.chunks_background_val),
+        ]:
+            combined = load_and_concatenate_chunks(args.output_dir, base, n_chunks)
+            out_path = os.path.join(args.output_dir, f"{base}_combined.npy")
+            np.save(out_path, combined)
+            print(f"Saved combined {base} to {out_path}")
+        print("All combined datasets saved.")
+    else:
+        datasets = [
+            ("glitch_train_scaled.npy", "glitch_train_scaled_mag_phase"),
+            ("background_train_scaled.npy", "background_train_scaled_mag_phase"),
+            ("glitch_val_scaled.npy", "glitch_val_scaled_mag_phase"),
+            ("background_val_scaled.npy", "background_val_scaled_mag_phase"),
+        ]
+        for in_name, out_name in datasets:
+            in_path = os.path.join(args.input_dir, in_name)
+            out_path = os.path.join(args.output_dir, out_name)
+            apply_stft_and_save(
+                in_path, out_path,
+                args.n_fft, args.hop_length, args.win_length, window,
+                args.chunk_size,
+            )
+        print("All STFT results saved.")
+
+
+if __name__ == "__main__":
+    main()

deepextractor/generation/generate_timeseries.py

@@ -0,0 +1,162 @@
+"""
+Generate synthetic time-domain training data.
+
+Usage::
+
+    deepextractor-generate --output-dir data/ --num-train 250000 --bilby-noise
+
+"""
+
+import argparse
+import os
+import pickle
+import random
+
+import numpy as np
+from sklearn.preprocessing import StandardScaler
+from tqdm import tqdm
+
+from deepextractor.generation.glitch_functions import (
+    generate_chirp,
+    generate_gaussian_pulse,
+    generate_sine,
+    generate_sine_gaussian,
+    ringdown,
+)
+from deepextractor.utils.signal import whitened_snr_scaling
+
+SAMPLE_RATE = 4096
+T = 2.0
+T_INJ = T / 2
+LENGTH = int(T * SAMPLE_RATE)
+SNR_MIN, SNR_MAX = 1, 250
+MINIMUM_FREQUENCY = 20.0
+SNR_SCALING_FACTOR_BILBY = 31.970149253731343
+
+SIGNAL_TYPES = ["chirp", "sine", "sine_gaussian", "gaussian_pulse", "ringdown"]
+SIGNAL_FUNCTION_MAP = {
+    "chirp": generate_chirp,
+    "sine": generate_sine,
+    "sine_gaussian": generate_sine_gaussian,
+    "gaussian_pulse": generate_gaussian_pulse,
+    "ringdown": ringdown,
+}
+
+
+def generate_gaussian_noise(mean, std_dev, num_samples, sample_shape, bilby_noise=False,
+                            sample_rate=SAMPLE_RATE, duration=T,
+                            minimum_frequency=MINIMUM_FREQUENCY):
+    """Generate Gaussian noise samples (pycbc or bilby)."""
+    if bilby_noise:
+        try:
+            import bilby
+        except ImportError as e:
+            raise ImportError(
+                "bilby is required for bilby noise generation. "
+                "Install it with: pip install deepextractor[generative]"
+            ) from e
+        gaussian_noise_samples = []
+        for i in tqdm(range(num_samples), desc="Generating bilby noise..."):
+            ifos = bilby.gw.detector.InterferometerList(["L1"])
+            for ifo in ifos:
+                ifo.minimum_frequency = minimum_frequency
+            ifos.set_strain_data_from_power_spectral_densities(
+                sampling_frequency=sample_rate,
+                duration=duration,
+                start_time=0,
+            )
+            white_time_domain_strain = list(ifos[0].whitened_time_domain_strain)
+            gaussian_noise_samples.append(white_time_domain_strain)
+        return np.asarray(gaussian_noise_samples)
+    else:
+        print("Generating pycbc noise...")
+        return np.random.normal(loc=mean, scale=std_dev, size=(num_samples, *sample_shape))
+
+
+def generate_synthetic_data(gaussian_noise_samples, bilby_noise=False, phase="train",
+                             t_min=0.125, t_max=2.0, snr_min=SNR_MIN, snr_max=SNR_MAX):
+    """Generate synthetic noisy glitch and background data arrays."""
+    noisy_glitch_ts = []
+    pure_noise_ts = []
+
+    for i in tqdm(range(len(gaussian_noise_samples)),
+                  desc=f"Generating Synthetic {phase.capitalize()} Data"):
+        background = gaussian_noise_samples[i]
+        noisy_glitch = background.copy()
+        n_injs = np.random.randint(1, 30)
+        for _ in range(n_injs):
+            snr_to_scale = np.random.uniform(snr_min, snr_max)
+            if bilby_noise:
+                snr_to_scale = snr_to_scale / SNR_SCALING_FACTOR_BILBY
+            duration = np.random.uniform(t_min, t_max)
+            s_type = random.choice(SIGNAL_TYPES)
+            _, signal_injection = SIGNAL_FUNCTION_MAP[s_type](duration)
+            len_glitch = len(signal_injection)
+            id_start = int((T_INJ * SAMPLE_RATE / LENGTH) * len(background)) - len_glitch // 2
+            glitch = signal_injection - np.mean(signal_injection)
+            glitch = whitened_snr_scaling(glitch, snr=snr_to_scale)
+            shift_int = np.random.randint(-id_start, len(background) - id_start - len_glitch)
+            noisy_glitch[id_start + shift_int:id_start + len_glitch + shift_int] += glitch
+
+        noisy_glitch_ts.append(noisy_glitch)
+        pure_noise_ts.append(background)
+
+    noisy_glitch_ts = np.asarray(noisy_glitch_ts)
+    pure_noise_ts = np.asarray(pure_noise_ts)
+
+    mask = ~np.any(
+        np.isnan(noisy_glitch_ts) | np.isinf(noisy_glitch_ts)
+        | (np.abs(noisy_glitch_ts) > np.finfo(np.float64).max),
+        axis=1,
+    )
+    return noisy_glitch_ts[mask], pure_noise_ts[mask]
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate synthetic time-domain training data",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--output-dir", type=str, default="data/", help="Root output directory.")
+    parser.add_argument("--num-train", type=int, default=250000)
+    parser.add_argument("--num-val", type=int, default=25000)
+    parser.add_argument(
+        "--bilby-noise", action="store_true", help="Use bilby noise instead of pycbc."
+    )
+    args = parser.parse_args()
+
+    bilby_noise = args.bilby_noise
+    noise_ext = "bilby_noise/" if bilby_noise else "pycbc_noise/"
+    ext = "bilby" if bilby_noise else "pycbc"
+    noise_type_path = os.path.join(args.output_dir, noise_ext)
+    domain_path = os.path.join(noise_type_path, "time_domain")
+    os.makedirs(domain_path, exist_ok=True)
+
+    mean = 0
+    std_dev = np.sqrt(SAMPLE_RATE)
+
+    train_noise = generate_gaussian_noise(mean, std_dev, args.num_train, (LENGTH,), bilby_noise)
+    val_noise = generate_gaussian_noise(mean, std_dev, args.num_val, (LENGTH,), bilby_noise)
+
+    glitch_train, bg_train = generate_synthetic_data(train_noise, bilby_noise, "train")
+    glitch_val, bg_val = generate_synthetic_data(val_noise, bilby_noise, "val")
+
+    scaler = StandardScaler()
+    glitch_train_scaled = scaler.fit_transform(glitch_train.reshape(-1, 1)).reshape(glitch_train.shape)
+    bg_train_scaled = scaler.transform(bg_train.reshape(-1, 1)).reshape(bg_train.shape)
+    glitch_val_scaled = scaler.transform(glitch_val.reshape(-1, 1)).reshape(glitch_val.shape)
+    bg_val_scaled = scaler.transform(bg_val.reshape(-1, 1)).reshape(bg_val.shape)
+
+    with open(os.path.join(noise_type_path, f"scaler_{ext}.pkl"), "wb") as f:
+        pickle.dump(scaler, f)
+
+    np.save(os.path.join(domain_path, "glitch_train_scaled"), glitch_train_scaled)
+    np.save(os.path.join(domain_path, "background_train_scaled"), bg_train_scaled)
+    np.save(os.path.join(domain_path, "glitch_val_scaled"), glitch_val_scaled)
+    np.save(os.path.join(domain_path, "background_val_scaled"), bg_val_scaled)
+
+    print("Done. Data saved to", domain_path)
+
+
+if __name__ == "__main__":
+    main()

deepextractor/generation/glitch_functions.py

@@ -0,0 +1,145 @@
+"""
+Synthetic glitch signal generators.
+
+The CDVGAN and gengli generators require optional dependencies:
+    pip install deepextractor[generative]
+"""
+
+import numpy as np
+from scipy.signal import chirp, gausspulse
+
+from deepextractor.utils.signal import quality_factor_conversion, rescale
+
+SRATE = 4096
+NYQUIST_FREQ = SRATE // 2
+
+
+def generate_chirp(duration, sample_rate=4096, f0_min=1, f0_max=NYQUIST_FREQ,
+                   f1_min=1, f1_max=NYQUIST_FREQ):
+    f0 = np.random.uniform(f0_min, f0_max)
+    f1 = np.random.uniform(f1_min, f1_max)
+    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+    signal = chirp(t, f0=f0, f1=f1, t1=duration, method="linear")
+    return t, signal
+
+
+def generate_sine(duration, sample_rate=4096, freq_min=1, freq_max=NYQUIST_FREQ):
+    frequency = np.random.uniform(freq_min, freq_max)
+    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+    signal = np.sin(2 * np.pi * frequency * t)
+    return t, signal
+
+
+def generate_sine_gaussian(duration, sample_rate=4096, freq_min=1, freq_max=NYQUIST_FREQ):
+    tau = np.random.uniform(duration / 200, duration / 4)
+    frequency = np.random.uniform(freq_min, freq_max)
+    t = np.linspace(-duration / 2, duration / 2, int(sample_rate * duration), endpoint=False)
+    signal = np.sin(2 * np.pi * frequency * t) * np.exp(-(t**2) / (2 * tau**2))
+    return t, signal
+
+
+def generate_gaussian_pulse(duration, sample_rate=4096, fc_min=1, fc_max=NYQUIST_FREQ,
+                             bw_min=0.1, bw_max=1.0, bwr_min=-10, bwr_max=0,
+                             tpr_min=0.5, tpr_max=2.0):
+    """
+    Generate a Gaussian pulse with random parameters.
+
+    Parameters
+    ----------
+    duration : float
+        Duration in seconds.
+    sample_rate : int
+        Sampling rate in Hz.
+    fc_min, fc_max : float
+        Range for the center frequency (Hz).
+    bw_min, bw_max : float
+        Range for the fractional bandwidth.
+    bwr_min, bwr_max : float
+        Range for the bandwidth reference level (dB).
+    tpr_min, tpr_max : float
+        Range for the taper reference level (dB).
+    """
+    bw = np.random.uniform(bw_min, bw_max)
+    fc = np.random.uniform(fc_min, fc_max)
+    bwr = np.random.uniform(bwr_min, bwr_max)
+    tpr = np.random.uniform(tpr_min, tpr_max)
+    t = np.linspace(-duration / 2, duration / 2, int(sample_rate * duration), endpoint=False)
+    signal = gausspulse(t, fc=fc, bw=bw, bwr=bwr, tpr=tpr)
+    return t, signal
+
+
+def ringdown(duration, sample_rate=4096, n_signals=1):
+    t = np.linspace(0, duration, int(sample_rate * duration))
+    phi = np.random.uniform(0, 2 * np.pi)
+    A = 1.0
+    f_0 = np.random.uniform(10, NYQUIST_FREQ, n_signals)
+    t_0 = np.random.uniform(t[-1] / 4, 3 * t[-1] / 4, n_signals)
+    Q = np.random.uniform(5, 150, n_signals)
+    tau = np.maximum(quality_factor_conversion(Q, f_0), 0.01)
+    f_0 = np.expand_dims(f_0, axis=1)
+    t_0 = np.expand_dims(t_0, axis=1)
+    tau = np.expand_dims(tau, axis=1)
+    h_1 = A * np.exp(-1.0 * ((t - t_0) / (tau))) * np.sin(2 * np.pi * f_0 * (t - t_0) + phi)
+    h_1 = ((t - t_0) > 0) * h_1
+    h_1 = rescale(h_1)
+    if np.random.rand() < 0.5:
+        h_1 = np.flip(h_1)
+    return t, h_1[0]
+
+
+def generate_gengli_glitch(ifo):
+    """
+    Generate a glitch sample using the gengli library.
+
+    Requires the ``[generative]`` optional dependencies:
+    ``pip install deepextractor[generative]``
+    """
+    try:
+        import gengli
+    except ImportError as e:
+        raise ImportError(
+            "gengli is required for this function. "
+            "Install it with: pip install deepextractor[generative]"
+        ) from e
+    g = gengli.glitch_generator(ifo)
+    glitch = g.get_glitch(1, srate=4096, snr=10, alpha=0.2, fhigh=1024)
+    return None, glitch
+
+
+def generate_cdvgan_glitch(gtype, cdvgan_generator):
+    """
+    Generate a glitch sample using a pretrained CDVGAN TensorFlow model.
+
+    Parameters
+    ----------
+    gtype : str
+        Glitch type: one of ``'blip'``, ``'tomte'``, ``'bbh'``,
+        ``'simplex'``, ``'uniform'``.
+    cdvgan_generator : tf.keras.Model
+        The loaded CDVGAN generator model.
+    """
+    try:
+        import tensorflow as tf
+    except ImportError as e:
+        raise ImportError(
+            "TensorFlow is required for CDVGAN glitches. "
+            "Install it with: pip install deepextractor[generative]"
+        ) from e
+
+    latent_dim = 100
+    random_ints = np.random.randint(0, 100, size=(1, 3))
+    simplex_classes = random_ints / np.sum(random_ints, axis=1).reshape(1, 1)
+    uniform_classes = np.random.uniform(low=0.0, high=1.0, size=(1, 3))
+
+    class_vector_map = {
+        "blip": [[1, 0, 0]],
+        "tomte": [[0, 1, 0]],
+        "bbh": [[0, 0, 1]],
+        "simplex": simplex_classes,
+        "uniform": uniform_classes,
+    }
+
+    class_vector = np.array(class_vector_map[gtype])
+    latent_vector = tf.random.normal(shape=(1, latent_dim))
+    generated_glitch = cdvgan_generator([latent_vector, class_vector]).numpy()
+    return None, generated_glitch