<commit_before># -*- coding: utf-8 -*-
# Author: <NAME> <<EMAIL>>

""" Classes for spectral analysis. """

from __future__ import absolute_import

import datetime

from random import randint
from itertools import izip
from copy import copy, deepcopy
from math import floor

import numpy as np
from numpy import ma

from scipy import ndimage

from matplotlib import pyplot as plt
from matplotlib.ticker import FuncFormatter, MaxNLocator, IndexLocator
from matplotlib.colorbar import Colorbar

from sunpy.time import parse_time
from sunpy.util.util import to_signed
from sunpy.spectra.spectrum import Spectrum

# This should not be necessary, as observations do not take more than a day
# but it is used for completeness' and extendibility's sake.
# XXX: Leap second?
SECONDS_PER_DAY = 86400

# Used for COPY_PROPERTIES
REFERENCE = 0
COPY = 1
DEEPCOPY = 2

# Maybe move to util.
def get_day(dt):
    """ Return datetime for the beginning of the day of given datetime. """
    return datetime.datetime(dt.year, dt.month, dt.day)


def min_delt(arr):
    deltas = (arr[:-1] - arr[1:])
    # Multiple values at the same frequency are just thrown away
    # in the process of linearizaion
    return deltas[deltas != 0].min()


def list_formatter(lst, fun=None):
    """ Return function that takes x, pos and returns fun(lst[x]) if
    fun is not None, else lst[x] or "" if x is out of range. """
    def _fun(x, pos):
        x = int(x)
        if x >= len(lst) or x < 0:
            return ""
        
        elem = lst[x]
        if fun is None:
            return elem
        return fun(elem)
    return _fun


def _union(sets):
    """ Return union of sets. """
    union = set()
    for s in sets:
        union |= s
    return union


class _AttrGetter(object):
    """ Helper class for frequency channel linearization.
    
    Parameters
    ----------
    arr : Spectrogram
        Spectrogram to linearize.
    delt : float
        Delta between frequency channels in linearized spectrogram. Defaults to
        (minimum delta / 2.) because of the Shannon sampling theorem.
    """
    def __init__(self, arr, delt=None):
        self.arr = arr
        if delt is None:
            # Nyquist–Shannon sampling theorem
            delt = min_delt(arr.freq_axis) / 2.
        
        self.delt = delt
        
        midpoints =(self.arr.freq_axis[:-1] + self.arr.freq_axis[1:]) / 2
        self.midpoints = np.concatenate([midpoints, arr.freq_axis[-1:]])
        
        self.shape = (len(self), arr.shape[1])
    
    def __len__(self):
        return 1 + (self.arr.freq_axis[0] - self.arr.freq_axis[-1]) / self.delt
    
    def __getitem__(self, item):
        freq = self.arr.freq_axis[0] - item * self.delt
        for n, mid in enumerate(self.midpoints):
            if mid <= freq:
                return self.arr[n, :]
        raise IndexError
    
    

# XXX: Find out why imshow(x) fails!
class Spectrogram(np.ndarray):
    """ Base class for spectral analysis in SunPy.
    
    Parameters
    ----------
    data : np.ndarray
        two-dimensional array of the image data of the spectrogram.
    time_axis : np.ndarray
        one-dimensional array containing the offset from the start
        for each column of data.
    freq_axis : np.ndarray
        one-dimensional array containing information about the
        frequencies each row of the image corresponds to.
    start : datetime
        starting time of the measurement
    end : datetime
        end time of the measurement
    t_init : int
        offset from the start of the day the measurement began. If None
        gets automatically set from start.
    t_label : str
        label for the time axis
    f_label : str
        label for the frequency axis
    content : str
        header for the image
    """
    # Contrary to what pylint may think, this is not an old-style class.
    # pylint: disable=E1002,W0142,R0902

    # This needs to list all attributes that need to be
    # copied to maintain the object and how to handle them.
    COPY_PROPERTIES = [
        ('time_axis', COPY),
        ('freq_axis', COPY),
        ('instruments', COPY),
        ('start', REFERENCE),
        ('end', REFERENCE),
        ('t_label', REFERENCE),
        ('f_label', REFERENCE),
        ('content', REFERENCE),
        ('t_init', REFERENCE),
    ]

    def as_class(self, cls):
        """ Implementation detail. """
        if not issubclass(cls, Spectrogram):
            raise ValueError

        dct = {}
        var = vars(self)
        for prop, _ in cls.COPY_PROPERTIES:
            if not prop in var:
                raise ValueError
            dct[prop] = var[prop]
        return cls(self, **dct)

    def get_params(self):
        """ Implementation detail. """
        return dict(
            (name, getattr(self, name)) for name, _ in self.COPY_PROPERTIES
        )
    
    def _slice(self, y_range, x_range):
        """ Return new spectrogram reduced to the values passed
        as slices. Implementation detail. """
        data = super(Spectrogram, self).__getitem__([y_range, x_range])
        params = self.get_params()

        soffset = 0 if x_range.start is None else x_range.start
        eoffset = self.shape[1] if x_range.stop is None else x_range.stop # pylint: disable=E1101
        eoffset -= 1

        fsoffset = 0 if y_range.start is None else y_range.start
        feoffset = self.shape[0] if y_range.stop is None else y_range.stop # pylint: disable=E1101
        
        params.update({
            'time_axis': self.time_axis[
                x_range.start:x_range.stop:x_range.step
            ] - self.time_axis[soffset],
            'freq_axis': self.freq_axis[
                y_range.start:y_range.stop:y_range.step],
            'start': self.start + datetime.timedelta(
                seconds=self.time_axis[soffset]),
            'end': self.start + datetime.timedelta(
                seconds=self.time_axis[eoffset]),
            't_init': self.t_init + self.time_axis[soffset],
        })
        return self.__class__(data, **params)

    # This accepting arbitrary arguments makes it easier to subclass this.
    def __new__(cls, data, *args, **kwargs):
        return np.asarray(data).view(cls)

    def __init__(self, data, time_axis, freq_axis, start, end, t_init=None,
        t_label="Time", f_label="Frequency", content="", instruments=None):
        # Because of how object creation works, there is no avoiding
        # unused arguments in this case.
        if t_init is None:
            diff = start - get_day(start)
            t_init = diff.seconds
        if instruments is None:
            instruments = set()
        
        self.start = start
        self.end = end

        self.t_label = t_label
        self.f_label = f_label

        self.t_init = t_init

        self.time_axis = time_axis
        self.freq_axis = freq_axis

        self.content = content
        self.instruments = instruments

    def time_formatter(self, x, pos):
        """ This returns the label for the tick of value x at
        a specified pos on the time axis. """
        # Callback, cannot avoid unused arguments.
        # pylint: disable=W0613
        x = int(x)
        if x >= len(self.time_axis) or x < 0:
            return ""
        return self.format_time(
            self.start + datetime.timedelta(
                seconds=self.time_axis[x]
            )
        )

    def __array_finalize__(self, obj):
        if self is obj:
            return

        for prop, cpy in self.COPY_PROPERTIES:
            elem = getattr(obj, prop, None)
            if cpy == COPY:
                elem = copy(elem)
            if cpy == DEEPCOPY:
                elem = deepcopy(elem)

            setattr(self, prop, elem)
    
    @staticmethod
    def format_time(time):
        """ Override to configure default plotting """
        return time.strftime("%H:%M:%S")
    
    @staticmethod
    def format_freq(freq):
        """ Override to configure default plotting """
        return "%.1f" % freq

    def show(self, *args, **kwargs):
        """ Draw spectrogram on figure with highest index or new one if
        none exists. For parameters see :py:meth:`plot`. """
        self.plot(*args, **kwargs).show()

    def plot(self, figure=None, overlays=[], colorbar=True, min_=None, max_=None,
             linear=True, showz=True, **matplotlib_args):
        """
        Plot spectrogram onto figure.
        
        Parameters
        ----------
        figure : matplotlib.figure.Figure
            Figure to plot the spectrogram on. If None, new Figure is created.
        overlays : list
            List of overlays (functions that receive figure and axes and return
            new ones) to be applied after drawing.
        colorbar : bool
            Flag that determines whether or not to draw a colorbar. If existing
            figure is passed, it is attempted to overdraw old colorbar.
        min_ : float
            Clip intensities lower than min_ before drawing.
        max_ : float
            Clip intensities higher than max_ before drawing.
        linear :  bool
            If set to True, "stretch" image to make frequency axis linear.
        showz : bool
            If set to True, the value of the pixel that is hovered with the
            mouse is shown in the bottom right corner.
        """
        # [] as default argument is okay here because it is only read.
        # pylint: disable=W0102,R0914
        if linear:
            data = _AttrGetter(self.clip(min_, max_))
            freqs = np.arange(
                self.freq_axis[0], self.freq_axis[-1], -data.delt
            )
        else:
            data = np.array(self.clip(min_, max_))
            freqs = self.freq_axis
        newfigure = figure is None
        if figure is None:
            figure = plt.figure(frameon=True)
            axes = figure.add_subplot(111)
        else:
            axes = figure.axes[0]
        
        params = {
            'origin': 'lower',
            'aspect': 'auto',
        }
        params.update(matplotlib_args)
        im = axes.imshow(data, **params)
        
        xa = axes.get_xaxis()
        ya = axes.get_yaxis()

        xa.set_major_formatter(
            FuncFormatter(self.time_formatter)
        )
        
        if linear:
            # XXX: Why is the x off by 0.5 with the IndexLocator?
            init = (self.freq_axis[0] % 5) / data.delt + 0.5
            nticks = 15.
            dist = (self.freq_axis[0] - self.freq_axis[-1]) / nticks
            dist = max(1, (dist // 10)) * 10
            
            ya.set_major_locator(<commit_msg>Update comment.<commit_after># -*- coding: utf-8 -*-
# Author: <NAME> <<EMAIL>>

""" Classes for spectral analysis. """

from __future__ import absolute_import

import datetime

from random import randint
from itertools import izip
from copy import copy, deepcopy
from math import floor

import numpy as np
from numpy import ma

from scipy import ndimage

from matplotlib import pyplot as plt
from matplotlib.ticker import FuncFormatter, MaxNLocator, IndexLocator
from matplotlib.colorbar import Colorbar

from sunpy.time import parse_time
from sunpy.util.util import to_signed
from sunpy.spectra.spectrum import Spectrum

# This should not be necessary, as observations do not take more than a day
# but it is used for completeness' and extendibility's sake.
# XXX: Leap second?
SECONDS_PER_DAY = 86400

# Used for COPY_PROPERTIES
REFERENCE = 0
COPY = 1
DEEPCOPY = 2

# Maybe move to util.
def get_day(dt):
    """ Return datetime for the beginning of the day of given datetime. """
    return datetime.datetime(dt.year, dt.month, dt.day)


def min_delt(arr):
    deltas = (arr[:-1] - arr[1:])
    # Multiple values at the same frequency are just thrown away
    # in the process of linearizaion
    return deltas[deltas != 0].min()


def list_formatter(lst, fun=None):
    """ Return function that takes x, pos and returns fun(lst[x]) if
    fun is not None, else lst[x] or "" if x is out of range. """
    def _fun(x, pos):
        x = int(x)
        if x >= len(lst) or x < 0:
            return ""
        
        elem = lst[x]
        if fun is None:
            return elem
        return fun(elem)
    return _fun


def _union(sets):
    """ Return union of sets. """
    union = set()
    for s in sets:
        union |= s
    return union


class _AttrGetter(object):
    """ Helper class for frequency channel linearization.
    
    Parameters
    ----------
    arr : Spectrogram
        Spectrogram to linearize.
    delt : float
        Delta between frequency channels in linearized spectrogram. Defaults to
        (minimum delta / 2.) because of the Shannon sampling theorem.
    """
    def __init__(self, arr, delt=None):
        self.arr = arr
        if delt is None:
            # Nyquist–Shannon sampling theorem
            delt = min_delt(arr.freq_axis) / 2.
        
        self.delt = delt
        
        midpoints =(self.arr.freq_axis[:-1] + self.arr.freq_axis[1:]) / 2
        self.midpoints = np.concatenate([midpoints, arr.freq_axis[-1:]])
        
        self.shape = (len(self), arr.shape[1])
    
    def __len__(self):
        return 1 + (self.arr.freq_axis[0] - self.arr.freq_axis[-1]) / self.delt
    
    def __getitem__(self, item):
        freq = self.arr.freq_axis[0] - item * self.delt
        for n, mid in enumerate(self.midpoints):
            if mid <= freq:
                return self.arr[n, :]
        raise IndexError
    
    

# XXX: Find out why imshow(x) fails!
class Spectrogram(np.ndarray):
    """ Base class for spectral analysis in SunPy.
    
    Parameters
    ----------
    data : np.ndarray
        two-dimensional array of the image data of the spectrogram.
    time_axis : np.ndarray
        one-dimensional array containing the offset from the start
        for each column of data.
    freq_axis : np.ndarray
        one-dimensional array containing information about the
        frequencies each row of the image corresponds to.
    start : datetime
        starting time of the measurement
    end : datetime
        end time of the measurement
    t_init : int
        offset from the start of the day the measurement began. If None
        gets automatically set from start.
    t_label : str
        label for the time axis
    f_label : str
        label for the frequency axis
    content : str
        header for the image
    """
    # Contrary to what pylint may think, this is not an old-style class.
    # pylint: disable=E1002,W0142,R0902

    # This needs to list all attributes that need to be
    # copied to maintain the object and how to handle them.
    COPY_PROPERTIES = [
        ('time_axis', COPY),
        ('freq_axis', COPY),
        ('instruments', COPY),
        ('start', REFERENCE),
        ('end', REFERENCE),
        ('t_label', REFERENCE),
        ('f_label', REFERENCE),
        ('content', REFERENCE),
        ('t_init', REFERENCE),
    ]

    def as_class(self, cls):
        """ Implementation detail. """
        if not issubclass(cls, Spectrogram):
            raise ValueError

        dct = {}
        var = vars(self)
        for prop, _ in cls.COPY_PROPERTIES:
            if not prop in var:
                raise ValueError
            dct[prop] = var[prop]
        return cls(self, **dct)

    def get_params(self):
        """ Implementation detail. """
        return dict(
            (name, getattr(self, name)) for name, _ in self.COPY_PROPERTIES
        )
    
    def _slice(self, y_range, x_range):
        """ Return new spectrogram reduced to the values passed
        as slices. Implementation detail. """
        data = super(Spectrogram, self).__getitem__([y_range, x_range])
        params = self.get_params()

        soffset = 0 if x_range.start is None else x_range.start
        eoffset = self.shape[1] if x_range.stop is None else x_range.stop # pylint: disable=E1101
        eoffset -= 1

        fsoffset = 0 if y_range.start is None else y_range.start
        feoffset = self.shape[0] if y_range.stop is None else y_range.stop # pylint: disable=E1101
        
        params.update({
            'time_axis': self.time_axis[
                x_range.start:x_range.stop:x_range.step
            ] - self.time_axis[soffset],
            'freq_axis': self.freq_axis[
                y_range.start:y_range.stop:y_range.step],
            'start': self.start + datetime.timedelta(
                seconds=self.time_axis[soffset]),
            'end': self.start + datetime.timedelta(
                seconds=self.time_axis[eoffset]),
            't_init': self.t_init + self.time_axis[soffset],
        })
        return self.__class__(data, **params)

    # This accepting arbitrary arguments makes it easier to subclass this.
    def __new__(cls, data, *args, **kwargs):
        return np.asarray(data).view(cls)

    def __init__(self, data, time_axis, freq_axis, start, end, t_init=None,
        t_label="Time", f_label="Frequency", content="", instruments=None):
        # Because of how object creation works, there is no avoiding
        # unused arguments in this case.
        if t_init is None:
            diff = start - get_day(start)
            t_init = diff.seconds
        if instruments is None:
            instruments = set()
        
        self.start = start
        self.end = end

        self.t_label = t_label
        self.f_label = f_label

        self.t_init = t_init

        self.time_axis = time_axis
        self.freq_axis = freq_axis

        self.content = content
        self.instruments = instruments

    def time_formatter(self, x, pos):
        """ This returns the label for the tick of value x at
        a specified pos on the time axis. """
        # Callback, cannot avoid unused arguments.
        # pylint: disable=W0613
        x = int(x)
        if x >= len(self.time_axis) or x < 0:
            return ""
        return self.format_time(
            self.start + datetime.timedelta(
                seconds=self.time_axis[x]
            )
        )

    def __array_finalize__(self, obj):
        if self is obj:
            return

        for prop, cpy in self.COPY_PROPERTIES:
            elem = getattr(obj, prop, None)
            if cpy == COPY:
                elem = copy(elem)
            if cpy == DEEPCOPY:
                elem = deepcopy(elem)

            setattr(self, prop, elem)
    
    @staticmethod
    def format_time(time):
        """ Override to configure default plotting """
        return time.strftime("%H:%M:%S")
    
    @staticmethod
    def format_freq(freq):
        """ Override to configure default plotting """
        return "%.1f" % freq

    def show(self, *args, **kwargs):
        """ Draw spectrogram on figure with highest index or new one if
        none exists. For parameters see :py:meth:`plot`. """
        self.plot(*args, **kwargs).show()

    def plot(self, figure=None, overlays=[], colorbar=True, min_=None, max_=None,
             linear=True, showz=True, **matplotlib_args):
        """
        Plot spectrogram onto figure.
        
        Parameters
        ----------
        figure : matplotlib.figure.Figure
            Figure to plot the spectrogram on. If None, new Figure is created.
        overlays : list
            List of overlays (functions that receive figure and axes and return
            new ones) to be applied after drawing.
        colorbar : bool
            Flag that determines whether or not to draw a colorbar. If existing
            figure is passed, it is attempted to overdraw old colorbar.
        min_ : float
            Clip intensities lower than min_ before drawing.
        max_ : float
            Clip intensities higher than max_ before drawing.
        linear :  bool
            If set to True, "stretch" image to make frequency axis linear.
        showz : bool
            If set to True, the value of the pixel that is hovered with the
            mouse is shown in the bottom right corner.
        """
        # [] as default argument is okay here because it is only read.
        # pylint: disable=W0102,R0914
        if linear:
            data = _AttrGetter(self.clip(min_, max_))
            freqs = np.arange(
                self.freq_axis[0], self.freq_axis[-1], -data.delt
            )
        else:
            data = np.array(self.clip(min_, max_))
            freqs = self.freq_axis
        newfigure = figure is None
        if figure is None:
            figure = plt.figure(frameon=True)
            axes = figure.add_subplot(111)
        else:
            axes = figure.axes[0]
        
        params = {
            'origin': 'lower',
            'aspect': 'auto',
        }
        params.update(matplotlib_args)
        im = axes.imshow(data, **params)
        
        xa = axes.get_xaxis()
        ya = axes.get_yaxis()

        xa.set_major_formatter(
            FuncFormatter(self.time_formatter)
        )
        
        if linear:
            # 0.5 is because matplotlib seems to center it.
            init = (self.freq_axis[0] % 5) / data.delt + 0.5
            nticks = 15.
            dist = (self.freq_axis[0] - self.freq_axis[-1]) / nticks
            dist = max(1, (dist // 10)) * 10
            
            ya.set_major_locator(
{'fraction_non_alphanumeric': 0.09328078078078078, 'fraction_numerical': 0.007882882882882882, 'mean_word_length': 1.9580846634281748, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, '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': '5029957', 'n_tokens_mistral': 6384, 'n_tokens_neox': 6156, 'n_words': 2015}