ALLEN INSTITUTE

# Source code for bmtk.simulator.filternet.lgnmodel.movie

    import numpy as np

    from .utilities import convert_tmin_tmax_framerate_to_trange
    from bmtk.simulator.filternet.io_tools import io

    class Movie(object):
        def __init__(self, data, row_range=None, col_range=None, labels=('time', 'y', 'x'),
                     units=('second', 'pixel', 'pixel'), frame_rate=None, t_range=None, padding=False):
            self.data = data
            self.labels = labels
            self.units = units
            self.padding=padding
            assert(units[0] == 'second')

            if t_range is None:
                self.frame_rate = float(frame_rate)
                self.t_range = np.arange(data.shape[0])*(1./self.frame_rate)
            else:
                self.t_range = np.array(t_range)
                self.frame_rate = 1./np.mean(np.diff(t_range))

            if row_range is None:
                self.row_range = np.arange(data.shape[1])
            else:
                self.row_range = np.array(row_range)
            if col_range is None:
                self.col_range = np.arange(data.shape[2])
            else:
                self.col_range = np.array(col_range)

        @property
        def ranges(self):
            return self.t_range, self.row_range, self.col_range

        def imshow_summary(self, ax=None, show=True, xlabel=None):
            import matplotlib.pyplot as plt

            if ax is None:
                _, ax = plt.subplots(1,1)

            t_vals = self.t_range.copy()
            y_vals = self.data.mean(axis=2).mean(axis=1)
            ax.plot(t_vals, y_vals)
            ax.set_ylim(y_vals.min()-np.abs(y_vals.min())*.05, y_vals.max()+np.abs(y_vals.max())*.05)

            if not xlabel is None:
                ax.set_xlabel(xlabel)

            ax.set_ylabel('Average frame intensity')

            if show:
                plt.show()

            return ax, (t_vals, y_vals)

        def imshow(self, t, show=True, vmin=-1, vmax=1, cmap=None):
            import matplotlib.pyplot as plt

            cmap = cmap or plt.cm.gray
            ti = int(t*self.frame_rate)
            data = self.data[ti, :, :]
            plt.imshow(data, vmin=vmin, vmax=vmax, cmap=cmap)
            plt.colorbar()
            if show:
                plt.show()

        def play(self, interval=None, skip=None, playback_speed=1.0, cmap=None, vmin=-1.0, vmax=1.0, save_as=None,
                 show=True):
            """

            :param interval:
            :param skip: skip every Nth frame, useful for very large data sets
            :param playback_speed:
            :param cmap: color map, default to gray
            :param save_as: If set to a valid (mp4 file) path animation will be saved. Otherwise it won't
            :param show: display the animation, default True
            """
            import matplotlib.pyplot as plt
            import matplotlib.animation as animation

            cmap = cmap or plt.cm.gray

            data = self.data
            times = self.t_range

            if skip is not None:
                data = data[::skip, :, :]
                times = self.t_range[::skip]

            if interval is None:
                interval = np.mean(np.diff(times))*1000

            frames = []
            fig = plt.figure()
            ax = fig.add_subplot(111)
            for i in range(len(times)):
                # TODO: Find a way to display a running clock as a x-label or title
                frame_img = ax.imshow(data[i, :, :], cmap=cmap, animated=True, vmin=vmin, vmax=vmax,
                                      extent=[self.col_range[0], self.col_range[-1], self.row_range[0], self.row_range[-1]])
                title = ax.text(0.99, 0.000, "{:05.2f} s".format(times[i]),
                                size=plt.rcParams["axes.titlesize"],
                                horizontalalignment='right',
                                verticalalignment='bottom',
                                transform=ax.transAxes, bbox={'boxstyle': 'round', 'facecolor': 'white'})

                frames.append([frame_img, title])

            ani = animation.ArtistAnimation(fig, frames, interval=interval*playback_speed, blit=True, repeat_delay=1000)

            if save_as is not None:
                ani.save(save_as)

            if show:
                plt.show()

        def __add__(self, other):
            assert self.labels == other.labels
            assert self.units == other.units
            assert self.frame_rate == other.frame_rate
            np.testing.assert_almost_equal(self.col_range, other.col_range)
            np.testing.assert_almost_equal(self.row_range, other.row_range)

            new_data = np.empty((len(self.t_range)+len(other.t_range)-1, len(self.row_range), len(self.col_range)))
            new_data[:len(self.t_range), :, :] = self.data[:, :, :]
            new_data[len(self.t_range):, :, :] = other.data[1:, :, :]

            return Movie(new_data, row_range=self.row_range.copy(), col_range=self.col_range.copy(), labels=self.labels,
                         units=self.units, frame_rate=self.frame_rate)

        def __getitem__(self, *args):
            return self.data.__getitem__(*args)

        @staticmethod
        def normalize_matrix(m_data, domain=None):
            """Attempts to take a numpy matrix "movie" and normalize the contrast values to range [-1.0, 1.0].

            :param m_data: A numpy matrix TxRxC containing an original "movie" data
            :param domain: If A tuple of values [min_contrast, max_contrast], will use those values to normalize. Otherwise
                will attempt to determine original contrast range from the m_data.
            :return: m_data that has been normalized
            """
            if isinstance(domain, (list, tuple, np.ndarray)):
                # Allow users to specify the contrast range of the original movie [c_min, c_max] that will be
                # converted to range [-1.0, +1.0]
                if len(domain) != 2:
                    raise ValueError(
                        'Unable to normalize movie; please specify to minimum and maximum contrast for original movie with '
                        'list [min_contrast, max_contrast]'
                    )
                c_min, c_max = domain[0], domain[1]
                if c_min >= c_max:
                    raise ValueError(
                        'Unable to normalize movie; original movie range must be defined with min_contrast < max_contrast'
                    )
                elif m_data.min() < c_min or m_data.max() > c_max:
                    raise ValueError(
                        'Unable to normalize movie; trying to convert movie from contrast values ' +
                        '[{}, {}] --> [-1.0, 1.0], '.format(c_min, c_max) +
                        'but movie contains values outside of original domain.'
                    )
            else:
                # Use the existing movie to try to best guess the domain the original movie is already in; in most
                # cases it will be [0.0, 1.0] or [0, 256]
                m_min, m_max = m_data.min(), m_data.max()
                if m_min >= 0.0 and m_max <= 1.0:
                    c_min, c_max = 0.0, 1.0
                elif m_min >= -1.0 and m_max <= 1.0:
                    # All values fall between [-1.0, 1.0], likely movie is likely already normalized.
                    io.log_debug('Input movie is already normalized between contrast values [-1.0, 1.0].')
                    return m_data
                elif m_min >= 0 and m_max < 256:
                    c_min, c_max = 0.0, 255.0
                else:
                    c_min, c_max = m_min, m_max

            if c_min == c_max:
                raise ValueError(
                    'Unable to normalize movie; please specify a range for the contrast min and max of the input movie.'
                )

            io.log_debug('Normalizing input movie from contrast range [{}, {}] --> [-1.0, 1.0].')
            return (m_data - c_min)*2.0/(c_max - c_min) - 1.0

    # TODO: Instead of using subclasses, convert the following movie type to a Factory class/function that returns a Movie
    #   object

    class FullFieldMovie(Movie):
        def __init__(self, f, rows, cols, frame_rate=24.0):
            if np.isscalar(rows):
                rows = np.arange(0, rows)

            if np.isscalar(cols):
                cols = np.arange(0, cols)

            self.row_range = rows
            self.col_range = cols
            self.frame_size = (len(self.row_range), len(self.col_range))
            self._frame_rate = frame_rate
            self.f = f

        @property
        def frame_rate(self):
            return self._frame_rate

        @property
        def data(self):
            return self

        def __getitem__(self, *args):
            t_inds, x_inds, y_inds = args[0]
            assert (len(x_inds) == len(y_inds)) and (len(y_inds) == len(t_inds))

            # Convert frame indices to times:
            t_vals = (1./self.frame_rate)*t_inds

            # Evaluate and return:
            return self.f(t_vals)

        def full(self, t_min=0.0, t_max=None):
            # Compute t_range
            t_range = convert_tmin_tmax_framerate_to_trange(t_min, t_max, self.frame_rate)

            nt = len(t_range)
            nr = len(self.row_range)
            nc = len(self.col_range)
            a, b, c = np.meshgrid(range(nt), range(nr), range(nc))
            af, bf, cf = map(lambda x: x.flatten(), [a, b, c])
            data = np.empty((nt, nr, nc))
            data[af, bf, cf] = self.f(t_range[af])

            return Movie(data, row_range=self.row_range, col_range=self.col_range, labels=('time', 'y', 'x'),
                         units=('second', 'pixel', 'pixel'), frame_rate=self.frame_rate)

        def create_movie(self, t_max, t_min=0.0):
            return self.full(t_min=t_min, t_max=t_max)

    class FullFieldFlashMovie(FullFieldMovie):
        def __init__(self, row_range, col_range, t_on, t_off, max_intensity=1, frame_rate=24):
            assert(t_on < t_off)

            def f(t):
                return np.piecewise(t, *zip(*[(t < t_on, 0), (np.logical_and(t_on <= t, t < t_off), max_intensity),
                                              (t_off <= t, 0)]))

            super(FullFieldFlashMovie, self).__init__(f, row_range, col_range, frame_rate=frame_rate)

    class GratingMovie(Movie):
        def __init__(self, row_size, col_size, frame_rate=1000.0):
            self.row_size = row_size  # in degrees
            self.col_size = col_size  # in degrees
            self.frame_rate = float(frame_rate)  # in Hz

        def create_movie(self, t_min=0, t_max=1, gray_screen_dur=0, cpd=0.05, temporal_f=4, theta=45,
                         phase=0., contrast=1.0, degrees_per_pixel=None, row_size_new=None, col_size_new=None):
            """Create the grating movie with the desired parameters

            :param t_min: start time in seconds
            :param t_max: end time in seconds
            :param gray_screen_dur: Duration of gray screen before grating stimulus starts
            :param cpd: cycles per degree
            :param temporal_f: in Hz
            :param theta: orientation angle, in degrees
            :param phase: temporal phase, in degrees
            :param contrast: the maximum constrast, must be between 0 and 1.0
            :param degrees_per_pixel: pixel pitch of the movie in degrees (default (if None) is 1.0/(cpd*10))
            :param row_size_new: Use to truncate screen, by default leaves original row dimension
            :param col_size_new: Use to truncate screen, by default leaves original col dimension
            :return: Movie object of grating with desired parameters
            """
            assert contrast <= 1, "Contrast must be <= 1"
            assert contrast > 0, "Contrast must be > 0"

            if degrees_per_pixel is None:  # default behavior when not specified
                degrees_per_pixel = 1.0 / (float(cpd) * 10)  # To make sure no aliasing occurs
            self.row_range = np.linspace(0, self.row_size, int(self.row_size/degrees_per_pixel), endpoint=True)
            self.col_range = np.linspace(0, self.col_size, int(self.col_size/degrees_per_pixel), endpoint=True)
            numberFramesNeeded = int(round(self.frame_rate * (t_max - gray_screen_dur))) + 1
            time_range = np.linspace(0, t_max - gray_screen_dur, numberFramesNeeded, endpoint=True)

            # Creates a drifting grating panel
            tt, yy, xx = np.meshgrid(time_range, self.row_range, self.col_range, indexing='ij')
            theta_rad = np.pi*(180 - theta) / 180.
            phase_rad = np.pi*(180 - phase) / 180.
            xy = xx * np.cos(theta_rad) + yy * np.sin(theta_rad)
            data = contrast*np.sin(2*np.pi*(cpd * xy + temporal_f *tt) + phase_rad)

            # truncate
            if row_size_new != None:
                self.row_range = np.linspace(0, row_size_new, data.shape[1], endpoint=True)

            if col_size_new != None:
                self.col_range = np.linspace(0, col_size_new, data.shape[2], endpoint=True)

            if gray_screen_dur > 0:
                # just adding one or two seconds to gray screen so flash never "happens"
                m_gray = FullFieldFlashMovie(self.row_range, self.col_range, gray_screen_dur + 1, gray_screen_dur + 2,
                                             frame_rate=self.frame_rate).full(t_max=gray_screen_dur)
                mov = m_gray + Movie(data, row_range=self.row_range, col_range=self.col_range, labels=('time', 'y', 'x'),
                                     units=('second', 'pixel', 'pixel'), frame_rate=self.frame_rate)
            else:
                mov = Movie(data, row_range=self.row_range, col_range=self.col_range, labels=('time', 'y', 'x'),
                            units=('second', 'pixel', 'pixel'), frame_rate=self.frame_rate)

            return mov

    class LoomingMovie(Movie):
        def __init__(self, row_size, col_size, frame_rate=1000.):
            self.row_size = row_size  # in degrees
            self.col_size = col_size  # in degrees
            self.frame_rate = float(frame_rate)  # in Hz

        def create_movie(self, t_looming=1, gray_screen_dur=0.5):
            """Create the looming movie with the desired parameters

            :param t_looming: duration of time looming
            :param gray_screen_dur:
            """
            degrees_per_pixel = 1.0  # To make sure no aliasing occurs
            self.row_range = np.linspace(0, self.row_size, int(self.row_size/degrees_per_pixel), endpoint=True)
            self.col_range = np.linspace(0, self.col_size, int(self.col_size/degrees_per_pixel), endpoint=True)
            loomingFramesNeeded = int(round(self.frame_rate * t_looming))
            grayScreenFrames = int(round(self.frame_rate * gray_screen_dur))
            time_range = np.linspace(0, t_looming, loomingFramesNeeded, endpoint=True)

            data = np.zeros((grayScreenFrames + loomingFramesNeeded, self.row_size, self.col_size))
            # mgrid is a mesh creation helper
            xx, yy = np.mgrid[:self.row_size, :self.col_size]
            # yy, xx = np.meshgrid(time_range, self.row_range, self.col_range, indexing='ij')
            # circles contains the squared distance to the (middle, middle) point
            # we are just using the circle equation learnt at school
            circle = (xx - self.row_size / 2) ** 2 + (yy - self.col_size / 2) ** 2
            radius = 25  # Cirlces grow to this size

            for i in range(loomingFramesNeeded):
                # Create filled circle
                data[grayScreenFrames + i, :, :] = (circle <= (i/(4./3)) % (radius**2))
            data *= -1   # Adjusting so have dark looming and not bright looming disk

            mov = Movie(data, row_range=self.row_range, col_range=self.col_range, labels=('time', 'y', 'x'),
                        units=('second', 'pixel', 'pixel'), frame_rate=self.frame_rate)

            return mov