<commit_before>from playground.base import logger
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.pyplot import cm
from scipy.interpolate import interp1d
from .core import spk_time_to_scv, firing_pos_from_scv
from ..base import SPKTAG



def info_bits(Fr, P):
    Fr[Fr==0] = 1e-25
    MFr = sum(P.ravel()*Fr.ravel())
    return sum(P.ravel()*(Fr.ravel()/MFr)*np.log2(Fr.ravel()/MFr))


def info_sparcity(Fr, P):
    Fr[Fr==0] = 1e-25
    MFr = sum(P.ravel()*Fr.ravel())
    return sum(P.ravel()*Fr.ravel()**2/MFr**2)


class place_field(object):
    """getting the place fields from subspaces"""
    def __init__(self, logfile=None, session_id=0, v_cutoff=5, maze_range=[[-100,100], [-100,100]], bin_size=4, sync=True):
        super(place_field, self).__init__()
        # default mode:
        # 1: ts, pos, dt 
        # 2: v_smoothed, v_cutoff
        # 3: maze_range, bin_size 
        if logfile is None:
            pass
             
        # logfile mode
        else:
            self.logfile = logfile
            self.log = logger(self.logfile, sync=sync)
            self.ts, self.pos = self.log.to_trajectory(session_id)
            self.pos[:,1] = -self.pos[:,1]
            

            # self.v_smoothed, self.v = self.log.get_speed(self.ts, self.pos, smooth_window=60, std=15)
            # self.v_cutoff = v_cutoff
            # self.get_still_idx() 

            # self.get_speed(smooth_window=60, std=15, v_cutoff=5) 
            # self.maze_range = maze_range
            # self.occupation_map(bin_size)
            self.initialize(bin_size=bin_size, v_cutoff=v_cutoff, maze_range=maze_range)


        ### place fields parameters ###
        self.n_fields = 0


    def interp_pos(self, t, pos, N=1):
        '''
        convert irregularly sampled pos into regularly sampled pos
        N is the dilution sampling factor. N=2 means half of the resampled pos
        example:
        >>> new_fs = 200.
        >>> pc.ts, pc.pos = pc.interp_pos(ts, pos, N=fs/new_fs)
        '''
        dt = np.mean(np.diff(t))
        x, y = interp1d(t, pos[:,0], fill_value="extrapolate"), interp1d(t, pos[:,1], fill_value="extrapolate")
        new_t = np.arange(0.0, dt*len(t), dt*N)
        new_pos = np.hstack((x(new_t).reshape(-1,1), y(new_t).reshape(-1,1)))
        return new_t, new_pos 


    def align_with_recording(self, recording_start_time, recording_end_time, replay_offset=0):
        '''
        ts before alignment   |--------------------|
        behavior  start:      |
        behavior    end:                           |
        recording start:         |------------
        recording   end:          ------------|
        replay_offset  :             |
        ts after alignment           |------------| 
        '''
        self.ts += replay_offset   # 0 if the ephys is not offset by replaying through neural signal generator
        self.pos = self.pos[np.logical_and(self.ts>recording_start_time, self.ts<recording_end_time)]
        self.ts  =  self.ts[np.logical_and(self.ts>recording_start_time, self.ts<recording_end_time)]
        self.t_start = self.ts[0]
        self.t_end   = self.ts[-1]


    def initialize(self, bin_size, v_cutoff, maze_range=None):
        self.dt = self.ts[1] - self.ts[0]
        self.get_maze_range(maze_range)
        self.get_speed(smooth_window=60, std=15, v_cutoff=v_cutoff) 
        self.occupation_map(bin_size)


    def get_maze_range(self, maze_range=None):
        if maze_range is None:
            self.maze_range = np.vstack((self.pos.min(axis=0), self.pos.max(axis=0))).T
            self._maze_original = self.maze_range[:,0] # the left, down corner location
        else:
            self.maze_range = np.array(maze_range)
            self._maze_original = self.maze_range[:,0] # the left, down corner location

    @property
    def maze_original(self):
        return self._maze_original


    def get_speed(self, smooth_window=59, std=6, v_cutoff=5):
        '''
        self.ts, self.pos is required
        '''
        v = np.linalg.norm(np.diff(self.pos, axis=0), axis=1)/np.diff(self.ts)
        w = signal.gaussian(smooth_window, std) # window size 59 frame (roughly 1 sec), std = 6 frame
        w /= sum(w)
        v_smoothed = np.convolve(v, w, mode='same')

        self.v = np.hstack((0.01, v))
        self.v_smoothed = np.hstack((0.01, v_smoothed))
        self.v_cutoff   = v_cutoff
        self.v_still_idx = np.where(self.v_smoothed < self.v_cutoff)[0]
        '''
        # check speed:
        f, ax = plt.subplots(1,1, figsize=(18,8))
        offset=20000
        plot(ts[offset:1000+offset], v[offset:1000+offset])
        plot(ts[offset:1000+offset], v_smoothed[offset:1000+offset])
        ax.axhline(5, c='m', ls='-.')
        '''
        # return v_smoothed, v

    def plot_speed(self, start, stop, thres=5):
        fig, ax = plt.subplots(1,1, figsize=(18,8))
        period = np.logical_and(self.ts>start, self.ts<stop)
        plt.plot(self.ts[period], self.v[period], alpha=.7)
        plt.plot(self.ts[period], self.v_smoothed[period], lw=3)
        ax.axhline(thres, c='m', ls='-.')
        sns.despine()
        return fig
        

    def occupation_map(self, bin_size=4, time_cutoff=None):
        '''
        f, ax = plt.subplots(1,2,figsize=(20,9))
        ax[0].plot(self.pos[:,0], self.pos[:,1])
        ax[0].plot(self.pos[0,0], self.pos[0,1], 'ro')
        ax[0].plot(self.pos[-1,0], self.pos[-1,1], 'ko')
        ax[0].pcolormesh(self.X, self.Y, self.O, cmap=cm.hot_r)
        sns.heatmap(self.O[::-1]*self.dt, annot=False, cbar=False, ax=ax[1])
        '''
        # if maze_range != 'auto':
        #     self.maze_range = maze_range
        self.maze_size = np.array([self.maze_range[0][1]-self.maze_range[0][0], self.maze_range[1][1]-self.maze_range[1][0]])
        self.bin_size  = bin_size
        self.nbins = self.maze_size/bin_size
        self.nbins = self.nbins.astype(int)
        # occupation, self.x_edges, self.y_edges = np.histogram2d(x=self.pos[1:,0], y=self.pos[1:,1], 
        #                                                         bins=self.nbins, range=self.maze_range)
        idx = np.where(self.v_smoothed >= self.v_cutoff)[0]
        if time_cutoff is not None:
            idx = np.delete(idx, np.where(self.ts[idx]>time_cutoff)[0])
        occupation, self.x_edges, self.y_edges = np.histogram2d(x=self.pos[idx,0], y=self.pos[idx,1], 
                                                                bins=self.nbins, range=self.maze_range)
        self.X, self.Y = np.meshgrid(self.x_edges, self.y_edges)
        self.O = occupation.T.astype(int)  # Let each row list bins with common y range.
        self.P = self.O/float(self.O.sum()) # occupation prabability

        #### parameter used to calculate the fields
        self.kernlen = 18
        self.kernstd = 2.5


    def plot_occupation_map(self):
        f, ax = plt.subplots(1,2,figsize=(20,9))
        ax[0].plot(self.pos[:,0], self.pos[:,1])
        ax[0].plot(self.pos[0,0], self.pos[0,1], 'ro')
        ax[0].plot(self.pos[-1,0], self.pos[-1,1], 'go')
        ax[0].pcolormesh(self.X, self.Y, self.O, cmap=cm.hot)
        # sns.heatmap(self.O[::-1]*self.dt, annot=False, cbar=False, ax=ax[1])
        ax[1].pcolormesh(self.X, self.Y, self.O, cmap=cm.hot)
        plt.show()

    @property
    def map_binned_size(self):
        return np.array(np.diff(self.maze_range)/self.bin_size, dtype=np.int).ravel()[::-1]
    
    @staticmethod
    def gkern(kernlen=21, std=2):
        """Returns a 2D Gaussian kernel array."""
        gkern1d = signal.gaussian(kernlen, std=std).reshape(kernlen, 1)
        gkern2d = np.outer(gkern1d, gkern1d)
        gkern2d /= gkern2d.sum()
        return gkern2d


    def _get_field(self, spk_times):
        spk_ts = np.searchsorted(self.ts, spk_times) - 1
        idx = np.setdiff1d(spk_ts, self.v_still_idx)
        self.firing_ts  = self.ts[spk_ts] #[:,1]
        self.firing_pos = self.pos[idx]        
        self.firing_map, x_edges, y_edges = np.histogram2d(x=self.firing_pos[:,0], y=self.firing_pos[:,1], 
                                                           bins=self.nbins, range=self.maze_range)
        self.firing_map = self.firing_map.T
        np.seterr(divide='ignore', invalid='ignore')
        self.FR = self.firing_map/self.O/self.dt
        # self.FR = np.nan_to_num(self.FR)
        self.FR[np.isnan(self.FR)] = 0
        self.FR[np.isinf(self.FR)] = 0
        self.FR_smoothed = signal.convolve2d(self.FR, self.gkern(self.kernlen, self.kernstd), boundary='symm', mode='same')
        return self.FR_smoothed


    def firing_map_from_scv(self, scv, section=[0,1]):
        '''
        firing heat map constructed from spike count vector (scv) and position
        '''
        # assert(scv.shape[1]==self.pos.shape[0])
        n_neurons, total_bin = scv.shape
        valid_bin = np.array(np.array(section)*total_bin, dtype=np.int)
        firing_map_smoothed = np.zeros((n_neurons, *self.map_binned_size))
        for neuron_id in range(n_neurons):
            firing_pos = firing_pos_from_scv(scv, self.pos, neuron_id, valid_bin)<commit_msg>pc.low_speed_idx when pc.initialize()<commit_after>from playground.base import logger
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.pyplot import cm
from scipy.interpolate import interp1d
from .core import spk_time_to_scv, firing_pos_from_scv
from ..base import SPKTAG



def info_bits(Fr, P):
    Fr[Fr==0] = 1e-25
    MFr = sum(P.ravel()*Fr.ravel())
    return sum(P.ravel()*(Fr.ravel()/MFr)*np.log2(Fr.ravel()/MFr))


def info_sparcity(Fr, P):
    Fr[Fr==0] = 1e-25
    MFr = sum(P.ravel()*Fr.ravel())
    return sum(P.ravel()*Fr.ravel()**2/MFr**2)


class place_field(object):
    """getting the place fields from subspaces"""
    def __init__(self, logfile=None, session_id=0, v_cutoff=5, maze_range=[[-100,100], [-100,100]], bin_size=4, sync=True):
        super(place_field, self).__init__()
        # default mode:
        # 1: ts, pos, dt 
        # 2: v_smoothed, v_cutoff
        # 3: maze_range, bin_size 
        if logfile is None:
            pass
             
        # logfile mode
        else:
            self.logfile = logfile
            self.log = logger(self.logfile, sync=sync)
            self.ts, self.pos = self.log.to_trajectory(session_id)
            self.pos[:,1] = -self.pos[:,1]
            

            # self.v_smoothed, self.v = self.log.get_speed(self.ts, self.pos, smooth_window=60, std=15)
            # self.v_cutoff = v_cutoff
            # self.get_still_idx() 

            # self.get_speed(smooth_window=60, std=15, v_cutoff=5) 
            # self.maze_range = maze_range
            # self.occupation_map(bin_size)
            self.initialize(bin_size=bin_size, v_cutoff=v_cutoff, maze_range=maze_range)


        ### place fields parameters ###
        self.n_fields = 0


    def interp_pos(self, t, pos, N=1):
        '''
        convert irregularly sampled pos into regularly sampled pos
        N is the dilution sampling factor. N=2 means half of the resampled pos
        example:
        >>> new_fs = 200.
        >>> pc.ts, pc.pos = pc.interp_pos(ts, pos, N=fs/new_fs)
        '''
        dt = np.mean(np.diff(t))
        x, y = interp1d(t, pos[:,0], fill_value="extrapolate"), interp1d(t, pos[:,1], fill_value="extrapolate")
        new_t = np.arange(0.0, dt*len(t), dt*N)
        new_pos = np.hstack((x(new_t).reshape(-1,1), y(new_t).reshape(-1,1)))
        return new_t, new_pos 


    def align_with_recording(self, recording_start_time, recording_end_time, replay_offset=0):
        '''
        ts before alignment   |--------------------|
        behavior  start:      |
        behavior    end:                           |
        recording start:         |------------
        recording   end:          ------------|
        replay_offset  :             |
        ts after alignment           |------------| 
        '''
        self.ts += replay_offset   # 0 if the ephys is not offset by replaying through neural signal generator
        self.pos = self.pos[np.logical_and(self.ts>recording_start_time, self.ts<recording_end_time)]
        self.ts  =  self.ts[np.logical_and(self.ts>recording_start_time, self.ts<recording_end_time)]
        self.t_start = self.ts[0]
        self.t_end   = self.ts[-1]


    def initialize(self, bin_size, v_cutoff, maze_range=None):
        self.dt = self.ts[1] - self.ts[0]
        self.get_maze_range(maze_range)
        self.get_speed(smooth_window=60, std=15, v_cutoff=v_cutoff) 
        self.occupation_map(bin_size)


    def get_maze_range(self, maze_range=None):
        if maze_range is None:
            self.maze_range = np.vstack((self.pos.min(axis=0), self.pos.max(axis=0))).T
            self._maze_original = self.maze_range[:,0] # the left, down corner location
        else:
            self.maze_range = np.array(maze_range)
            self._maze_original = self.maze_range[:,0] # the left, down corner location

    @property
    def maze_original(self):
        return self._maze_original


    def get_speed(self, smooth_window=59, std=6, v_cutoff=5):
        '''
        self.ts, self.pos is required
        '''
        v = np.linalg.norm(np.diff(self.pos, axis=0), axis=1)/np.diff(self.ts)
        w = signal.gaussian(smooth_window, std) # window size 59 frame (roughly 1 sec), std = 6 frame
        w /= sum(w)
        v_smoothed = np.convolve(v, w, mode='same')

        self.v = np.hstack((0.01, v))
        self.v_smoothed = np.hstack((0.01, v_smoothed))
        self.v_cutoff   = v_cutoff
        self.low_speed_idx = np.where(self.v_smoothed < self.v_cutoff)[0]
        '''
        # check speed:
        f, ax = plt.subplots(1,1, figsize=(18,8))
        offset=20000
        plot(ts[offset:1000+offset], v[offset:1000+offset])
        plot(ts[offset:1000+offset], v_smoothed[offset:1000+offset])
        ax.axhline(5, c='m', ls='-.')
        '''
        # return v_smoothed, v

    def plot_speed(self, start, stop, thres=5):
        fig, ax = plt.subplots(1,1, figsize=(18,8))
        period = np.logical_and(self.ts>start, self.ts<stop)
        plt.plot(self.ts[period], self.v[period], alpha=.7)
        plt.plot(self.ts[period], self.v_smoothed[period], lw=3)
        ax.axhline(thres, c='m', ls='-.')
        sns.despine()
        return fig
        

    def occupation_map(self, bin_size=4, time_cutoff=None):
        '''
        f, ax = plt.subplots(1,2,figsize=(20,9))
        ax[0].plot(self.pos[:,0], self.pos[:,1])
        ax[0].plot(self.pos[0,0], self.pos[0,1], 'ro')
        ax[0].plot(self.pos[-1,0], self.pos[-1,1], 'ko')
        ax[0].pcolormesh(self.X, self.Y, self.O, cmap=cm.hot_r)
        sns.heatmap(self.O[::-1]*self.dt, annot=False, cbar=False, ax=ax[1])
        '''
        # if maze_range != 'auto':
        #     self.maze_range = maze_range
        self.maze_size = np.array([self.maze_range[0][1]-self.maze_range[0][0], self.maze_range[1][1]-self.maze_range[1][0]])
        self.bin_size  = bin_size
        self.nbins = self.maze_size/bin_size
        self.nbins = self.nbins.astype(int)
        # occupation, self.x_edges, self.y_edges = np.histogram2d(x=self.pos[1:,0], y=self.pos[1:,1], 
        #                                                         bins=self.nbins, range=self.maze_range)
        idx = np.where(self.v_smoothed >= self.v_cutoff)[0]
        if time_cutoff is not None:
            idx = np.delete(idx, np.where(self.ts[idx]>time_cutoff)[0])
        occupation, self.x_edges, self.y_edges = np.histogram2d(x=self.pos[idx,0], y=self.pos[idx,1], 
                                                                bins=self.nbins, range=self.maze_range)
        self.X, self.Y = np.meshgrid(self.x_edges, self.y_edges)
        self.O = occupation.T.astype(int)  # Let each row list bins with common y range.
        self.P = self.O/float(self.O.sum()) # occupation prabability

        #### parameter used to calculate the fields
        self.kernlen = 18
        self.kernstd = 2.5


    def plot_occupation_map(self):
        f, ax = plt.subplots(1,2,figsize=(20,9))
        ax[0].plot(self.pos[:,0], self.pos[:,1])
        ax[0].plot(self.pos[0,0], self.pos[0,1], 'ro')
        ax[0].plot(self.pos[-1,0], self.pos[-1,1], 'go')
        ax[0].pcolormesh(self.X, self.Y, self.O, cmap=cm.hot)
        # sns.heatmap(self.O[::-1]*self.dt, annot=False, cbar=False, ax=ax[1])
        ax[1].pcolormesh(self.X, self.Y, self.O, cmap=cm.hot)
        plt.show()

    @property
    def map_binned_size(self):
        return np.array(np.diff(self.maze_range)/self.bin_size, dtype=np.int).ravel()[::-1]
    
    @staticmethod
    def gkern(kernlen=21, std=2):
        """Returns a 2D Gaussian kernel array."""
        gkern1d = signal.gaussian(kernlen, std=std).reshape(kernlen, 1)
        gkern2d = np.outer(gkern1d, gkern1d)
        gkern2d /= gkern2d.sum()
        return gkern2d


    def _get_field(self, spk_times):
        spk_ts = np.searchsorted(self.ts, spk_times) - 1
        idx = np.setdiff1d(spk_ts, self.low_speed_idx)
        self.firing_ts  = self.ts[spk_ts] #[:,1]
        self.firing_pos = self.pos[idx]        
        self.firing_map, x_edges, y_edges = np.histogram2d(x=self.firing_pos[:,0], y=self.firing_pos[:,1], 
                                                           bins=self.nbins, range=self.maze_range)
        self.firing_map = self.firing_map.T
        np.seterr(divide='ignore', invalid='ignore')
        self.FR = self.firing_map/self.O/self.dt
        # self.FR = np.nan_to_num(self.FR)
        self.FR[np.isnan(self.FR)] = 0
        self.FR[np.isinf(self.FR)] = 0
        self.FR_smoothed = signal.convolve2d(self.FR, self.gkern(self.kernlen, self.kernstd), boundary='symm', mode='same')
        return self.FR_smoothed


    def firing_map_from_scv(self, scv, section=[0,1]):
        '''
        firing heat map constructed from spike count vector (scv) and position
        '''
        # assert(scv.shape[1]==self.pos.shape[0])
        n_neurons, total_bin = scv.shape
        valid_bin = np.array(np.array(section)*total_bin, dtype=np.int)
        firing_map_smoothed = np.zeros((n_neurons, *self.map_binned_size))
        for neuron_id in range(n_neurons):
            firing_pos = firing_pos_from_scv(scv, self.pos, neuron_id, valid_bin)
{'fraction_non_alphanumeric': 0.15032040838492453, 'fraction_numerical': 0.023351797545345932, 'mean_word_length': 2.3031390134529146, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '224705', 'n_tokens_mistral': 6936, 'n_tokens_neox': 6664, 'n_words': 1151}