deepextractor.model

High-level DeepExtractor model wrapper for inference.

Module Contents

deepextractor.model.logger

class deepextractor.model.DeepExtractorModel(checkpoint: str = 'DeepExtractor_257', checkpoint_filename: str = CHECKPOINT_BILBY, checkpoint_dir: str | None = None, device: str | torch.device | None = None, scaler_path: str | None = None, n_fft: int = 512, win_length: int = 64, hop_length: int = 32)

High-level wrapper around a pretrained DeepExtractor UNET2D model.

Bundles the PyTorch model, StandardScaler, and STFT parameters into a single object so callers don’t need to manage them separately.

Parameters:

• checkpoint (str) – Model name / checkpoint key. Defaults to "DeepExtractor_257".

• checkpoint_filename (str) – Checkpoint file name within the model subdirectory on HuggingFace Hub or local checkpoint_dir. Defaults to CHECKPOINT_BILBY.

• checkpoint_dir (str | None) – Local directory to search for checkpoint files before falling back to HuggingFace Hub. Pass None to always use the Hub.

• device (str | torch.device | None) – Compute device. Auto-detects CUDA if available when None.

• scaler_path (str | None) – Path to the scaler .pkl file. Defaults to the bundled assets/scaler_bilby.pkl.

• n_fft (int) – STFT FFT size. Default 512.

• win_length (int) – STFT window length. Default 64.

• hop_length (int) – STFT hop length. Default 32.

device

n_fft = 512

win_length = 64

hop_length = 32

background(noisy_input: numpy.ndarray) → numpy.ndarray

Estimate the background (noise-only) component.

Parameters:

noisy_input (np.ndarray) – 1-D array of shape (T,) or 2-D batch of shape (N, T).

Returns:

Background estimate, same shape as noisy_input.

Return type:

np.ndarray

reconstruct(noisy_input: numpy.ndarray) → numpy.ndarray

Extract the transient signal by subtracting the predicted background.

Parameters:

noisy_input (np.ndarray) – 1-D array of shape (T,) or 2-D batch of shape (N, T).

Returns:

Reconstructed signal, same shape as noisy_input.

Return type:

np.ndarray

class deepextractor.model.SeparationResult

Bases: NamedTuple

Outputs of DeepExtractorSeparator.separate().

All arrays have shape (T,) for single inputs or (N, T) for batches.

h1_signal: numpy.ndarray

l1_signal: numpy.ndarray

h1_background: numpy.ndarray

l1_background: numpy.ndarray

class deepextractor.model.DeepExtractorSeparator(checkpoint_path: str | pathlib.Path, scaler=None, device: str | torch.device | None = None, model_kwargs: dict | None = None)

Two-detector time-domain signal/glitch separator.

Wraps a pretrained UNET1D_LSTM_ATT model and a ChannelStandardScaler to expose a clean inference API for separating H1+L1 strain into signal and background components in the time domain.

Parameters:

• checkpoint_path (str | Path) – Path to the .pth.tar checkpoint saved during training.

• scaler (ChannelStandardScaler | str | Path | None) – Per-channel input scaler. Pass a fitted ChannelStandardScaler instance, a path to a pickled scaler, or None to skip scaling (not recommended — the model expects standard-scaled inputs).

• device (str | torch.device | None) – Compute device. Auto-detects CUDA if available when None.

• model_kwargs (dict | None) – Override keyword arguments forwarded to UNET1D_LSTM_ATT. By default uses in_channels=2, out_channels=4.

device

separate(h1: numpy.ndarray, l1: numpy.ndarray) → SeparationResult

Separate H1 and L1 strain into signal and background components.

Parameters:

• h1 (np.ndarray) – H1 strain. Shape (T,) or (N, T).

• l1 (np.ndarray) – L1 strain. Same shape as h1.

Returns:

Named tuple with fields h1_signal, l1_signal, h1_background, l1_background, each of the same shape as the inputs.

Return type:

SeparationResult