Exemplo n.º 1
0
def test_nocond():
    f = ConditionalDispatch()
    # Multiply even numbers by two.
    f.add(lambda x: 2 * x, lambda x: x % 2 == 0)
    with pytest.raises(TypeError) as exc_info:
        f(3)
        assert "Your input did not fulfill the condition for any function." in str(
            exc_info.value)
Exemplo n.º 2
0
def test_else2():
    # This verifies else branches do not catch cases that are covered
    # by cases added later.
    f = ConditionalDispatch()
    # Because gcd(2, 3) == 1, 2 | x and 3 | x are mutually exclusive.
    f.add(lambda x: 2 * x, lambda x: x % 2 == 0)
    f.add(lambda x: 3 * x)
    f.add(lambda x: 4 * x, lambda x: x % 3 == 0)
    assert f(2) == 4
    assert f(3) == 12
    assert f(5) == 15
Exemplo n.º 3
0
def test_else():
    f = ConditionalDispatch()
    # Multiply even numbers by two.
    f.add(lambda x: 2 * x, lambda x: x % 2 == 0)
    f.add(lambda x: 3 * x)
    assert f(2) == 4
    assert f(3) == 9
Exemplo n.º 4
0
def oddeven(request):
    f = ConditionalDispatch()
    # Multiply even numbers by two.
    f.add(lambda x: 2 * x, lambda x: x % 2 == 0)
    # Multiply odd numbers by three.
    f.add(lambda x: 3 * x, lambda x: x % 2 == 1)
    return f
Exemplo n.º 5
0
class SWavesSpectrogram(LinearTimeSpectrogram):
    _create = ConditionalDispatch.from_existing(LinearTimeSpectrogram._create)
    create = classmethod(_create.wrapper())
    COPY_PROPERTIES = LinearTimeSpectrogram.COPY_PROPERTIES + [
        ('bg', REFERENCE)
    ]

    @staticmethod
    def swavesfile_to_date(filename):
        _, name = os.path.split(filename)
        date = name.split('_')[2]
        return datetime.datetime(
            int(date[0:4]), int(date[4:6]), int(date[6:])
        )

    @classmethod
    def read(cls, filename, **kwargs):
        """
        Read in FITS file and return a new SWavesSpectrogram.
        """
        data = np.genfromtxt(filename, skip_header=2)
        time_axis = data[:, 0] * 60.
        data = data[:, 1:].transpose()
        header = np.genfromtxt(filename, skip_footer=time_axis.size)
        freq_axis = header[0, :]
        bg = header[1, :]
        start = cls.swavesfile_to_date(filename)
        end = start + datetime.timedelta(seconds=time_axis[-1])
        t_delt = 60.
        t_init = (start - get_day(start)).seconds
        content = ''
        t_label = 'Time [UT]'
        f_label = 'Frequency [KHz]'

        freq_axis = freq_axis[::-1]
        data = data[::-1, :]

        return cls(data, time_axis, freq_axis, start, end, t_init, t_delt,
                   t_label, f_label, content, bg)

    def __init__(self, data, time_axis, freq_axis, start, end,
                 t_init, t_delt, t_label, f_label, content, bg):
        # Because of how object creation works, there is no avoiding
        # unused arguments in this case.
        # pylint: disable=W0613

        super(SWavesSpectrogram, self).__init__(
            data, time_axis, freq_axis, start, end,
            t_init, t_delt, t_label, f_label,
            content, set(["SWAVES"])
        )
        self.bg = bg
Exemplo n.º 6
0
class EOVSASpectrogram(LinearTimeSpectrogram):
    _create = ConditionalDispatch.from_existing(LinearTimeSpectrogram._create)
    create = classmethod(_create.wrapper())
    COPY_PROPERTIES = LinearTimeSpectrogram.COPY_PROPERTIES + [
        ('bg', REFERENCE)
    ]

    @classmethod
    def read(cls, filename, **kwargs):
        """
        Read in FITS file and return a new SWavesSpectrogram.
        """
        fl = fits.open(filename, **kwargs)
        data = fl[0].data
        freq_axis = fl[1].data.sfreq
        tdata = fl[2].data
        mjd = tdata.mjd
        sec = tdata.time / 1000.
        start = Time(mjd[0] + sec[0] / 86400., format='mjd').datetime
        end = Time(mjd[-1] + sec[-1] / 86400., format='mjd').datetime
        time_axis = sec - sec[0]
        header = fl[0].header
        t_delt = 1.0
        t_init = (start - get_day(start)).seconds
        content = 'EOVSA'
        t_label = 'Time [UT]'
        f_label = 'Frequency [GHz]'
        #reverse the frequency axis
        freq_axis = freq_axis[::-1]
        data = data[::-1, :]

        return cls(data, time_axis, freq_axis, start, end, t_init, t_delt,
                   t_label, f_label, content)

    def __init__(self, data, time_axis, freq_axis, start, end, t_init, t_delt,
                 t_label, f_label, content):
        # Because of how object creation works, there is no avoiding
        # unused arguments in this case.
        # pylint: disable=W0613

        super(EOVSASpectrogram,
              self).__init__(data, time_axis, freq_axis, start, end, t_init,
                             t_delt, t_label, f_label, content, set(["EOVSA"]))
Exemplo n.º 7
0
class Spectrogram(Parent):
    """
    Spectrogram Class.

    .. warning:: This module is under development! Use at your own risk.

    Attributes
    ----------
    data : `~numpy.ndarray`
        two-dimensional array of the image data of the spectrogram.
    time_axis : `~numpy.ndarray`
        one-dimensional array containing the offset from the start
        for each column of data.
    freq_axis : `~numpy.ndarray`
        one-dimensional array containing information about the
        frequencies each row of the image corresponds to.
    start : `~datetime.datetime`
        starting time of the measurement
    end : `~datetime.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
    instruments : str array
        instruments that recorded the data, may be more than one if
        it was constructed using combine_frequencies or join_many.
    """
    # 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),
    ]
    _create = ConditionalDispatch.from_existing(Parent._create)

    @property
    def shape(self):
        return self.data.shape

    @property
    def dtype(self):
        return self.data.dtype

    def _get_params(self):
        """
        Implementation detail.
        """
        return {
            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 = self.data[y_range, x_range]
        params = self._get_params()

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

        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)

    def _with_data(self, data):
        new = copy(self)
        new.data = data
        return new

    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.
        self.data = data

        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=float(self.time_axis[x])
            )
        )

    @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 "{freq:0.1f}".format(freq=freq)

    def peek(self, *args, **kwargs):
        """
        Plot spectrum onto current axes.

        Parameters
        ----------
        *args : dict

        **kwargs : dict
            Any additional plot arguments that should be used
            when plotting.

        Returns
        -------
        fig : `~matplotlib.Figure`
            A plot figure.
        """
        figure()
        ret = self.plot(*args, **kwargs)
        plt.show()
        return ret

    def plot(self, figure=None, overlays=[], colorbar=True, vmin=None,
             vmax=None, linear=True, showz=True, yres=DEFAULT_YRES,
             max_dist=None, **matplotlib_args):
        """
        Plot spectrogram onto figure.

        Parameters
        ----------
        figure : `~matplotlib.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.
        vmin : float
            Clip intensities lower than vmin before drawing.
        vmax : float
            Clip intensities higher than vmax 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.
        yres : int or None
            To be used in combination with linear=True. If None, sample the
            image with half the minimum frequency delta. Else, sample the
            image to be at most yres pixels in vertical dimension. Defaults
            to 1080 because that's a common screen size.
        max_dist : float or None
            If not None, mask elements that are further than max_dist away
            from actual data points (ie, frequencies that actually have data
            from the receiver and are not just nearest-neighbour interpolated).
        """
        # [] as default argument is okay here because it is only read.
        # pylint: disable=W0102,R0914
        if linear:
            delt = yres
            if delt is not None:
                delt = max(
                    (self.freq_axis[0] - self.freq_axis[-1]) / (yres - 1),
                    _min_delt(self.freq_axis) / 2.
                )
                delt = float(delt)

            data = _LinearView(self.clip_values(vmin, vmax), delt)
            freqs = np.arange(
                self.freq_axis[0], self.freq_axis[-1], -data.delt
            )
        else:
            data = np.array(self.clip_values(vmin, vmax))
            freqs = self.freq_axis

        figure = plt.gcf()

        if figure.axes:
            axes = figure.axes[0]
        else:
            axes = figure.add_subplot(111)

        params = {
            'origin': 'lower',
            'aspect': 'auto',
        }
        params.update(matplotlib_args)
        if linear and max_dist is not None:
            toplot = ma.masked_array(data, mask=data.make_mask(max_dist))
        else:
            toplot = data
        im = axes.imshow(toplot, **params)

        xa = axes.get_xaxis()
        ya = axes.get_yaxis()

        xa.set_major_formatter(
            FuncFormatter(self.time_formatter)
        )

        if linear:
            # Start with a number that is divisible by 5.
            init = (self.freq_axis[0] % 5) / data.delt
            nticks = 15.
            # Calculate MHz difference between major ticks.
            dist = (self.freq_axis[0] - self.freq_axis[-1]) / nticks
            # Round to next multiple of 10, at least ten.
            dist = max(round(dist, -1), 10)
            # One pixel in image space is data.delt MHz, thus we can convert
            # our distance between the major ticks into image space by dividing
            # it by data.delt.

            ya.set_major_locator(
                IndexLocator(
                    dist / data.delt, init
                )
            )
            ya.set_minor_locator(
                IndexLocator(
                    dist / data.delt / 10, init
                )
            )

            def freq_fmt(x, pos):
                # This is necessary because matplotlib somehow tries to get
                # the mid-point of the row, which we do not need here.
                x = x + 0.5
                return self.format_freq(self.freq_axis[0] - x * data.delt)
        else:
            freq_fmt = _list_formatter(freqs, self.format_freq)
            ya.set_major_locator(MaxNLocator(integer=True, steps=[1, 5, 10]))

        ya.set_major_formatter(
            FuncFormatter(freq_fmt)
        )

        axes.set_xlabel(self.t_label)
        axes.set_ylabel(self.f_label)
        # figure.suptitle(self.content)

        figure.suptitle(
            ' '.join([
                get_day(self.start).strftime("%d %b %Y"),
                'Radio flux density',
                '(' + ', '.join(self.instruments) + ')',
            ])
        )

        for tl in xa.get_ticklabels():
            tl.set_fontsize(10)
            tl.set_rotation(30)
        figure.add_axes(axes)
        figure.subplots_adjust(bottom=0.2)
        figure.subplots_adjust(left=0.2)

        if showz:
            axes.format_coord = self._mk_format_coord(
                data, figure.gca().format_coord)

        if colorbar:
            if len(figure.axes) > 1:
                Colorbar(figure.axes[1], im).set_label("Intensity")
            else:
                figure.colorbar(im).set_label("Intensity")

        for overlay in overlays:
            figure, axes = overlay(figure, axes)

        for ax in figure.axes:
            ax.autoscale()
        if isinstance(figure, SpectroFigure):
            figure._init(self, freqs)
        return axes

    def __getitem__(self, key):
        only_y = not isinstance(key, tuple)

        if only_y:
            return self.data[int(key)]
        elif isinstance(key[0], slice) and isinstance(key[1], slice):
            return self._slice(key[0], key[1])
        elif isinstance(key[1], slice):
            # return Spectrum( # XXX: Right class
            #     super(Spectrogram, self).__getitem__(key),
            #     self.time_axis[key[1].start:key[1].stop:key[1].step]
            # )
            return np.array(self.data[key])
        elif isinstance(key[0], slice):
            return Spectrum(
                self.data[key],
                self.freq_axis[key[0].start:key[0].stop:key[0].step]
            )

        return self.data[int(key)]

    def clip_freq(self, vmin=None, vmax=None):
        """
        Return a new spectrogram only consisting of frequencies in the interval
        [vmin, vmax].

        Parameters
        ----------
        vmin : float
            All frequencies in the result are greater or equal to this.
        vmax : float
            All frequencies in the result are smaller or equal to this.
        """
        left = 0
        if vmax is not None:
            while self.freq_axis[left] > vmax:
                left += 1

        right = len(self.freq_axis) - 1

        if vmin is not None:
            while self.freq_axis[right] < vmin:
                right -= 1

        return self[left:right + 1, :]

    def auto_find_background(self, amount=0.05):
        """
        Automatically find the background. This is done by first subtracting
        the average value in each channel and then finding those times which
        have the lowest standard deviation.

        Parameters
        ----------
        amount : float
            The percent amount (out of 1) of lowest standard deviation to
            consider.
        """
        # pylint: disable=E1101,E1103
        data = self.data.astype(to_signed(self.dtype))
        # Subtract average value from every frequency channel.
        tmp = (data - np.average(self.data, 1).reshape(self.shape[0], 1))
        # Get standard deviation at every point of time.
        # Need to convert because otherwise this class's __getitem__
        # is used which assumes two-dimensionality.
        sdevs = np.asarray(np.std(tmp, 0))

        # Get indices of values with lowest standard deviation.
        cand = sorted(list(range(self.shape[1])), key=lambda y: sdevs[y])
        # Only consider the best 5 %.
        return cand[:max(1, int(amount * len(cand)))]

    def auto_const_bg(self):
        """
        Automatically determine background.
        """
        realcand = self.auto_find_background()
        bg = np.average(self.data[:, realcand], 1)
        return bg.reshape(self.shape[0], 1)

    def subtract_bg(self):
        """
        Perform constant background subtraction.
        """
        return self._with_data(self.data - self.auto_const_bg())

    def randomized_auto_const_bg(self, amount):
        """
        Automatically determine background. Only consider a randomly chosen
        subset of the image.

        Parameters
        ----------
        amount : int
            Size of random sample that is considered for calculation of
            the background.
        """
        cols = [randint(0, self.shape[1] - 1) for _ in range(amount)]

        # pylint: disable=E1101,E1103
        data = self.data.astype(to_signed(self.dtype))
        # Subtract average value from every frequency channel.
        tmp = (data - np.average(self.data, 1).reshape(self.shape[0], 1))
        # Get standard deviation at every point of time.
        # Need to convert because otherwise this class's __getitem__
        # is used which assumes two-dimensionality.
        tmp = tmp[:, cols]
        sdevs = np.asarray(np.std(tmp, 0))

        # Get indices of values with lowest standard deviation.
        cand = sorted(list(range(amount)), key=lambda y: sdevs[y])
        # Only consider the best 5 %.
        realcand = cand[:max(1, int(0.05 * len(cand)))]

        # Average the best 5 %
        bg = np.average(self[:, [cols[r] for r in realcand]], 1)

        return bg.reshape(self.shape[0], 1)

    def randomized_subtract_bg(self, amount):
        """
        Perform randomized constant background subtraction. Does not produce
        the same result every time it is run.

        Parameters
        ----------
        amount : int
            Size of random sample that is considered for calculation of
            the background.
        """
        return self._with_data(self.data - self.randomized_auto_const_bg(amount))

    def clip_values(self, vmin=None, vmax=None, out=None):
        """
        Clip intensities to be in the interval [vmin, vmax].

        Any values greater than the maximum will be assigned the maximum,
        any values lower than the minimum will be assigned the minimum.
        If either is left out or None, do not clip at that side of the interval.

        Parameters
        ----------
        min : int or float
            New minimum value for intensities.
        max : int or float
            New maximum value for intensities
        """
        # pylint: disable=E1101
        if vmin is None:
            vmin = int(self.data.min())

        if vmax is None:
            vmax = int(self.data.max())

        return self._with_data(self.data.clip(vmin, vmax, out))

    def rescale(self, vmin=0, vmax=1, dtype=np.dtype('float32')):
        """
        Rescale intensities to [vmin, vmax]. Note that vmin ≠ vmax and
        spectrogram.min() ≠ spectrogram.max().

        Parameters
        ----------
        vmin : float or int
            New minimum value in the resulting spectrogram.
        vmax : float or int
            New maximum value in the resulting spectrogram.
        dtype : `numpy.dtype`
            Data-type of the resulting spectrogram.
        """
        if vmax == vmin:
            raise ValueError("Maximum and minimum must be different.")
        if self.data.max() == self.data.min():
            raise ValueError("Spectrogram needs to contain distinct values.")
        data = self.data.astype(dtype)  # pylint: disable=E1101
        return self._with_data(
            vmin + (vmax - vmin) * (data - self.data.min()) /  # pylint: disable=E1101
            (self.data.max() - self.data.min())  # pylint: disable=E1101
        )

    def interpolate(self, frequency):
        """
        Linearly interpolate intensity at unknown frequency using linear
        interpolation of its two neighbours.

        Parameters
        ----------
        frequency : float or int
            Unknown frequency for which to linearly interpolate the intensities.
            freq_axis[0] >= frequency >= self_freq_axis[-1]
        """
        lfreq, lvalue = None, None
        for freq, value in zip(self.freq_axis, self.data[:, :]):
            if freq < frequency:
                break
            lfreq, lvalue = freq, value
        else:
            raise ValueError("Frequency not in interpolation range")
        if lfreq is None:
            raise ValueError("Frequency not in interpolation range")
        diff = frequency - freq  # pylint: disable=W0631
        ldiff = lfreq - frequency
        return (ldiff * value + diff * lvalue) / (diff + ldiff)  # pylint: disable=W0631

    def linearize_freqs(self, delta_freq=None):
        """
        Rebin frequencies so that the frequency axis is linear.

        Parameters
        ----------
        delta_freq : float
            Difference between consecutive values on the new frequency axis.
            Defaults to half of smallest delta in current frequency axis.
            Compare Nyquist-Shannon sampling theorem.
        """
        if delta_freq is None:
            # Nyquist–Shannon sampling theorem
            delta_freq = _min_delt(self.freq_axis) / 2.
        nsize = int((self.freq_axis.max() - self.freq_axis.min()) /
                    delta_freq + 1)
        new = np.zeros((int(nsize), self.shape[1]), dtype=self.data.dtype)

        freqs = self.freq_axis - self.freq_axis.max()
        freqs = freqs / delta_freq

        midpoints = np.round((freqs[:-1] + freqs[1:]) / 2)
        fillto = np.concatenate(
            [midpoints - 1, np.round([freqs[-1]]) - 1]
        )
        fillfrom = np.concatenate(
            [np.round([freqs[0]]), midpoints - 1]
        )

        fillto = np.abs(fillto)
        fillfrom = np.abs(fillfrom)

        for row, from_, to_ in zip(self, fillfrom, fillto):
            new[int(from_): int(to_)] = row

        vrs = self._get_params()
        vrs.update({
            'freq_axis': np.linspace(
                self.freq_axis.max(), self.freq_axis.min(), nsize
            )
        })

        return self.__class__(new, **vrs)

    def freq_overlap(self, other):
        """
        Get frequency range present in both spectrograms. Returns (min, max)
        tuple.

        Parameters
        ----------
        other : Spectrogram
            other spectrogram with which to look for frequency overlap
        """
        lower = max(self.freq_axis[-1], other.freq_axis[-1])
        upper = min(self.freq_axis[0], other.freq_axis[0])
        if lower > upper:
            raise ValueError("No overlap.")
        return lower, upper

    def time_to_x(self, time):
        """
        Return x-coordinate in spectrogram that corresponds to the passed
        `~datetime.datetime` value.

        Parameters
        ----------
        time : `~sunpy.time.parse_time` compatible str
            `~datetime.datetime` to find the x coordinate for.
        """
        diff = time - self.start
        diff_s = SECONDS_PER_DAY * diff.days + diff.seconds
        if self.time_axis[-1] < diff_s < 0:
            raise ValueError("Out of bounds")
        for n, elem in enumerate(self.time_axis):
            if diff_s < elem:
                return n - 1
        # The last element is the searched one.
        return n

    def at_freq(self, freq):
        return self[np.nonzero(self.freq_axis == freq)[0], :]

    @staticmethod
    def _mk_format_coord(spec, fmt_coord):
        def format_coord(x, y):
            shape = list(map(int, spec.shape))

            xint, yint = int(x), int(y)
            if 0 <= xint < shape[1] and 0 <= yint < shape[0]:
                pixel = spec[yint][xint]
            else:
                pixel = ""

            return '{!s} z={!s}'.format(fmt_coord(x, y), pixel)

        return format_coord
Exemplo n.º 8
0
class CallistoSpectrogram(LinearTimeSpectrogram):
    """
    Class used for dynamic spectra coming from the Callisto network.

    Attributes
    ----------
    header : `~astropy.io.fits.Header`
        main header of the FITS file
    axes_header : `~astropy.io.fits.Header`
        header for the axes table
    swapped : bool
        flag that specifies whether originally in the file the x-axis was
        frequency
    """
    # XXX: Determine those from the data.
    SIGMA_SUM = 75
    SIGMA_DELTA_SUM = 20
    _create = ConditionalDispatch.from_existing(LinearTimeSpectrogram._create)
    create = classmethod(_create.wrapper())
    # 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 = LinearTimeSpectrogram.COPY_PROPERTIES + [
        ('header', REFERENCE), ('swapped', REFERENCE),
        ('axes_header', REFERENCE)
    ]

    # List of instruments retrieved in July 2012 from
    # http://soleil.i4ds.ch/solarradio/data/2002-20yy_Callisto/
    INSTRUMENTS = {
        'ALASKA', 'ALMATY', 'BIR', 'DARO', 'HB9SCT', 'HUMAIN', 'HURBANOVO',
        'KASI', 'KENYA', 'KRIM', 'MALAYSIA', 'MRT1', 'MRT2', 'OOTY', 'OSRA',
        'SWMC', 'TRIEST', 'UNAM'
    }

    def save(self, filepath):
        """
        Save modified spectrogram back to filepath.

        Parameters
        ----------
        filepath : str
            path to save the spectrogram to
        """
        main_header = self.get_header()
        data = fits.PrimaryHDU(self, header=main_header)
        # XXX: Update axes header.

        freq_col = fits.Column(name="frequency",
                               format="D8.3",
                               array=self.freq_axis)
        time_col = fits.Column(name="time",
                               format="D8.3",
                               array=self.time_axis)
        cols = fits.ColDefs([freq_col, time_col])
        table = fits.new_table(cols, header=self.axes_header)

        hdulist = fits.HDUList([data, table])
        hdulist.writeto(filepath)

    def get_header(self):
        """
        Returns the updated header.
        """
        header = self.header.copy()

        if self.swapped:
            header['NAXIS2'] = self.shape[1]  # pylint: disable=E1101
            header['NAXIS1'] = self.shape[0]  # pylint: disable=E1101
        else:
            header['NAXIS1'] = self.shape[1]  # pylint: disable=E1101
            header['NAXIS2'] = self.shape[0]  # pylint: disable=E1101
        return header

    @classmethod
    def read(cls, filename, **kwargs):
        """
        Reads in FITS file and return a new CallistoSpectrogram. Any unknown
        (i.e. any except filename) keyword arguments get passed to fits.open.

        Parameters
        ----------
        filename : str
            path of the file to read
        """
        fl = fits.open(filename, **kwargs)
        data = fl[0].data
        axes = fl[1]
        header = fl[0].header

        start = _parse_header_time(
            header['DATE-OBS'], header.get('TIME-OBS',
                                           header.get('TIME$_OBS')))
        end = _parse_header_time(
            header['DATE-END'], header.get('TIME-END',
                                           header.get('TIME$_END')))

        swapped = "time" not in header["CTYPE1"].lower()

        # Swap dimensions so x-axis is always time.
        if swapped:
            t_delt = header["CDELT2"]
            t_init = header["CRVAL2"] - t_delt * header["CRPIX2"]
            t_label = header["CTYPE2"]

            f_delt = header["CDELT1"]
            f_init = header["CRVAL1"] - t_delt * header["CRPIX1"]
            f_label = header["CTYPE1"]
            data = data.transpose()
        else:
            t_delt = header["CDELT1"]
            t_init = header["CRVAL1"] - t_delt * header["CRPIX1"]
            t_label = header["CTYPE1"]

            f_delt = header["CDELT2"]
            f_init = header["CRVAL2"] - t_delt * header["CRPIX2"]
            f_label = header["CTYPE2"]

        # Table may contain the axes data. If it does, the other way of doing
        # it might be very wrong.
        if axes is not None:
            try:
                # It's not my fault. Neither supports __contains__ nor .get
                tm = axes.data['time']
            except KeyError:
                tm = None
            try:
                fq = axes.data['frequency']
            except KeyError:
                fq = None

        if tm is not None:
            # Fix dimensions (whyever they are (1, x) in the first place)
            time_axis = np.squeeze(tm)
        else:
            # Otherwise, assume it's linear.
            time_axis = \
                np.linspace(0, data.shape[1] - 1) * t_delt + t_init  # pylint: disable=E1101

        if fq is not None:
            freq_axis = np.squeeze(fq)
        else:
            freq_axis = \
                np.linspace(0, data.shape[0] - 1) * f_delt + f_init  # pylint: disable=E1101

        content = header["CONTENT"]
        instruments = {header["INSTRUME"]}

        fl.close()
        return cls(data, time_axis, freq_axis, start, end, t_init, t_delt,
                   t_label, f_label, content, instruments, header, axes.header,
                   swapped)

    def __init__(self,
                 data,
                 time_axis,
                 freq_axis,
                 start,
                 end,
                 t_init=None,
                 t_delt=None,
                 t_label="Time",
                 f_label="Frequency",
                 content="",
                 instruments=None,
                 header=None,
                 axes_header=None,
                 swapped=False):
        # Because of how object creation works, there is no avoiding
        # unused arguments in this case.
        # pylint: disable=W0613

        super(CallistoSpectrogram,
              self).__init__(data, time_axis, freq_axis, start, end, t_init,
                             t_delt, t_label, f_label, content, instruments)

        self.header = header
        self.axes_header = axes_header
        self.swapped = swapped

    @classmethod
    def is_datasource_for(cls, header):
        """
        Check if class supports data from the given FITS file.

        Parameters
        ----------
        header : `~astropy.io.fits.Header`
            main header of the FITS file
        """
        return header.get('instrume', '').strip() in cls.INSTRUMENTS

    def remove_border(self):
        """
        Remove duplicate entries on the borders.
        """
        left = 0
        while self.freq_axis[left] == self.freq_axis[0]:
            left += 1
        right = self.shape[0] - 1
        while self.freq_axis[right] == self.freq_axis[-1]:
            right -= 1
        return self[left - 1:right + 2, :]

    @classmethod
    def read_many(cls, filenames, sort_by=None):
        """
        Returns a list of CallistoSpectrogram objects read from filenames.

        Parameters
        ----------
        filenames : list of str
            list of paths to read from
        sort_by : str
            optional attribute of the resulting objects to sort from, e.g.
            start to sort by starting time.
        """
        objs = list(map(cls.read, filenames))
        if sort_by is not None:
            objs.sort(key=lambda x: getattr(x, sort_by))
        return objs

    @classmethod
    def from_range(cls, instrument, start, end, **kwargs):
        """
        Automatically download data from instrument between start and end and
        join it together.

        Parameters
        ----------
        instrument : str
            instrument to retrieve the data from
        start : `~sunpy.time.parse_time` compatible
            start of the measurement
        end : `~sunpy.time.parse_time` compatible
            end of the measurement
        """
        kw = {
            'maxgap': None,
            'fill': cls.JOIN_REPEAT,
        }

        kw.update(kwargs)
        if SUNPY_LT_1:
            start = parse_time(start)
            end = parse_time(end)
        else:
            start = parse_time(start).datetime
            end = parse_time(end).datetime
        urls = query(start, end, [instrument])
        data = list(map(cls.from_url, urls))
        freq_buckets = defaultdict(list)
        for elem in data:
            freq_buckets[tuple(elem.freq_axis)].append(elem)
        try:
            return cls.combine_frequencies(
                [cls.join_many(elem, **kw) for elem in freq_buckets.values()])
        except ValueError:
            raise ValueError("No data found.")

    def _overlap(self, other):
        """
        Find frequency and time overlap of two spectrograms.
        """
        one, two = self.intersect_time([self, other])
        ovl = one.freq_overlap(two)
        return one.clip_freq(*ovl), two.clip_freq(*ovl)

    @staticmethod
    def _to_minimize(a, b):
        """
        Function to be minimized for matching to frequency channels.
        """
        def _fun(p):
            if p[0] <= 0.2 or abs(p[1]) >= a.max():
                return float("inf")
            return a - (p[0] * b + p[1])

        return _fun

    def _homogenize_params(self, other, maxdiff=1):
        """
        Return triple with a tuple of indices (in self and other,
        respectively), factors and constants at these frequencies.

        Parameters
        ----------
        other : `radiospectra.CallistoSpectrogram`
            Spectrogram to be homogenized with the current one.
        maxdiff : float
            Threshold for which frequencies are considered equal.
        """

        pairs_indices = [
            (x, y)
            for x, y, d in minimal_pairs(self.freq_axis, other.freq_axis)
            if d <= maxdiff
        ]

        pairs_data = [(self[n_one, :], other[n_two, :])
                      for n_one, n_two in pairs_indices]

        # XXX: Maybe unnecessary.
        pairs_data_gaussian = [(gaussian_filter1d(a, 15),
                                gaussian_filter1d(b, 15))
                               for a, b in pairs_data]

        # If we used integer arithmetic, we would accept more invalid
        # values.
        pairs_data_gaussian64 = np.float64(pairs_data_gaussian)
        least = [
            leastsq(self._to_minimize(a, b), [1, 0])[0]
            for a, b in pairs_data_gaussian64
        ]

        factors = [x for x, y in least]
        constants = [y for x, y in least]

        return pairs_indices, factors, constants

    def homogenize(self, other, maxdiff=1):
        """
        Return overlapping part of self and other as (self, other) tuple.
        Homogenize intensities so that the images can be used with
        combine_frequencies. Note that this works best when most of the picture
        is signal, so use :py:meth:`in_interval` to select the subset of your
        image before applying this method.

        Parameters
        ----------
        other : `radiospectra.CallistoSpectrogram`
            Spectrogram to be homogenized with the current one.
        maxdiff : float
            Threshold for which frequencies are considered equal.
        """
        one, two = self._overlap(other)
        pairs_indices, factors, constants = one._homogenize_params(
            two, maxdiff)
        # XXX: Maybe (xd.freq_axis[x] + yd.freq_axis[y]) / 2.
        pairs_freqs = [one.freq_axis[x] for x, y in pairs_indices]

        # XXX: Extrapolation does not work this way.
        # XXX: Improve.
        f1 = np.polyfit(pairs_freqs, factors, 3)
        f2 = np.polyfit(pairs_freqs, constants, 3)

        return (one, two * np.polyval(f1, two.freq_axis)[:, np.newaxis] +
                np.polyval(f2, two.freq_axis)[:, np.newaxis])

    def extend(self, minutes=15, **kwargs):
        """
        Requests subsequent files from the server.

        If minutes is negative, retrieve preceding files.
        """
        if len(self.instruments) != 1:
            raise ValueError

        instrument = next(iter(self.instruments))
        if minutes > 0:
            data = CallistoSpectrogram.from_range(
                instrument, self.end,
                self.end + datetime.timedelta(minutes=minutes))
        else:
            data = CallistoSpectrogram.from_range(
                instrument, self.start - datetime.timedelta(minutes=-minutes),
                self.start)

        data = data.clip_freq(self.freq_axis[-1], self.freq_axis[0])
        return CallistoSpectrogram.join_many([self, data], **kwargs)

    @classmethod
    def from_url(cls, url):
        """
        Returns CallistoSpectrogram read from URL.

        Parameters
        ----------
        url : str
            URL to retrieve the data from

        Returns
        -------
        newSpectrogram : `radiospectra.CallistoSpectrogram`
        """
        return cls.read(url)
Exemplo n.º 9
0
def test_types():
    f = ConditionalDispatch()
    f.add(lambda x: 2 * x, lambda x: x % 2 == 0, [int])
    with pytest.raises(TypeError):
        f(2.0)