preprocess.py

"""

Preprocess and clean individual fish trajectories, and extract group polarisation order parameter.


Author: Arshed Nabeel
Copyright: TEE-Lab, Center for Ecological Sciences, Indian Institute of Science, Bengaluru

"""

import glob
import os
import warnings
from typing import Tuple, List

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import MaxNLocator
from scipy.sparse.csgraph import connected_components
from scipy.spatial.distance import pdist, squareform

from tqdm.notebook import tqdm, trange


# --- MAIN TRAJECTORY PROCESSOR CLASS ---

class TrajectoryProcessor(object):

    def __init__(self, source: str, format: str = 'trex', n_fish: int = None, dt: float = None):

        self.source_name = source.rstrip(os.sep)

        if format == 'trex':
            self.timestamps, self.positions = self.load_trex_trajectories(source)
        elif format == 'golden-shiner':
            self.timestamps, self.positions = self.load_golden_shiner_trajectories(source, n_fish, dt)
        elif format == 'hemi':
            self.timestamps, self.positions = self.load_hemigrammus_trajectories(source, n_fish, dt)
        elif format == 'fasttrack':
            self.timestamps, self.positions = self.load_fasttrack_trajectories(source, dt)
        elif format == 'idtracker':
            self.timestamps, self.positions = self.load_idtracker_trajectories(source)
        else:
            raise TypeError("Unsupported format. Currently supported: 'trex', 'fasttrack', 'idtracker'.")

        self.n_timepoints, self.n_fish, _ = self.positions.shape
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            self.velocities = self.compute_velocities()
            self.polarisation = self.compute_polarisation()
            self.speeds = np.linalg.norm(self.velocities, axis=2)

        self.metadata = []

    def summary(self,
                x_range: Tuple = None,
                y_range: Tuple = None,
                vx_range: Tuple = None,
                vy_range: Tuple = None,
                speed_range: Tuple = None) -> None:
        """
        Visualize the position, velocity, and polarisation distributions of the data, to look for anomalies.

        Args:
            x_range: Range for plotting x-coordinates of positions
            y_range: Range for plotting y-coordinates of positions
            vx_range: Range for plotting x-coordinates of velocities
            vy_range: Range for plotting y-coordinates of velocities
            speed_range: Range for plotting speed histogram
        """

        fig, ax = plt.subplots(2, 3, figsize=(12, 8))
        ax_pos = ax[0, 0]
        ax_vel, ax_speed = ax[0, 1], ax[0, 2]
        ax_pol, ax_pol_mag = ax[1, 1], ax[1, 2]
        ax_n_fish = ax[1, 0]

        missing_fish = (np.isnan(self.positions[..., 0]) | np.isnan(self.positions[..., 1]))
        n_detected_fish = np.sum(~missing_fish, axis=1)

        if x_range is None:
            x_range = np.nanmin(self.positions[:, :, 0]), np.nanmax(self.positions[:, :, 0])

        if y_range is None:
            y_range = np.nanmin(self.positions[:, :, 1]), np.nanmax(self.positions[:, :, 1])

        if vx_range is None:
            vx_range = np.nanmin(self.velocities[:, :, 0]), np.nanmax(self.velocities[:, :, 0])

        if vy_range is None:
            vy_range = np.nanmin(self.velocities[:, :, 1]), np.nanmax(self.velocities[:, :, 1])

        if speed_range is None:
            speed_range = (np.nanmin(self.speeds), np.nanmax(self.speeds))

        pol_mag = np.linalg.norm(self.polarisation, axis=1)
        ax_pos.hist2d(self.positions[:, :, 0].flatten(), self.positions[:, :, 1].flatten(),
                      range=(x_range, y_range), bins=51, density=True, cmap='inferno')
        ax_vel.hist2d(self.velocities[:, :, 0].flatten(), self.velocities[:, :, 1].flatten(),
                      range=(vx_range, vy_range), bins=51, density=True, cmap='inferno')
        ax_speed.hist(self.speeds.flatten(), range=speed_range, bins=100, density=True)
        ax_pol.hist2d(self.polarisation[:, 0].flatten(), self.polarisation[:, 1].flatten(),
                      range=((-1, 1), (-1, 1)), bins=51, density=True, cmap='inferno')
        ax_pol_mag.hist(pol_mag, range=(0, 1), bins=100, density=True)
        ax_n_fish.hist(n_detected_fish, bins=np.arange(.5, self.n_fish + 1.5), density=True)

        ax_pos.set(title='Individual Position', xlabel='$x$', ylabel='$y$')
        ax_vel.set(title='Individual Velocity', xlabel='$v_x$', ylabel='$v_y$')
        ax_speed.set(title='Individual Speed', xlabel='$|v|$', ylabel='$P(|v|)$')
        ax_pol.set(title='Group Polarisation', xlabel='$m_x$', ylabel='$m_y$')
        ax_pol_mag.set(title='Group Polarisation (Magnitude)', xlabel='$|m|$', ylabel='$P(|m|)$')
        ax_n_fish.set(title='Number of detected fish per frame', xlabel='$N_{det}$', ylabel='$P(N_{det})$')
        ax_n_fish.xaxis.set_major_locator(MaxNLocator(integer=True))

        plt.tight_layout()
        plt.show()

    def preprocess(self, preprocessors: List = None):

        for processor in preprocessors:
            processor(self)

        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            self.velocities = self.compute_velocities()
            self.polarisation = self.compute_polarisation()
            self.speeds = np.linalg.norm(self.velocities, axis=2)

        n_valid_frames = np.sum(np.all(~np.isnan(self.polarisation), axis=1))
        msg = f'{n_valid_frames}/{self.n_timepoints} valid frames ({100 * n_valid_frames / self.n_timepoints:.2f}%) ' \
              f'after preprocessing.\n'
        print(msg)
        self.metadata.append(msg)

    def group_split_diagnostics(self, threshold: float = 6):
        """ Diagnostics function to examine group-splitting dynamics and identify the correct distance threshold
        for clustering.
        Args:
            threshold: Distance threshold for clustering. Two fish will be considered neighbours if they are within
            these distance.
        """

        self.igs = IdentifyGroupSplit(threshold)
        n_clusters = self.igs(self, preview=True)
        split_idx = np.where(n_clusters > 1)[0]
        no_split_idx = np.where(n_clusters == 1)[0]

        def _show_individuals(ax, coords, xbounds, ybounds):
            # circles = PatchCollection([Circle((x_, y_), r) for (x_, y_) in coords],
            #                           alpha=0.2)
            ax.plot(coords[..., 0], coords[..., 1], '.')
            ax.set(xlim=xbounds, ylim=ybounds)
            # ax.add_collection(circles)

        xmin, ymin = np.nanmin(self.positions, axis=(0, 1))
        xmax, ymax = np.nanmax(self.positions, axis=(0, 1))

        fig, ax_ = plt.subplots(4, 4, figsize=(12, 12), sharex='all', sharey='all',)
        idx_to_plot = np.random.choice(no_split_idx, 16)
        for idx, ax in zip(idx_to_plot, ax_.flatten()):
            _show_individuals(ax, self.positions[idx, ...],
                              xbounds=(xmin, xmax), ybounds=(ymin, ymax))
        fig.suptitle('Example frames with no splitting')
        # plt.tight_layout()

        fig, ax_ = plt.subplots(4, 4, figsize=(12, 12), sharex='all', sharey='all',)
        idx_to_plot = np.random.choice(split_idx, 16)
        for idx, ax in zip(idx_to_plot, ax_.flatten()):
            _show_individuals(ax, self.positions[idx, ...],
                              xbounds=(xmin, xmax), ybounds=(ymin, ymax))
            ax.set_title(f'{int(n_clusters[idx])} clusters \n (Frame: {idx}, Time: {self.timestamps[idx]})',
                         size='small')
        fig.suptitle('Example frames with splitting')
        # plt.tight_layout()

        plt.show()

    def compute_velocities(self) -> np.ndarray:
        """
        Returns:
            np.ndarray: (n_timepoints, 2) array of velocity vectors.
        """
        velocities = np.nan * np.empty_like(self.positions)
        velocities[:-1, ...] = ((self.positions[1:, ...] - self.positions[:-1, ...]).T /
                                (self.timestamps[1:, ...] - self.timestamps[:-1, ...])).T
        return velocities

    def compute_polarisation(self) -> np.ndarray:
        """
        Returns:
            np.ndarray: (n_timepoints, 2) array of the (vector) polarisation order parameter.
        """

        v_magnitude = np.linalg.norm(self.velocities, axis=2)
        pol = np.nanmean((self.velocities.T / v_magnitude.T).T, axis=1)
        return pol

    def get_median_nnd(self):
        """
        Returns:
            The median nearest neighbour distance, computed across individuals and across time.
        """
        nnd = np.full((self.n_timepoints, self.n_fish), np.nan)
        for t in range(self.n_timepoints):
            D = squareform(pdist(self.positions[t]))
            np.fill_diagonal(D, np.nan)
            nnd[t] = np.nanmin(D, axis=0)

        return np.nanmedian(nnd.ravel())
    
    def save_data(self, dest: str = None, filename: str = None) -> None:
        if filename is None:
            filename = os.path.splitext(os.path.basename(self.source_name))[0]

        if dest is None:
            dest = '.'
        pol_file = os.path.join(dest, f'{filename}_pol.csv')
        meta_file = os.path.join(dest, f'{filename}_metadata.txt')
        np.savetxt(pol_file,
                   np.vstack((self.timestamps, self.polarisation.T)).T,
                   header='time,pol_x,pol_y',
                   delimiter=',',
                   comments='', )

        with open(meta_file, 'w') as f:
            f.writelines(self.metadata)

        print(f'Polarisation data: {pol_file}')
        print(f'Metadata: {meta_file}')

    def save_data(self, dest: str = None, filename: str = None) -> None:
        """ Save trajectory data (position and velocity information) to a npz file. """

        if filename is None:
            filename = os.path.splitext(os.path.basename(self.source_name))[0]

        if dest is None:
            dest = '.'

        pol_file = os.path.join(dest, f'{filename}_pol.csv')
        traj_file = os.path.join(dest, f'{filename}_traj.npz')
        meta_file = os.path.join(dest, f'{filename}_metadata.txt')

        np.savetxt(pol_file,
                   np.vstack((self.timestamps, self.polarisation.T)).T,
                   header='time,pol_x,pol_y',
                   delimiter=',',
                   comments='', )
        np.savez(traj_file, positions=self.positions, velocities=self.velocities)
        with open(meta_file, 'w') as f:
            f.writelines(self.metadata)

        print(f'Polarisation data: {pol_file}')
        print(f'Trajectory data: {traj_file}')
        print(f'Metadata: {meta_file}')

    @staticmethod
    def load_trex_trajectories(source: str) -> Tuple[np.ndarray, np.ndarray]:
        """ Load trex datafiles from the speficied source directory.
        Args:
            source: Path of the source directory containing the trex tracked files.

        Returns:
            np.ndarray: (n_timepoints, n_fishes, 2) array containing individual fish trajectories
        """

        data_files = glob.glob(f'{source}/*_fish*.csv')

        # Find the total number of time points and pre-allocate array.
        n_timepoints = -np.inf
        timestamps = None
        for filename in tqdm(data_files, desc='Loading'):
            df = pd.read_csv(filename)
            # print(df.shape[0])
            if df.shape[0] > n_timepoints:
                n_timepoints = df.shape[0]
                timestamps = df[['time']].to_numpy().squeeze()

        trajectories = np.empty((n_timepoints, len(data_files), 2))
        # print(trajectories.shape)
        # print(timestamps.shape)
        # print(n_timepoints)

        for i, filename in enumerate(tqdm(data_files, desc='Loading')):
            df = pd.read_csv(filename)
            # print(df.shape)
            df_timestamps = df[['time']].to_numpy().squeeze()
            start_idx, end_idx = np.argmax(timestamps == df_timestamps[0]), np.argmax(
                timestamps == df_timestamps[-1]) + 1
            # print(start_idx, end_idx)
            trajectories[start_idx:end_idx, i, :] = df[['X#wcentroid (cm)', 'Y#wcentroid (cm)']]
            # trajectories[np.where(timestamps = df_timestamps[0])]

        trajectories[np.isinf(trajectories)] = np.nan

        return timestamps, trajectories

    @staticmethod
    def load_fasttrack_trajectories(source: str, dt: float):
        df = pd.read_csv(source, delimiter='\t')
        n_fish = df.id.max() + 1
        n_frames = df.imageNumber.max() + 1

        positions = np.nan * np.empty((n_frames, n_fish, 2))

        for i in range(n_fish):
            positions[df[df.id == i].imageNumber, i, 0] = df[df.id == i].xBody
            positions[df[df.id == i].imageNumber, i, 1] = df[df.id == i].yBody

        timestamps = np.linspace(0, dt * n_frames, n_frames)
        return timestamps, positions

    @staticmethod
    def load_idtracker_trajectories(source: str):
        data = np.load(source, allow_pickle=True).item()
        dt = 1 / data['frames_per_second']
        positions = data['trajectories']
        n_frames = positions.shape[0]
        timestamps = np.linspace(0, dt * n_frames, n_frames)
        return timestamps, positions


# --- INDIVIDUAL PREPROCESSOR CLASSES ---

class PreprocessorBase(object):
    """ Base-class for preprocessor classes. """

    def __call__(self, tp: TrajectoryProcessor):
        raise NotImplementedError


class SetArenaBounds(PreprocessorBase):
    """ Set upper and lower bounds for position coordinates, and discard points outside bounds. """

    def __init__(self, xbounds, ybounds, type='circle'):
        assert type == 'circle', "Only arena type 'circle' is currently supported. "

        self.xmin, self.xmax = xbounds
        self.ymin, self.ymax = ybounds

    def __call__(self, tp):
        msg = f'Setting arena bounds: X: ({self.xmin}, {self.xmax}), ({self.ymin}, {self.ymax})\n'
        tp.positions[tp.positions[:, :, 0] < self.xmin, ...] = np.nan
        tp.positions[tp.positions[:, :, 0] > self.xmax, ...] = np.nan
        tp.positions[tp.positions[:, :, 1] < self.ymin, ...] = np.nan
        tp.positions[tp.positions[:, :, 1] > self.ymax, ...] = np.nan

        print(msg)
        msg += '\n'
        tp.metadata.append(msg)


class SetSpeedBounds(PreprocessorBase):
    """ Set an upper bound for individual speed, discard points where speed exceeds this. """

    def __init__(self, speed_bounds):
        self.speed_bound = speed_bounds

    def __call__(self, tp):
        msg = f'Setting speed bounds: {self.speed_bound}\n'
        tp.positions[tp.speeds < self.speed_bound[0], ...] = np.nan
        tp.positions[tp.speeds > self.speed_bound[1], ...] = np.nan

        print(msg)
        msg += '\n'
        tp.metadata.append(msg)


class ScaleAndCenter(PreprocessorBase):
    """ Center and optionally rescale the positions so that the center of the arena is at (0, 0), and (optional) the
    x and y coordinates range from -1 to 1. """

    def __init__(self, xbounds=(None, None), ybounds=(None, None), scale=False, xscale=1, yscale=1):

        self.xmin, self.xmax = xbounds
        self.ymin, self.ymax = ybounds
        self.scale = scale
        self.xscale = xscale
        self.yscale = yscale

    def __call__(self, tp):
        if self.xmin is None:
            self.xmin = np.nanmin(tp.positions[:, :, 0])
        if self.xmax is None:
            self.xmax = np.nanmax(tp.positions[:, :, 0])
        if self.ymin is None:
            self.ymin = np.nanmin(tp.positions[:, :, 1])
        if self.ymax is None:
            self.ymax = np.nanmax(tp.positions[:, :, 1])

        xcenter = (self.xmin + self.xmax) / 2
        ycenter = (self.ymin + self.ymax) / 2
        xscale = (self.xmax - self.xmin) / 2
        yscale = (self.ymax - self.ymin) / 2

        msg = f'Recentering data: original center: ({xcenter}, {ycenter})\n'
        tp.positions[:, :, 0] -= xcenter
        tp.positions[:, :, 1] -= ycenter

        if self.scale:
            msg += f'Rescaling data: scale factors: ({xscale}, {yscale})\n'
            tp.positions[:, :, 0] = self.xscale * tp.positions[:, :, 0] / xscale
            tp.positions[:, :, 1] = self.yscale * tp.positions[:, :, 1] / yscale

        print(msg)
        msg += '\n'
        tp.metadata.append(msg)


class IdentifyGroupSplit(PreprocessorBase):
    """ Identifies frames where the fish have split into multiple groups, and drop all fish which are not part of the
        largest cluster. """

    def __init__(self, threshold=6):
        """

        Args:
            threshold: Distance threshold for clustering
        """

        self.thresh = threshold
        self.threshold = None
        self.n_clusters = None
        self.largest_cluster = None

    def __call__(self, tp, preview=False):
        msg = f"Choosing {self.thresh} * median NND as the default threshold. \n"
        self.threshold = self.thresh * tp.get_median_nnd()
        
        if not preview:
            msg += f"Dropping fish that diverge from the main group. \n" \
                   f"Distance threshold = {self.threshold}\n"
        else:
            msg += f"Distance threshold = {self.threshold}\n"

        print(msg)

        if tp.igs is not None:
            self = tp.igs
        if self.n_clusters is None:
            self.n_clusters = np.empty(tp.n_timepoints)
            self.largest_cluster = np.zeros((tp.n_timepoints, tp.n_fish), dtype=bool)
            # n_fish_largest_cluster = np.empty(tp.n_timepoints)
            for t in trange(tp.n_timepoints, desc='Identifying clusters'):
                pw_dist = squareform(pdist(tp.positions[t]))
                neighbour_graph = (pw_dist <= self.threshold)
                n, c = connected_components(neighbour_graph)
                self.n_clusters[t] = n
                largest_cluster_id = np.unique(c)[0]
                self.largest_cluster[t, c == largest_cluster_id] = True

            self.n_clusters -= np.sum(np.isnan(tp.positions[:, :, 0]), axis=1)
        
        msg_ = f'{np.sum(self.n_clusters > 1)}/{tp.n_timepoints} frames identified with group splitting.\n'
        print(msg_)
        msg += msg_

        if preview:
            return self.n_clusters
        else:
            msg += '\n'
            tp.metadata.append(msg)
            tp.positions[~self.largest_cluster, :] = np.nan


class DropFramesWithMissingFish(PreprocessorBase):
    """ Drop frames where more than a certain fraction of fish are missing. """

    def __init__(self, fraction):
        self.fraction = fraction

    def __call__(self, tp):
        msg = f'Dropping frames with more than {self.fraction} fraction ' \
              f'({self.fraction * tp.n_fish:.0f} fish) missing.\n'
        missing_fish = (np.isnan(tp.positions[..., 0]) | np.isnan(tp.positions[..., 1]))
        n_missing_fish = np.sum(missing_fish, axis=1)
        # n_missing_fish = np.sum(np.isnan(tp.positions), axis=(1, ))
        # plt.plot(n_missing_fish)
        # plt.show()
        drop_idx = (n_missing_fish > self.fraction * tp.n_fish)
        tp.positions[drop_idx, ...] = np.nan
        msg += f'{np.sum(drop_idx)}/{tp.n_timepoints} frames dropped.\n'

        print(msg)
        msg += '\n'
        tp.metadata.append(msg)


class RemoveHabituationTime(PreprocessorBase):
    """ Drop frames corresponding to a user specified habituation time. """

    def __init__(self, hab_time):
        """
        Args:
            hab_time: Habituation time (in minutes)
        """
        self.hab_time = hab_time

    def __call__(self, tp):

        dt = tp.timestamps[1] - tp.timestamps[0]
        hab_frames = int(60 * self.hab_time / dt)

        msg = f'Removing habituation time: {self.hab_time} minutes ({hab_frames} frames).\n'

        tp.timestamps = tp.timestamps[hab_frames:]
        tp.positions = tp.positions[hab_frames:, ...]
        tp.velocities = tp.velocities[hab_frames:, ...]
        tp.polarisation = tp.polarisation[hab_frames:, ...]
        tp.speeds = tp.speeds[hab_frames:, ...]
        tp.n_timepoints -= hab_frames

        print(msg)
        msg += '\n'
        tp.metadata.append(msg)


if __name__ == '__main__':
    tp = TrajectoryProcessor('/Users/nabeel/Data/Guy-Hemmigramus-data/10-corrected/10G1902/10G1902.corrected.MTS.txt',
                             format='hemi', dt=1/25., n_fish=10)
    # tp.group_split_diagnostics(threshold=5)