deepextractor.model
===================

.. py:module:: deepextractor.model

.. autoapi-nested-parse::

   High-level DeepExtractor model wrapper for inference.







Module Contents
---------------

.. py:data:: logger

.. py:class:: 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.

   :param checkpoint: Model name / checkpoint key. Defaults to ``"DeepExtractor_257"``.
   :type checkpoint: str
   :param checkpoint_filename: Checkpoint file name within the model subdirectory on HuggingFace Hub
                               or local ``checkpoint_dir``. Defaults to ``CHECKPOINT_BILBY``.
   :type checkpoint_filename: str
   :param checkpoint_dir: Local directory to search for checkpoint files before falling back to
                          HuggingFace Hub. Pass ``None`` to always use the Hub.
   :type checkpoint_dir: str | None
   :param device: Compute device. Auto-detects CUDA if available when ``None``.
   :type device: str | torch.device | None
   :param scaler_path: Path to the scaler ``.pkl`` file. Defaults to the bundled
                       ``assets/scaler_bilby.pkl``.
   :type scaler_path: str | None
   :param n_fft: STFT FFT size. Default 512.
   :type n_fft: int
   :param win_length: STFT window length. Default 64.
   :type win_length: int
   :param hop_length: STFT hop length. Default 32.
   :type hop_length: int


   .. py:attribute:: device


   .. py:attribute:: n_fft
      :value: 512



   .. py:attribute:: win_length
      :value: 64



   .. py:attribute:: hop_length
      :value: 32



   .. py:method:: background(noisy_input: numpy.ndarray) -> numpy.ndarray

      Estimate the background (noise-only) component.

      :param noisy_input: 1-D array of shape ``(T,)`` or 2-D batch of shape ``(N, T)``.
      :type noisy_input: np.ndarray

      :returns: Background estimate, same shape as ``noisy_input``.
      :rtype: np.ndarray



   .. py:method:: reconstruct(noisy_input: numpy.ndarray) -> numpy.ndarray

      Extract the transient signal by subtracting the predicted background.

      :param noisy_input: 1-D array of shape ``(T,)`` or 2-D batch of shape ``(N, T)``.
      :type noisy_input: np.ndarray

      :returns: Reconstructed signal, same shape as ``noisy_input``.
      :rtype: np.ndarray



.. py:class:: SeparationResult

   Bases: :py:obj:`NamedTuple`


   Outputs of :meth:`DeepExtractorSeparator.separate`.

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


   .. py:attribute:: h1_signal
      :type:  numpy.ndarray


   .. py:attribute:: l1_signal
      :type:  numpy.ndarray


   .. py:attribute:: h1_background
      :type:  numpy.ndarray


   .. py:attribute:: l1_background
      :type:  numpy.ndarray


.. py:class:: DeepExtractorSeparator(checkpoint_path: Union[str, pathlib.Path], scaler=None, device: Union[str, torch.device, None] = None, model_kwargs: dict | None = None)

   Two-detector time-domain signal/glitch separator.

   Wraps a pretrained :class:`~deepextractor.models.UNET1D_LSTM_ATT` model and
   a :class:`~deepextractor.data.ChannelStandardScaler` to expose a clean
   inference API for separating H1+L1 strain into signal and background
   components in the time domain.

   :param checkpoint_path: Path to the ``.pth.tar`` checkpoint saved during training.
   :type checkpoint_path: str | Path
   :param scaler: Per-channel input scaler. Pass a fitted
                  :class:`~deepextractor.data.ChannelStandardScaler` instance, a path to
                  a pickled scaler, or ``None`` to skip scaling (not recommended —
                  the model expects standard-scaled inputs).
   :type scaler: ChannelStandardScaler | str | Path | None
   :param device: Compute device. Auto-detects CUDA if available when ``None``.
   :type device: str | torch.device | None
   :param model_kwargs: Override keyword arguments forwarded to :class:`UNET1D_LSTM_ATT`.
                        By default uses ``in_channels=2, out_channels=4``.
   :type model_kwargs: dict | None


   .. py:attribute:: device


   .. py:method:: separate(h1: numpy.ndarray, l1: numpy.ndarray) -> SeparationResult

      Separate H1 and L1 strain into signal and background components.

      :param h1: H1 strain. Shape ``(T,)`` or ``(N, T)``.
      :type h1: np.ndarray
      :param l1: L1 strain. Same shape as ``h1``.
      :type l1: np.ndarray

      :returns: Named tuple with fields ``h1_signal``, ``l1_signal``,
                ``h1_background``, ``l1_background``, each of the same shape as
                the inputs.
      :rtype: SeparationResult