<reponame>pauldmccarthy/fsleyes<filename>fsleyes/plotting/powerspectrumseries.py
#!/usr/bin/env python
#
# powerspectrumseries.py - Classes used by the PowerSpectrumPanel.
#
# Author: <NAME> <<EMAIL>>
#
"""This module provides :class:`.DataSeries` sub-classes which are used
by the :class:`.PowerSpectrumPanel` for plotting power spectra.

The following classes are provided:

.. autosummary::
   :nosignatures:

   PowerSpectrumSeries
   VoxelPowerSpectrumSeries
   ComplexPowerSpectrumSeries
   ImaginaryPowerSpectrumSeries
   MagnitudePowerSpectrumSeries
   PhasePowerSpectrumSeries
   MelodicPowerSpectrumSeries
"""


import logging

import numpy     as np
import numpy.fft as fft

import fsl.data.image        as fslimage
import fsl.data.melodicimage as fslmelimage
import fsleyes_props         as props
import fsleyes.colourmaps    as fslcm
import fsleyes.strings       as strings
from . import                   dataseries


log = logging.getLogger(__name__)


def calcPowerSpectrum(data):
    """Calculates a power spectrum for the given one-dimensional data array.

    :arg data:    Numpy array containing the time series data

    :returns:     If ``data`` contains real values, the magnitude of the power
                  spectrum is returned. If ``data`` contains complex values,
                  the complex power spectrum is returned.
    """

    # Fourier transform on complex data
    if np.issubdtype(data.dtype, np.complexfloating):
        data = fft.fft(data)
        data = fft.fftshift(data)

    # Fourier transform on real data - we
    # calculate and return the magnitude.
    # We also drop the first (zero-frequency)
    # term (see the rfft docs) as it is
    # kind of useless for display purposes
    else:
        data = fft.rfft(data)[1:]
        data = magnitude(data)

    return data


def calcFrequencies(nsamples, sampleTime, dtype):
    """Calculates the frequencies of the power spectrum for the given
    data.

    :arg nsamples:   Number of samples in the input time series data
    :arg sampleTime: Time between each data point
    :arg dtype:      Data type - the calculation differs depending on
                     whether the data is real or complex.
    :returns:        A ``numpy`` array containing the frequencies of the
                     power spectrum for ``data``
    """

    if np.issubdtype(dtype, np.complexfloating):
        xdata = fft.fftfreq(nsamples, sampleTime)
        xdata = fft.fftshift(xdata)
    else:
        # Drop the zero-frequency term
        # (see calcPowerSpectrum)
        xdata = fft.rfftfreq(nsamples, sampleTime)[1:]

    return xdata


def magnitude(data):
    """Returns the magnitude of the given complex data. """
    real  = data.real
    imag  = data.imag
    return np.sqrt(real ** 2 + imag ** 2)


def phase(data):
    """Returns the phase of the given complex data. """
    real = data.real
    imag = data.imag
    return np.arctan2(imag, real)


def normalise(data, dmin=None, dmax=None, nmin=-1, nmax=1):
    """Returns ``data``, rescaled to the range [nmin, nmax].

    If dmin and dmax are provided, the data is normalised with respect to them,
    rather than being normalised by the data minimum/maximum.
    """
    if dmin is None: dmin = data.min()
    if dmax is None: dmax = data.max()
    return nmin + (nmax - nmin) * (data - dmin) / (dmax - dmin)


def phaseCorrection(spectrum, freqs, p0, p1):
    """Applies phase correction to the given complex power spectrum.

    :arg spectrum: Complex-valued power spectrum
    :arg freqs:    Spectrum frequency bins
    :arg p0:       Zero order phase correction term
    :arg p1:       First order phase correction term
    :returns:      The corrected power spectrum.
    """
    exp = 1j * 2 * np.pi * (p0 / 360 + freqs * p1)
    return np.exp(exp) * spectrum


class PowerSpectrumSeries:
    """The ``PowerSpectrumSeries`` encapsulates a power spectrum data series
    from an overlay. The ``PowerSpectrumSeries`` class is a base mixin class
    for all other classes in this module.
    """


    varNorm  = props.Boolean(default=False)
    """If ``True``, the fourier-transformed data is normalised to the range
    [0, 1] before plotting.

    .. note:: The :class:`ComplexPowerSpectrumSeries` applies normalisation
              differently.
    """


    @property
    def sampleTime(self):
        """Returns the time between time series samples for the overlay
        data. """
        if isinstance(self.overlay, fslmelimage.MelodicImage):
            return self.overlay.tr
        elif isinstance(self.overlay, fslimage.Image):
            return self.overlay.pixdim[3]
        else:
            return 1


class VoxelPowerSpectrumSeries(dataseries.VoxelDataSeries,
                               PowerSpectrumSeries):
    """The ``VoxelPowerSpectrumSeries`` class encapsulates the power spectrum
    of a single voxel from a 4D :class:`.Image` overlay. The voxel is dictated
    by the :attr:`.DisplayContext.location` property.
    """


    def __init__(self, *args, **kwargs):
        """Create a ``VoxelPowerSpectrumSeries``.  All arguments are passed
        to the :meth:`VoxelDataSeries.__init__` method. A :exc:`ValueError`
        is raised if the overlay is not a 4D :class:`.Image`.
        """

        dataseries.VoxelDataSeries.__init__(self, *args, **kwargs)

        if self.overlay.ndim < 4:
            raise ValueError('Overlay is not a 4D image')


    def currentVoxelData(self, location):
        """Overrides :meth:`.VoxelDataSeries.currentVoxelData`. Retrieves
        the data at the specified location, then performs a fourier transform
        on it and returnes the result.
        """

        data = dataseries.VoxelDataSeries.currentVoxelData(self, location)
        data = calcPowerSpectrum(data)
        return data


    def getData(self):
        """Returns the ``(xdata, ydata)`` to be plotted from the current voxel
        location.
        """
        overlay = self.overlay
        ydata   = self.dataAtCurrentVoxel()
        xdata   = calcFrequencies(overlay.shape[3],
                                  self.sampleTime,
                                  overlay.dtype)

        if self.varNorm:
            ydata = normalise(ydata)

        return xdata, ydata


class ComplexPowerSpectrumSeries(VoxelPowerSpectrumSeries):
    """This class is the frequency-spectrum equivalent of the
    :class:`.ComplexTimeSeries` class - see it for more details.
    """

    plotReal      = props.Boolean(default=True)
    plotImaginary = props.Boolean(default=False)
    plotMagnitude = props.Boolean(default=False)
    plotPhase     = props.Boolean(default=False)


    zeroOrderPhaseCorrection = props.Real(default=0)
    """Apply zero order phase correction to the power spectrum of the complex
    data.
    """


    firstOrderPhaseCorrection = props.Real(default=0)
    """Apply first order phase correction to the power spectrum of the complex
    data.
    """


    def __init__(self, overlay, overlayList, displayCtx, plotCanvas):
        """Create a ``ComplexPowerSpectrumSeries``. All arguments are
        passed through to the :class:`VoxelPowerSpectrumSeries` constructor.
        """

        VoxelPowerSpectrumSeries.__init__(
            self, overlay, overlayList, displayCtx, plotCanvas)

        # Separate DataSeries for the imaginary/
        # magnitude/phase signals, returned by
        # the extraSeries method
        self.__imagps = ImaginaryPowerSpectrumSeries(
            self, overlay, overlayList, displayCtx, plotCanvas)
        self.__magps = MagnitudePowerSpectrumSeries(
            self, overlay, overlayList, displayCtx, plotCanvas)
        self.__phaseps = PhasePowerSpectrumSeries(
            self, overlay, overlayList, displayCtx, plotCanvas)

        for ps in (self.__imagps, self.__magps, self.__phaseps):
            ps.colour = fslcm.randomDarkColour()
            ps.bindProps('alpha',     self)
            ps.bindProps('lineWidth', self)
            ps.bindProps('lineStyle', self)


    def makeLabelBase(self):
        """Returns a string to be used as the label prefix for this
        ``ComplexPowerSpectrumSeries`` instance, and for the imaginary,
        magnitude, and phase child series.
        """
        return VoxelPowerSpectrumSeries.makeLabel(self)


    def makeLabel(self):
        """Returns a label to use for this data series. """
        return '{} ({})'.format(self.makeLabelBase(), strings.labels[self])


    def getData(self, component='real'):
        """If :attr:`plotReal` is true, returns the real component of the power
        spectrum of the data at the current voxel. Otherwise returns ``(None,
        None)``.

        Every time this method is called, the power spectrum is retrieved (see
        the :class:`VoxelPowerSpectrumSeries` class), phase correction is
        applied if set, andthe data is normalised, if set. A tuple containing
        the ``(xdata, ydata)`` is returned, with ``ydata`` containing the
        requested ``component`` ( ``'real'``, ``'imaginary'``,
        ``'magnitude'``, or ``'phase'``).

        This method is called by the :class:`ImaginarySpectrumPowerSeries`,
        :class:`MagnitudeSpectrumPowerSeries`, and
        :class:`PhasePowerSpectrumPowerSeries` instances that are associated
        with this data series.
        """

        if ((component == 'real')      and (not self.plotReal))      or \
           ((component == 'imaginary') and (not self.plotImaginary)) or \
           ((component == 'magnitude') and (not self.plotMagnitude)) or \
           ((component == 'phase')     and (not self.plotPhase)):
            return None, None

        # See VoxelPowerSpectrumSeries - the data
        # is already fourier-transformed
        ydata = self.dataAtCurrentVoxel()

        if ydata is None:
            return None, None

        # All of the calculations below are repeated
        # for each real/imag/mag/phase series that
        # gets plotted. But keeping the code together
        # and clean is currently more important than
        # performance, as there is not really any
        # performance hit.
        overlay = self.overlay
        xdata   = calcFrequencies(overlay.shape[3],
                                  self.sampleTime,
                                  overlay.dtype)

        if self.zeroOrderPhaseCorrection  != 0 or \
           self.firstOrderPhaseCorrection != 0:
            ydata = phaseCorrection(ydata,
                                    xdata,
                                    self.zeroOrderPhaseCorrection,
                                    self.firstOrderPhaseCorrection)

        # Normalise magnitude, real, imaginary
        # components with respect to magnitude.
        # Normalise phase independently.
        if self.varNorm:
            mag = magnitude(ydata)
            mr  = mag.min(), mag.max()
            if   component == 'phase':     ydata = normalise(phase(ydata))
            elif component == 'magnitude': ydata = normalise(mag)
            elif component == 'real':      ydata = normalise(ydata.real, *mr)
            elif component == 'imaginary': ydata = normalise(ydata.imag, *mr)

        elif component == 'real':      ydata = ydata.real
        elif component == 'imaginary': ydata = ydata.imag
        elif component == 'magnitude': ydata = magnitude(ydata)
        elif component == 'phase':     ydata = phase(ydata)

        return xdata, ydata


    def extraSeries(self):
        """Returns a list of additional series to be plotted, based
        on the values of the :attr:`plotImaginary`, :attr:`plotMagnitude`
        and :attr:`plotPhase` properties.
        """

        extras = []
        if self.plotImaginary: extras.append(self.__imagps)
        if self.plotMagnitude: extras.append(self.__magps)
        if self.plotPhase:     extras.append(self.__phaseps)
        return extras


class ImaginaryPowerSpectrumSeries(dataseries.DataSeries):
    """An ``ImaginaryPowerSpectrumSeries`` represents the power spectrum of the
    imaginary component of a complex-valued image.
    ``ImaginaryPowerSpectrumSeries`` instances are created by
    :class:`ComplexPowerSpectrumSeries` instances.
    """


    def __init__(self, parent, *args, **kwargs):
        """Create an ``ImaginaryPowerSpectrumSeries``.

        :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this
                     ``ImaginaryPowerSpectrumSeries``.

        All other arguments are passed through to the :class:`DataSeries`
        constructor.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)
        self.__parent = parent


    def makeLabel(self):
        """Returns a string representation of this
        ``ImaginaryPowerSpectrumSeries`` instance.
        """
        return '{} ({})'.format(self.__parent.makeLabelBase(),
                                strings.labels[self])


    def getData(self):
        """Returns the imaginary component of the power spectrum. """
        return self.__parent.getData('imaginary')


class MagnitudePowerSpectrumSeries(dataseries.DataSeries):
    """An ``MagnitudePowerSpectrumSeries`` represents the magnitude of a
    complex-valued image. ``MagnitudePowerSpectrumSeries`` instances are
    created by :class:`ComplexPowerSpectrumSeries` instances.
    """


    def __init__(self, parent, *args, **kwargs):
        """Create an ``ImaginaryPowerSpectrumSeries``.

        :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this
                     ``ImaginaryPowerSpectrumSeries``.

        All other arguments are passed through to the :class:`DataSeries`
        constructor.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)
        self.__parent = parent


    def makeLabel(self):
        """Returns a string representation of this
        ``MagnitudePowerSpectrumSeries`` instance.
        """
        return '{} ({})'.format(self.__parent.makeLabelBase(),
                                strings.labels[self])


    def getData(self):
        """Returns the magnitude of the complex power spectrum. """
        return self.__parent.getData('magnitude')


class PhasePowerSpectrumSeries(dataseries.DataSeries):
    """An ``PhasePowerSpectrumSeries`` represents the phase of a complex-valued
    image. ``PhasePowerSpectrumSeries`` instances are created by
    :class:`ComplexPowerSpectrumSeries` instances.
    """


    def __init__(self, parent, *args, **kwargs):
        """Create an ``ImaginaryPowerSpectrumSeries``.

        :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this
                     ``ImaginaryPowerSpectrumSeries``.

        All other arguments are passed through to the :class:`DataSeries`
        constructor.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)
        self.__parent = parent


    def makeLabel(self):
        """Returns a string representation of this ``PhasePowerSpectrumSeries``
        instance.
        """
        return '{} ({})'.format(self.__parent.makeLabelBase(),
                                strings.labels[self])


    def getData(self):
        """Returns the phase of the complex power spectrum. """
        return self.__parent.getData('phase')


class MelodicPowerSpectrumSeries(dataseries.DataSeries,
                                 PowerSpectrumSeries):
    """The ``MelodicPowerSpectrumSeries`` class encapsulates the power spectrum
    of the time course for a single component of a :class:`.MelodicImage`. The
    component is dictated by the :attr:`.NiftiOpts.volume` property.
    """


    def __init__(self, *args, **kwargs):
        """Create a ``MelodicPowerSpectrumSeries``. All arguments are passed
        through to the :meth:`PowerSpectrumSeries.__init__` method.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)

        if not isinstance(self.overlay, fslmelimage.MelodicImage):
            raise ValueError('Overlay is not a MelodicImage')

        self.varNorm = False
        self.disableProperty('varNorm')


    def makeLabel(self):
        """Returns a label that can be used for this
        ``MelodicPowerSpectrumSeries``.
        """
        display   = self.displayCtx.getDisplay(self.overlay)
        opts      = display.opts
        component = opts.volume

        return '{} [component {}]'.format(display.name, component + 1)


    def getData(self):
        """Returns the power spectrum for the current component of the
        :class:`.MelodicImage`, as defined by the :attr:`.NiftiOpts.volume`
        property.
        """

        overlay   = self.overlay
        opts      = self.displayCtx.getOpts(overlay)
        component = opts.volume

        ydata     = overlay.getComponentPowerSpectrum(component)
        nsamples  = len(overlay.getComponentTimeSeries(component))
        xdata     = calcFrequencies(nsamples, overlay.tr, np.float32)

        return xdata, ydata


class MeshPowerSpectrumSeries(dataseries.DataSeries,
                              PowerSpectrumSeries):
    """A ``MeshPowerSpectrumSeries`` object encapsulates the power spectrum for
    the data from a :class:`.Mesh` overlay which has some time series
    vertex data associated with it. See the :attr:`.MeshOpts.vertexData`
    property.
    """


    def __init__(self, *args, **kwargs):
        """Create a ``MeshPowerSpectrumSeries`` instance. All arguments are
        passed through to  :meth:`.DataSeries.__init__`.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)


    def makeLabel(self):
        """Returns a label to use for this ``MeshPowerSpectrumSeries`` on the
        legend.
        """

        display = self.displayCtx.getDisplay(self.overlay)

        if self.__haveData():
            opts = display.opts
            vidx = opts.getVertex()
            return '{} [{}]'.format(display.name, vidx)

        else:
            return display.name


    def __haveData(self):
        """Returns ``True`` if there is currently time series data to show
        for this ``MeshPowerSpectrumSeries``, ``False`` otherwise.
        """
        opts = self.displayCtx.getOpts(self.overlay)
        vidx = opts.getVertex()
        vd   = opts.getVertexData()

        return vidx is not None and vd is not None and vd.shape[1] > 1


    def getData(self):
        """Returns the power spectrum of the data at the current location for
        the :class:`.Mesh`, or ``None, None`` if there is no data.
        """

        if not self.__haveData():
            return None, None

        opts = self.displayCtx.getOpts(self.overlay)
        vidx = opts.getVertex()

        if vidx is None:
            return None, None

        vd    = opts.getVertexData()
        data  = vd[vidx, :]
        xdata = calcFrequencies(  len(data), self.sampleTime, data.dtype)
        ydata = calcPowerSpectrum(data)

        return xdata, ydata

#!/usr/bin/env python
#
# powerspectrumseries.py - Classes used by the PowerSpectrumPanel.
#
# Author: <NAME> <<EMAIL>>
#
"""This module provides :class:`.DataSeries` sub-classes which are used
by the :class:`.PowerSpectrumPanel` for plotting power spectra.

The following classes are provided:

.. autosummary::
   :nosignatures:

   PowerSpectrumSeries
   VoxelPowerSpectrumSeries
   ComplexPowerSpectrumSeries
   ImaginaryPowerSpectrumSeries
   MagnitudePowerSpectrumSeries
   PhasePowerSpectrumSeries
   MelodicPowerSpectrumSeries
"""


import logging

import numpy     as np
import numpy.fft as fft

import fsl.data.image        as fslimage
import fsl.data.melodicimage as fslmelimage
import fsleyes_props         as props
import fsleyes.colourmaps    as fslcm
import fsleyes.strings       as strings
from . import                   dataseries


log = logging.getLogger(__name__)


def calcPowerSpectrum(data):
    """Calculates a power spectrum for the given one-dimensional data array.

    :arg data:    Numpy array containing the time series data

    :returns:     If ``data`` contains real values, the magnitude of the power
                  spectrum is returned. If ``data`` contains complex values,
                  the complex power spectrum is returned.
    """

    # Fourier transform on complex data
    if np.issubdtype(data.dtype, np.complexfloating):
        data = fft.fft(data)
        data = fft.fftshift(data)

    # Fourier transform on real data - we
    # calculate and return the magnitude.
    # We also drop the first (zero-frequency)
    # term (see the rfft docs) as it is
    # kind of useless for display purposes
    else:
        data = fft.rfft(data)[1:]
        data = magnitude(data)

    return data


def calcFrequencies(nsamples, sampleTime, dtype):
    """Calculates the frequencies of the power spectrum for the given
    data.

    :arg nsamples:   Number of samples in the input time series data
    :arg sampleTime: Time between each data point
    :arg dtype:      Data type - the calculation differs depending on
                     whether the data is real or complex.
    :returns:        A ``numpy`` array containing the frequencies of the
                     power spectrum for ``data``
    """

    if np.issubdtype(dtype, np.complexfloating):
        xdata = fft.fftfreq(nsamples, sampleTime)
        xdata = fft.fftshift(xdata)
    else:
        # Drop the zero-frequency term
        # (see calcPowerSpectrum)
        xdata = fft.rfftfreq(nsamples, sampleTime)[1:]

    return xdata


def magnitude(data):
    """Returns the magnitude of the given complex data. """
    real  = data.real
    imag  = data.imag
    return np.sqrt(real ** 2 + imag ** 2)


def phase(data):
    """Returns the phase of the given complex data. """
    real = data.real
    imag = data.imag
    return np.arctan2(imag, real)


def normalise(data, dmin=None, dmax=None, nmin=-1, nmax=1):
    """Returns ``data``, rescaled to the range [nmin, nmax].

    If dmin and dmax are provided, the data is normalised with respect to them,
    rather than being normalised by the data minimum/maximum.
    """
    if dmin is None: dmin = data.min()
    if dmax is None: dmax = data.max()
    return nmin + (nmax - nmin) * (data - dmin) / (dmax - dmin)


def phaseCorrection(spectrum, freqs, p0, p1):
    """Applies phase correction to the given complex power spectrum.

    :arg spectrum: Complex-valued power spectrum
    :arg freqs:    Spectrum frequency bins
    :arg p0:       Zero order phase correction term
    :arg p1:       First order phase correction term
    :returns:      The corrected power spectrum.
    """
    exp = 1j * 2 * np.pi * (p0 / 360 + freqs * p1)
    return np.exp(exp) * spectrum


class PowerSpectrumSeries:
    """The ``PowerSpectrumSeries`` encapsulates a power spectrum data series
    from an overlay. The ``PowerSpectrumSeries`` class is a base mixin class
    for all other classes in this module.
    """


    varNorm  = props.Boolean(default=False)
    """If ``True``, the fourier-transformed data is normalised to the range
    [0, 1] before plotting.

    .. note:: The :class:`ComplexPowerSpectrumSeries` applies normalisation
              differently.
    """


    @property
    def sampleTime(self):
        """Returns the time between time series samples for the overlay
        data. """
        if isinstance(self.overlay, fslmelimage.MelodicImage):
            return self.overlay.tr
        elif isinstance(self.overlay, fslimage.Image):
            return self.overlay.pixdim[3]
        else:
            return 1


class VoxelPowerSpectrumSeries(dataseries.VoxelDataSeries,
                               PowerSpectrumSeries):
    """The ``VoxelPowerSpectrumSeries`` class encapsulates the power spectrum
    of a single voxel from a 4D :class:`.Image` overlay. The voxel is dictated
    by the :attr:`.DisplayContext.location` property.
    """


    def __init__(self, *args, **kwargs):
        """Create a ``VoxelPowerSpectrumSeries``.  All arguments are passed
        to the :meth:`VoxelDataSeries.__init__` method. A :exc:`ValueError`
        is raised if the overlay is not a 4D :class:`.Image`.
        """

        dataseries.VoxelDataSeries.__init__(self, *args, **kwargs)

        if self.overlay.ndim < 4:
            raise ValueError('Overlay is not a 4D image')


    def currentVoxelData(self, location):
        """Overrides :meth:`.VoxelDataSeries.currentVoxelData`. Retrieves
        the data at the specified location, then performs a fourier transform
        on it and returnes the result.
        """

        data = dataseries.VoxelDataSeries.currentVoxelData(self, location)
        data = calcPowerSpectrum(data)
        return data


    def getData(self):
        """Returns the ``(xdata, ydata)`` to be plotted from the current voxel
        location.
        """
        overlay = self.overlay
        ydata   = self.dataAtCurrentVoxel()
        xdata   = calcFrequencies(overlay.shape[3],
                                  self.sampleTime,
                                  overlay.dtype)

        if self.varNorm:
            ydata = normalise(ydata)

        return xdata, ydata


class ComplexPowerSpectrumSeries(VoxelPowerSpectrumSeries):
    """This class is the frequency-spectrum equivalent of the
    :class:`.ComplexTimeSeries` class - see it for more details.
    """

    plotReal      = props.Boolean(default=True)
    plotImaginary = props.Boolean(default=False)
    plotMagnitude = props.Boolean(default=False)
    plotPhase     = props.Boolean(default=False)


    zeroOrderPhaseCorrection = props.Real(default=0)
    """Apply zero order phase correction to the power spectrum of the complex
    data.
    """


    firstOrderPhaseCorrection = props.Real(default=0)
    """Apply first order phase correction to the power spectrum of the complex
    data.
    """


    def __init__(self, overlay, overlayList, displayCtx, plotCanvas):
        """Create a ``ComplexPowerSpectrumSeries``. All arguments are
        passed through to the :class:`VoxelPowerSpectrumSeries` constructor.
        """

        VoxelPowerSpectrumSeries.__init__(
            self, overlay, overlayList, displayCtx, plotCanvas)

        # Separate DataSeries for the imaginary/
        # magnitude/phase signals, returned by
        # the extraSeries method
        self.__imagps = ImaginaryPowerSpectrumSeries(
            self, overlay, overlayList, displayCtx, plotCanvas)
        self.__magps = MagnitudePowerSpectrumSeries(
            self, overlay, overlayList, displayCtx, plotCanvas)
        self.__phaseps = PhasePowerSpectrumSeries(
            self, overlay, overlayList, displayCtx, plotCanvas)

        for ps in (self.__imagps, self.__magps, self.__phaseps):
            ps.colour = fslcm.randomDarkColour()
            ps.bindProps('alpha',     self)
            ps.bindProps('lineWidth', self)
            ps.bindProps('lineStyle', self)


    def makeLabelBase(self):
        """Returns a string to be used as the label prefix for this
        ``ComplexPowerSpectrumSeries`` instance, and for the imaginary,
        magnitude, and phase child series.
        """
        return VoxelPowerSpectrumSeries.makeLabel(self)


    def makeLabel(self):
        """Returns a label to use for this data series. """
        return '{} ({})'.format(self.makeLabelBase(), strings.labels[self])


    def getData(self, component='real'):
        """If :attr:`plotReal` is true, returns the real component of the power
        spectrum of the data at the current voxel. Otherwise returns ``(None,
        None)``.

        Every time this method is called, the power spectrum is retrieved (see
        the :class:`VoxelPowerSpectrumSeries` class), phase correction is
        applied if set, andthe data is normalised, if set. A tuple containing
        the ``(xdata, ydata)`` is returned, with ``ydata`` containing the
        requested ``component`` ( ``'real'``, ``'imaginary'``,
        ``'magnitude'``, or ``'phase'``).

        This method is called by the :class:`ImaginarySpectrumPowerSeries`,
        :class:`MagnitudeSpectrumPowerSeries`, and
        :class:`PhasePowerSpectrumPowerSeries` instances that are associated
        with this data series.
        """

        if ((component == 'real')      and (not self.plotReal))      or \
           ((component == 'imaginary') and (not self.plotImaginary)) or \
           ((component == 'magnitude') and (not self.plotMagnitude)) or \
           ((component == 'phase')     and (not self.plotPhase)):
            return None, None

        # See VoxelPowerSpectrumSeries - the data
        # is already fourier-transformed
        ydata = self.dataAtCurrentVoxel()

        if ydata is None:
            return None, None

        # All of the calculations below are repeated
        # for each real/imag/mag/phase series that
        # gets plotted. But keeping the code together
        # and clean is currently more important than
        # performance, as there is not really any
        # performance hit.
        overlay = self.overlay
        xdata   = calcFrequencies(overlay.shape[3],
                                  self.sampleTime,
                                  overlay.dtype)

        if self.zeroOrderPhaseCorrection  != 0 or \
           self.firstOrderPhaseCorrection != 0:
            ydata = phaseCorrection(ydata,
                                    xdata,
                                    self.zeroOrderPhaseCorrection,
                                    self.firstOrderPhaseCorrection)

        # Normalise magnitude, real, imaginary
        # components with respect to magnitude.
        # Normalise phase independently.
        if self.varNorm:
            mag = magnitude(ydata)
            mr  = mag.min(), mag.max()
            if   component == 'phase':     ydata = normalise(phase(ydata))
            elif component == 'magnitude': ydata = normalise(mag)
            elif component == 'real':      ydata = normalise(ydata.real, *mr)
            elif component == 'imaginary': ydata = normalise(ydata.imag, *mr)

        elif component == 'real':      ydata = ydata.real
        elif component == 'imaginary': ydata = ydata.imag
        elif component == 'magnitude': ydata = magnitude(ydata)
        elif component == 'phase':     ydata = phase(ydata)

        return xdata, ydata


    def extraSeries(self):
        """Returns a list of additional series to be plotted, based
        on the values of the :attr:`plotImaginary`, :attr:`plotMagnitude`
        and :attr:`plotPhase` properties.
        """

        extras = []
        if self.plotImaginary: extras.append(self.__imagps)
        if self.plotMagnitude: extras.append(self.__magps)
        if self.plotPhase:     extras.append(self.__phaseps)
        return extras


class ImaginaryPowerSpectrumSeries(dataseries.DataSeries):
    """An ``ImaginaryPowerSpectrumSeries`` represents the power spectrum of the
    imaginary component of a complex-valued image.
    ``ImaginaryPowerSpectrumSeries`` instances are created by
    :class:`ComplexPowerSpectrumSeries` instances.
    """


    def __init__(self, parent, *args, **kwargs):
        """Create an ``ImaginaryPowerSpectrumSeries``.

        :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this
                     ``ImaginaryPowerSpectrumSeries``.

        All other arguments are passed through to the :class:`DataSeries`
        constructor.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)
        self.__parent = parent


    def makeLabel(self):
        """Returns a string representation of this
        ``ImaginaryPowerSpectrumSeries`` instance.
        """
        return '{} ({})'.format(self.__parent.makeLabelBase(),
                                strings.labels[self])


    def getData(self):
        """Returns the imaginary component of the power spectrum. """
        return self.__parent.getData('imaginary')


class MagnitudePowerSpectrumSeries(dataseries.DataSeries):
    """An ``MagnitudePowerSpectrumSeries`` represents the magnitude of a
    complex-valued image. ``MagnitudePowerSpectrumSeries`` instances are
    created by :class:`ComplexPowerSpectrumSeries` instances.
    """


    def __init__(self, parent, *args, **kwargs):
        """Create an ``ImaginaryPowerSpectrumSeries``.

        :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this
                     ``ImaginaryPowerSpectrumSeries``.

        All other arguments are passed through to the :class:`DataSeries`
        constructor.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)
        self.__parent = parent


    def makeLabel(self):
        """Returns a string representation of this
        ``MagnitudePowerSpectrumSeries`` instance.
        """
        return '{} ({})'.format(self.__parent.makeLabelBase(),
                                strings.labels[self])


    def getData(self):
        """Returns the magnitude of the complex power spectrum. """
        return self.__parent.getData('magnitude')


class PhasePowerSpectrumSeries(dataseries.DataSeries):
    """An ``PhasePowerSpectrumSeries`` represents the phase of a complex-valued
    image. ``PhasePowerSpectrumSeries`` instances are created by
    :class:`ComplexPowerSpectrumSeries` instances.
    """


    def __init__(self, parent, *args, **kwargs):
        """Create an ``ImaginaryPowerSpectrumSeries``.

        :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this
                     ``ImaginaryPowerSpectrumSeries``.

        All other arguments are passed through to the :class:`DataSeries`
        constructor.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)
        self.__parent = parent


    def makeLabel(self):
        """Returns a string representation of this ``PhasePowerSpectrumSeries``
        instance.
        """
        return '{} ({})'.format(self.__parent.makeLabelBase(),
                                strings.labels[self])


    def getData(self):
        """Returns the phase of the complex power spectrum. """
        return self.__parent.getData('phase')


class MelodicPowerSpectrumSeries(dataseries.DataSeries,
                                 PowerSpectrumSeries):
    """The ``MelodicPowerSpectrumSeries`` class encapsulates the power spectrum
    of the time course for a single component of a :class:`.MelodicImage`. The
    component is dictated by the :attr:`.NiftiOpts.volume` property.
    """


    def __init__(self, *args, **kwargs):
        """Create a ``MelodicPowerSpectrumSeries``. All arguments are passed
        through to the :meth:`PowerSpectrumSeries.__init__` method.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)

        if not isinstance(self.overlay, fslmelimage.MelodicImage):
            raise ValueError('Overlay is not a MelodicImage')

        self.varNorm = False
        self.disableProperty('varNorm')


    def makeLabel(self):
        """Returns a label that can be used for this
        ``MelodicPowerSpectrumSeries``.
        """
        display   = self.displayCtx.getDisplay(self.overlay)
        opts      = display.opts
        component = opts.volume

        return '{} [component {}]'.format(display.name, component + 1)


    def getData(self):
        """Returns the power spectrum for the current component of the
        :class:`.MelodicImage`, as defined by the :attr:`.NiftiOpts.volume`
        property.
        """

        overlay   = self.overlay
        opts      = self.displayCtx.getOpts(overlay)
        component = opts.volume

        ydata     = overlay.getComponentPowerSpectrum(component)
        nsamples  = len(overlay.getComponentTimeSeries(component))
        xdata     = calcFrequencies(nsamples, overlay.tr, np.float32)

        return xdata, ydata


class MeshPowerSpectrumSeries(dataseries.DataSeries,
                              PowerSpectrumSeries):
    """A ``MeshPowerSpectrumSeries`` object encapsulates the power spectrum for
    the data from a :class:`.Mesh` overlay which has some time series
    vertex data associated with it. See the :attr:`.MeshOpts.vertexData`
    property.
    """


    def __init__(self, *args, **kwargs):
        """Create a ``MeshPowerSpectrumSeries`` instance. All arguments are
        passed through to  :meth:`.DataSeries.__init__`.
        """
        dataseries.DataSeries.__init__(self, *args, **kwargs)


    def makeLabel(self):
        """Returns a label to use for this ``MeshPowerSpectrumSeries`` on the
        legend.
        """

        display = self.displayCtx.getDisplay(self.overlay)

        if self.__haveData():
            opts = display.opts
            vidx = opts.getVertex()
            return '{} [{}]'.format(display.name, vidx)

        else:
            return display.name


    def __haveData(self):
        """Returns ``True`` if there is currently time series data to show
        for this ``MeshPowerSpectrumSeries``, ``False`` otherwise.
        """
        opts = self.displayCtx.getOpts(self.overlay)
        vidx = opts.getVertex()
        vd   = opts.getVertexData()

        return vidx is not None and vd is not None and vd.shape[1] > 1


    def getData(self):
        """Returns the power spectrum of the data at the current location for
        the :class:`.Mesh`, or ``None, None`` if there is no data.
        """

        if not self.__haveData():
            return None, None

        opts = self.displayCtx.getOpts(self.overlay)
        vidx = opts.getVertex()

        if vidx is None:
            return None, None

        vd    = opts.getVertexData()
        data  = vd[vidx, :]
        xdata = calcFrequencies(  len(data), self.sampleTime, data.dtype)
        ydata = calcPowerSpectrum(data)

        return xdata, ydata
#!/usr/bin/env python # # powerspectrumseries.py - Classes used by the PowerSpectrumPanel. # # Author: <NAME> <<EMAIL>> # This module provides :class:`.DataSeries` sub-classes which are used by the :class:`.PowerSpectrumPanel` for plotting power spectra.  The following classes are provided:  .. autosummary::    :nosignatures:     PowerSpectrumSeries    VoxelPowerSpectrumSeries    ComplexPowerSpectrumSeries    ImaginaryPowerSpectrumSeries    MagnitudePowerSpectrumSeries    PhasePowerSpectrumSeries    MelodicPowerSpectrumSeries Calculates a power spectrum for the given one-dimensional data array.      :arg data:    Numpy array containing the time series data      :returns:     If ``data`` contains real values, the magnitude of the power                   spectrum is returned. If ``data`` contains complex values,                   the complex power spectrum is returned. # Fourier transform on complex data # Fourier transform on real data - we # calculate and return the magnitude. # We also drop the first (zero-frequency) # term (see the rfft docs) as it is # kind of useless for display purposes Calculates the frequencies of the power spectrum for the given     data.      :arg nsamples:   Number of samples in the input time series data     :arg sampleTime: Time between each data point     :arg dtype:      Data type - the calculation differs depending on                      whether the data is real or complex.     :returns:        A ``numpy`` array containing the frequencies of the                      power spectrum for ``data`` # Drop the zero-frequency term # (see calcPowerSpectrum) Returns the magnitude of the given complex data. Returns the phase of the given complex data. Returns ``data``, rescaled to the range [nmin, nmax].      If dmin and dmax are provided, the data is normalised with respect to them,     rather than being normalised by the data minimum/maximum. Applies phase correction to the given complex power spectrum.      :arg spectrum: Complex-valued power spectrum     :arg freqs:    Spectrum frequency bins     :arg p0:       Zero order phase correction term     :arg p1:       First order phase correction term     :returns:      The corrected power spectrum. The ``PowerSpectrumSeries`` encapsulates a power spectrum data series     from an overlay. The ``PowerSpectrumSeries`` class is a base mixin class     for all other classes in this module. If ``True``, the fourier-transformed data is normalised to the range     [0, 1] before plotting.      .. note:: The :class:`ComplexPowerSpectrumSeries` applies normalisation               differently. Returns the time between time series samples for the overlay         data. The ``VoxelPowerSpectrumSeries`` class encapsulates the power spectrum     of a single voxel from a 4D :class:`.Image` overlay. The voxel is dictated     by the :attr:`.DisplayContext.location` property. Create a ``VoxelPowerSpectrumSeries``.  All arguments are passed         to the :meth:`VoxelDataSeries.__init__` method. A :exc:`ValueError`         is raised if the overlay is not a 4D :class:`.Image`. Overrides :meth:`.VoxelDataSeries.currentVoxelData`. Retrieves         the data at the specified location, then performs a fourier transform         on it and returnes the result. Returns the ``(xdata, ydata)`` to be plotted from the current voxel         location. This class is the frequency-spectrum equivalent of the     :class:`.ComplexTimeSeries` class - see it for more details. Apply zero order phase correction to the power spectrum of the complex     data. Apply first order phase correction to the power spectrum of the complex     data. Create a ``ComplexPowerSpectrumSeries``. All arguments are         passed through to the :class:`VoxelPowerSpectrumSeries` constructor. # Separate DataSeries for the imaginary/ # magnitude/phase signals, returned by # the extraSeries method Returns a string to be used as the label prefix for this         ``ComplexPowerSpectrumSeries`` instance, and for the imaginary,         magnitude, and phase child series. Returns a label to use for this data series. If :attr:`plotReal` is true, returns the real component of the power         spectrum of the data at the current voxel. Otherwise returns ``(None,         None)``.          Every time this method is called, the power spectrum is retrieved (see         the :class:`VoxelPowerSpectrumSeries` class), phase correction is         applied if set, andthe data is normalised, if set. A tuple containing         the ``(xdata, ydata)`` is returned, with ``ydata`` containing the         requested ``component`` ( ``'real'``, ``'imaginary'``,         ``'magnitude'``, or ``'phase'``).          This method is called by the :class:`ImaginarySpectrumPowerSeries`,         :class:`MagnitudeSpectrumPowerSeries`, and         :class:`PhasePowerSpectrumPowerSeries` instances that are associated         with this data series. # See VoxelPowerSpectrumSeries - the data # is already fourier-transformed # All of the calculations below are repeated # for each real/imag/mag/phase series that # gets plotted. But keeping the code together # and clean is currently more important than # performance, as there is not really any # performance hit. # Normalise magnitude, real, imaginary # components with respect to magnitude. # Normalise phase independently. Returns a list of additional series to be plotted, based         on the values of the :attr:`plotImaginary`, :attr:`plotMagnitude`         and :attr:`plotPhase` properties. An ``ImaginaryPowerSpectrumSeries`` represents the power spectrum of the     imaginary component of a complex-valued image.     ``ImaginaryPowerSpectrumSeries`` instances are created by     :class:`ComplexPowerSpectrumSeries` instances. Create an ``ImaginaryPowerSpectrumSeries``.          :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this                      ``ImaginaryPowerSpectrumSeries``.          All other arguments are passed through to the :class:`DataSeries`         constructor. Returns a string representation of this         ``ImaginaryPowerSpectrumSeries`` instance. Returns the imaginary component of the power spectrum. An ``MagnitudePowerSpectrumSeries`` represents the magnitude of a     complex-valued image. ``MagnitudePowerSpectrumSeries`` instances are     created by :class:`ComplexPowerSpectrumSeries` instances. Create an ``ImaginaryPowerSpectrumSeries``.          :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this                      ``ImaginaryPowerSpectrumSeries``.          All other arguments are passed through to the :class:`DataSeries`         constructor. Returns a string representation of this         ``MagnitudePowerSpectrumSeries`` instance. Returns the magnitude of the complex power spectrum. An ``PhasePowerSpectrumSeries`` represents the phase of a complex-valued     image. ``PhasePowerSpectrumSeries`` instances are created by     :class:`ComplexPowerSpectrumSeries` instances. Create an ``ImaginaryPowerSpectrumSeries``.          :arg parent: The :class:`ComplexPowerSpectrumSeries` which owns this                      ``ImaginaryPowerSpectrumSeries``.          All other arguments are passed through to the :class:`DataSeries`         constructor. Returns a string representation of this ``PhasePowerSpectrumSeries``         instance. Returns the phase of the complex power spectrum. The ``MelodicPowerSpectrumSeries`` class encapsulates the power spectrum     of the time course for a single component of a :class:`.MelodicImage`. The     component is dictated by the :attr:`.NiftiOpts.volume` property. Create a ``MelodicPowerSpectrumSeries``. All arguments are passed         through to the :meth:`PowerSpectrumSeries.__init__` method. Returns a label that can be used for this         ``MelodicPowerSpectrumSeries``. Returns the power spectrum for the current component of the         :class:`.MelodicImage`, as defined by the :attr:`.NiftiOpts.volume`         property. A ``MeshPowerSpectrumSeries`` object encapsulates the power spectrum for     the data from a :class:`.Mesh` overlay which has some time series     vertex data associated with it. See the :attr:`.MeshOpts.vertexData`     property. Create a ``MeshPowerSpectrumSeries`` instance. All arguments are         passed through to  :meth:`.DataSeries.__init__`. Returns a label to use for this ``MeshPowerSpectrumSeries`` on the         legend. Returns ``True`` if there is currently time series data to show         for this ``MeshPowerSpectrumSeries``, ``False`` otherwise. Returns the power spectrum of the data at the current location for         the :class:`.Mesh`, or ``None, None`` if there is no data.