<commit_before>"""
The main module of sslf (simple spectral line finder).
The provided Spectrum class is intended to facilitate all functionality.
"""

# Python 2 and 3 compatibility
from future.builtins import range
from future.utils import raise_with_traceback

import copy
import logging

import numpy as np
import numpy.ma as ma
from scipy import signal


logger = logging.getLogger(__name__)


def find_background_rms(array, num_chunks, use_chunks):
    """
    In order to obtain a better estimate of a spectrum's RMS, break the input
    array into evenly sized chunks (or as even as possible), sort the chunks
    by RMS, then return the average of some of these chunks as the true RMS.

    Parameters
    ----------
    array : ndarray
        The 1-dimensional spectrum we wish to determine the RMS for.

    num_chunks : int, optional
        Split `array` into this many chunks for RMS calculation.

    use_chunks : int, optional
        Use this many chunks from `array` for RMS calculation.
    """

    chunks = np.array_split(array, num_chunks)
    # The following requires np.errstate, because numpy.ma apparently throws
    # floating point errors erroneously.
    # https://github.com/numpy/numpy/issues/4895
    with np.errstate(under="ignore"):
        sorted_by_rms = sorted([np.std(x) for x in chunks])
    if logger.isEnabledFor(logging.DEBUG):
        logger.debug("rms chunks = %s" % sorted_by_rms)
    mean = np.mean(sorted_by_rms[:use_chunks])
    if mean == 0:
        raise_with_traceback(ValueError("RMS is 0, which cannot be used."))
    return mean


def _blank_spectrum_part(spectrum, point, radius, value=0):
    lower = max([0, point+int(-radius)])
    upper = min([len(spectrum), point+int(radius)])
    spectrum[lower:upper] = value


class _Peak(object):
    def __init__(self, channel, snr, width):
        self.channel = channel
        self.snr = snr
        self.width = width


class Spectrum(object):
    def __init__(self, spectrum, vel=None, num_chunks=5, use_chunks=3):
        """
        Provide a spectrum to find lines on, and/or remove the bandpass from.
        The optional vel parameter essentially provides an "x-axis" to the
        data, such that peak positions can be determined in terms of this
        axis rather than merely the channel position in the spectrum.
        Note that any NaN values in the spectrum are filtered.
        """
        # Make sure the data are in numpy arrays.
        if isinstance(spectrum, list):
            spectrum = np.array(spectrum)
        if isinstance(vel, list):
            vel = np.array(vel)

        # Filter any NaNs, scipy doesn't like them.
        if np.any(np.isnan(spectrum)):
            nan_indices = np.isnan(spectrum)
            self.original = spectrum[~nan_indices]
            if vel is not None:
                self.vel = vel[~nan_indices]
            else:
                self.vel = None
        else:
            self.original = spectrum
            self.vel = vel

        self.rms = find_background_rms(spectrum, num_chunks=num_chunks, use_chunks=use_chunks)


    def find_cwt_peaks(self, scales, snr=6.5, wavelet=signal.ricker,
                       num_chunks=5, use_chunks=3):
        """
        From the input spectrum (and a range of scales to search):
        - perform a CWT
        - find a significant peak in the CWT matrix
        - mask this peak in wavelet space, for all scales
        - loop from step 2, until no significant peaks remain
        - return the list of peaks

        In my experience, an SNR of 6.5 is a good compromise for reducing the number
        of false positives found while reliably finding real, significant peaks.

        It may be worthwhile to smooth the spectrum before performing the CWT.
        """<commit_msg>Added and updated documentation.<commit_after>"""
The main module of sslf (simple spectral line finder).
The provided Spectrum class is intended to facilitate all functionality.
"""

# Python 2 and 3 compatibility
from future.builtins import range
from future.utils import raise_with_traceback

import copy
import logging

import numpy as np
import numpy.ma as ma
from scipy import signal


logger = logging.getLogger(__name__)


def find_background_rms(array, num_chunks, use_chunks):
    """
    In order to obtain a better estimate of a spectrum's RMS, break the input
    array into evenly sized chunks (or as even as possible), sort the chunks
    by RMS, then return the average of some of these chunks as the true RMS.

    Parameters
    ----------
    array : array_like
        The 1-dimensional spectrum we wish to determine the RMS for.

    num_chunks : int, optional
        Split `array` into this many chunks for RMS calculation.

    use_chunks : int, optional
        Use this many chunks from `array` for RMS calculation.
    """

    chunks = np.array_split(array, num_chunks)
    # The following requires np.errstate, because numpy.ma apparently throws
    # floating point errors erroneously.
    # https://github.com/numpy/numpy/issues/4895
    with np.errstate(under="ignore"):
        sorted_by_rms = sorted([np.std(x) for x in chunks])
    if logger.isEnabledFor(logging.DEBUG):
        logger.debug("rms chunks = %s" % sorted_by_rms)
    mean = np.mean(sorted_by_rms[:use_chunks])
    if mean == 0:
        raise_with_traceback(ValueError("RMS is 0, which cannot be used."))
    return mean


def _blank_spectrum_part(spectrum, point, radius, value=0):
    lower = max([0, point+int(-radius)])
    upper = min([len(spectrum), point+int(radius)])
    spectrum[lower:upper] = value


class _Peak(object):
    def __init__(self, channel, snr, width):
        self.channel = channel
        self.snr = snr
        self.width = width


class Spectrum(object):
    def __init__(self, spectrum, vel=None, num_chunks=5, use_chunks=3):
        """
        Provide a spectrum to find lines on, and/or remove the bandpass from.
        The optional vel parameter essentially provides an "x-axis" to the
        data, such that peak positions can be determined in terms of this
        axis rather than merely the channel position in the spectrum.
        Note that any NaN values in the spectrum are filtered.

        Parameters
        ----------
        spectrum : array_like
            The 1-dimensional spectrum we will performing line-finding on.

        vel : array_like, optional
            Supply an equally-sized array to `spectrum` to have `sslf` report
            the velocities (or some other arbitrary interpretation of the x-axis
            of `spectrum`) of found lines.

        num_chunks : int, optional
            When calculating the RMS of `spectrum`, break `spectrum` into
            equally-sized chunks according to `num_chunks`.

        use_chunks : int, optional
            When calculating the RMS of `spectrum`, use this many equally-sized
            chunks from when `spectrum` was split according to `num_chunks`.
        """

        # Make sure the data are in numpy arrays.
        if isinstance(spectrum, list):
            spectrum = np.array(spectrum)
        if isinstance(vel, list):
            vel = np.array(vel)

        # Filter any NaNs, scipy doesn't like them.
        if np.any(np.isnan(spectrum)):
            nan_indices = np.isnan(spectrum)
            self.original = spectrum[~nan_indices]
            if vel is not None:
                self.vel = vel[~nan_indices]
            else:
                self.vel = None
        else:
            self.original = spectrum
            self.vel = vel

        self.rms = find_background_rms(spectrum, num_chunks=num_chunks, use_chunks=use_chunks)


    def find_cwt_peaks(self, scales, snr=6.5, wavelet=signal.ricker,
                       num_chunks=5, use_chunks=3):
        """
        From the input spectrum (and a range of scales to search):
        - perform a continuous wavelet transformation (CWT)
        - find a significant peak in the CWT matrix
        - mask this peak in wavelet space, for all scales
        - loop from step 2, until no significant peaks remain
        - return the list of peaks

        In my experience, an SNR of 6.5 is a good compromise for reducing the number
        of false positives found while reliably finding real, significant peaks.

        It may be worthwhile to smooth the spectrum before performing the CWT.

        Parameters
        ----------
        scales : array_like
            The channel-widths to use when performing the CWT, and thus the
            channel-widths we use when searching for spectral lines.

        snr : float, optional
            The significance threshold of any found spectral line.

        wavelet : function
            The wavelet to use when performing the CWT. The wavelet used should
            have a similar profile to any spectral lines that are being
            searched.

        num_chunks : int, optional
            When calculating the RMS of `spectrum`, break `spectrum` into
            equally-sized chunks according to `num_chunks`.

        use_chunks : int, optional
            When calculating the RMS of `spectrum`, use this many equally-sized
            chunks from when `spectrum` was split according to `num_chunks`.
        """
{'fraction_non_alphanumeric': 0.07458175930922828, 'fraction_numerical': 0.004425256341068538, 'mean_word_length': 2.4382189239332095, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2072719', 'n_tokens_mistral': 2575, 'n_tokens_neox': 2443, 'n_words': 994}