results.py

from abc import ABC, abstractmethod
import copy
import datetime
import os

import h5py
import numpy as np

from fluegg.resultsrecorder import QuantileResultsRecorder


class Results(ABC):

    def __init__(self, results):

        self._configuration = results.configuration()
        self._domain_length = results.domain_length()
        self._time = results.time(slice(None))

    def configuration(self):
        """Returns the configuration dictionary

        Returns
        -------
        dict

        """

        return copy.deepcopy(self._configuration)

    def domain_length(self):
        """Returns the length of the computational domain

        Returns
        -------
        float

        """

        return self._domain_length

    def frac_below_depth(self, time, depth):
        """Fraction of particles below given depth at a time

        Parameters
        ----------
        time : float
            Simulation time, in seconds
        depth : float
            Fractional position along the z-axis, where 0 is the
            water surface and 1 is the bottom of the channel.
        bins : int or array_like, optional
            Number of bins or array of bin edges. The default is 20.

        Returns
        -------
        float
            Fraction of eggs below position

        """
        # calculate quantiles of every particle location
        positions, quantiles = self.position_cfrac(time, 2)
        # make positions positive, so 0 is the surface and 1 is the bottom
        positions = np.abs(positions)
        # find index of position closest to desired depth
        quantile_index = (np.abs(positions - depth)).argmin()
        # use found index to get quantile of depth position
        fraction_below = quantiles[quantile_index]

        return fraction_below

    @classmethod
    def from_hdf(cls, file_path):
        """Initialize a results instance from an HDF file

        Parameters
        ----------
        file_path : str
            Path to HDF file containing results data

        Returns
        -------
        QuantileResults

        """
        members = cls.__dict__['_members']
        attrib = {}

        with h5py.File(file_path, 'r') as f:
            for k in members:
                attrib[k] = f.get(k[1:])[()]

        if '_configuration' in members:
            configuration = attrib['_configuration']
            config_split = configuration.split(", '")
            configuration = []
            for pair in config_split:
                pair = pair.replace('{', '')
                pair = pair.replace('}', '')
                pair = pair.replace('"', '')
                pair = pair.replace("'", '')
                configuration.append(pair.split(': '))
            configuration = dict(configuration)
            attrib['_configuration'] = configuration

        results = cls.__new__(cls)
        for k, v in attrib.items():
            setattr(results, k, v)

        return results

    def particle_count(self):
        """Returns the number of particles.

        Returns
        -------
        int
            Number of particles
        """
        return int(self._configuration["num_eggs"])

    @abstractmethod
    def position_cfrac(self, time, position_axis=0):

        pass

    def position_frac(self, time, position_axis=0, bins=20):
        """Fraction of particles within bins defined by position
        at a given time

        Parameters
        ----------
        time : float
            Simulation time, in seconds
        position_axis : {0, 1, 2}, optional
            Position axis of simulation domain. 0 for x, 1 for y,
            2 for z. The default is 0.
        bins : int or array_like, optional
            Number of bins or array of bin edges. The default is 20.

        Returns
        -------
        numpy.ndarray, numpy.ndarray
            edges, frac

        """

        pos, cfrac = self.position_cfrac(time, position_axis)

        if np.isscalar(bins):
            edges = np.linspace(pos.min(), pos.max(), bins + 1)
        else:
            edges = np.array(bins)

        interp_p = np.interp(edges, pos, cfrac)
        frac = np.diff(interp_p)

        return edges, frac

    @abstractmethod
    def quantile(self, q, position_axis=0):

        pass

    def time(self, index):
        """Simulation time

        Parameters
        ----------
        index : int or slice
            Index of time to return

        Returns
        -------
        float
            Simulation time, in seconds

        """

        return self._time[index]

    @abstractmethod
    def time_cfrac(self, position, position_axis=0):

        pass

    def time_frac(self, position, bin_width, position_axis=0):
        """Fraction of particles within a position bin through time

        Parameters
        ----------
        position : float
            Fractional position in simulation domain
        bin_width : float
            Width of position bin, in meters

        Returns
        -------
        numpy.ndarray, numpy.ndarray
            time, frac_in_bin

        """

        pos_lo = position - bin_width / 2 / self.domain_length()
        pos_hi = position + bin_width / 2 / self.domain_length()

        _, cfrac_lo = self.time_cfrac(pos_lo, position_axis=position_axis)
        _, cfrac_hi = self.time_cfrac(pos_hi, position_axis=position_axis)

        frac_in_bin = cfrac_lo - cfrac_hi

        return self._time.copy(), frac_in_bin

    def to_hdf(self, file_path):
        """Store results data in an HDF file

        Parameters
        ----------
        file_path : str
            Path to HDF file

        """

        if os.path.exists(file_path):
            now = datetime.datetime.now()
            date_string = now.strftime('%Y-%m-%d-%H-%M-%S')
            file_path = file_path[:-3] + str(date_string) + file_path[-3:]

        with h5py.File(file_path, 'w') as f:
            f.create_dataset('type', data=self.__class__.__name__)
            for k in self.__class__.__dict__['_members']:
                data = self.__dict__[k]
                if isinstance(data, dict):
                    data = str(data)
                f.create_dataset(k[1:], data=data)


class FullResults(Results):
    """Simulation results stored with full information

    Parameters
    ----------
    results : FullResultsRecorder

    """

    _members = (
        '_positions',
        '_depth',
        '_width',
        '_domain_length',
        '_configuration',
        '_time')

    def __init__(self, results):

        super().__init__(results)

        self._positions = results.positions()
        self._depth = results.depth()
        self._width = results.width()

    def _normalize_axis(self, position_axis):

        if position_axis == 0:
            positions = \
                self._positions[:, :, position_axis] / self._domain_length
        elif position_axis == 1:
            positions = \
                self._positions[:, :, position_axis] / self._width + 0.5
        elif position_axis == 2:
            positions = \
                self._positions[:, :, position_axis] / self._depth

        return positions

    def in_domain(self, time):
        """Returns a results instance containing particles that are
        in the simulation domain at a given time

        Parameters
        ----------
        time : float
            Simulation time, in seconds. `time` must be finite.

        Returns
        -------
        Full Results or None
            Returns None if no particles are in the domain at `time`

        """

        if not np.isfinite(time):
            raise ValueError("time must be a finite value")

        positions_at_time = np.apply_along_axis(
            lambda p: np.interp(
                time,
                self._time,
                p),
            0,
            self._positions)

        particles_in_domain = positions_at_time[:, 0] <= self._domain_length

        new_config = self.configuration()
        new_config['num_eggs'] = sum(x for x in particles_in_domain)

        if np.any(particles_in_domain):
            cls = self.__class__
            res = cls.__new__(cls)
            res._positions = self._positions[:, particles_in_domain, :]
            res._depth = self._depth[:, particles_in_domain]
            res._width = self._width[:, particles_in_domain]
            res._domain_length = self._domain_length
            res._configuration = new_config
            res._time = self._time.copy()
        else:
            res = None

        return res

    def positions(self, normalize=()):
        """Returns particle position array

        Parameters
        ----------
        normalize : tuple of int, optional
            Normalize position array axes (the default is an empty,
            tuple, which doesn't normalize any axes).

        Returns
        -------
        numpy.ndarray
            Position array

        """

        axes = {0, 1, 2}

        # keep these axes in absolute length
        absolute = axes.difference(normalize)

        positions = np.zeros_like(self._positions)

        for a in normalize:
            positions[:, :, a] = self._normalize_axis(a)

        for a in absolute:
            positions[:, :, a] = self._positions[:, :, a]

        return positions

    def position_cfrac(self, time, position_axis=0):
        """Cumulative fraction of particles past positions at a
        given time

        Paramters
        ---------
        time : float
            Simulation time, in seconds
        position_axis : {0, 1, 2}, optional
            Position axis of simulation domain. 0 for x, 1 for y,
            2 for z. The default is 0.

        Returns
        -------
        numpy.ndarray, numpy.ndarray
            position, cfrac

        """

        particle_positions = self._normalize_axis(position_axis)

        positions = np.apply_along_axis(
            lambda p: np.interp(
                time,
                self._time,
                p,
                left=np.nan,
                right=np.nan),
            0,
            particle_positions)

        position = np.sort(positions)
        n_particles = position.shape[0]
        cfrac = np.arange(1, n_particles + 1) / n_particles

        return position, cfrac

    def quantile(self, q, position_axis=0):
        """Returns position quantiles over time

        Parameters
        ----------
        q : array_like of float
            Quantile or sequence of quantiles to compute
        position_axis : {0, 1, 2}, optional
            Position axis of simulation domain. 0 for x, 1 for y,
            2 for z. The default is 0.

        Returns
        -------
        numpy.ndarray

        """

        particle_positions = self._normalize_axis(position_axis)

        quantiles = np.quantile(particle_positions, q, axis=1)

        return self._time.copy(), quantiles

    def time_cfrac(self, position, position_axis=0):
        """Cumulative fraction of particles with time past a given position

        Parameters
        ----------
        position : float
            Fractional position of simulation domain
        position_axis : {0, 1, 2}, optional
            Position axis of simulation domain. 0 for x, 1 for y,
            2 for z. The default is 0.

        Returns
        -------
        numpy.ndarray, numpy.ndarray
            time, cumulative_fraction

        """
        particle_positions = self._normalize_axis(position_axis)

        particles_through_time = np.greater(
            particle_positions,
            position).sum(
            axis=1)

        n_particles = particle_positions.shape[1]
        cfrac = particles_through_time / n_particles

        return self._time.copy(), cfrac


class QuantileResults(Results):
    """Simulation results stored as quantiles to optimize memory

    Parameters
    ----------
    results : ResultsRecorder,

    """

    _members = (
        '_configuration',
        '_time',
        '_positions',
        '_quantiles',
        '_domain_length')

    def __init__(self, results):

        super().__init__(results)

        # if an AttributeError is raised, assume results is not a
        # QuantileResultsRecorder
        try:
            self._quantiles = results.quantiles()
            self._positions = results.positions()
        except AttributeError:
            # use same quantiles as QuantileResultsRecorder
            self._quantiles = QuantileResultsRecorder.quantiles()
            positions = results.positions((0, 1, 2))
            self._positions = np.quantile(positions, self._quantiles, axis=1)

    def position_cfrac(self, time, position_axis=0):
        """Cumulative fraction of particles past positions at a
        given time

        Paramters
        ---------
        time : float
            Simulation time, in seconds
        position_axis : {0, 1, 2}, optional
            Position axis of simulation domain. 0 for x, 1 for y,
            2 for z. The default is 0.

        Returns
        -------
        numpy.ndarray, numpy.ndarray
            position, cfrac

        """

        time_axis = 1

        position = np.apply_along_axis(
            lambda p: np.interp(
                time,
                self._time,
                p,
                left=np.nan,
                right=np.nan),
            time_axis,
            self._positions[:, :, position_axis])

        return position, self._quantiles.copy()

    def quantile(self, q, position_axis=0):
        """Returns position quantiles over time

        Parameters
        ----------
        q : array_like of float
            Quantile or sequence of quantiles to compute
        position_axis : {0, 1, 2}, optional
            Position axis of simulation domain. 0 for x, 1 for y,
            2 for z. The default is 0.

        Returns
        -------
        numpy.ndarray, numpy.ndarray
            time, quantiles

        """

        quantile_axis = 0
        q = np.array(q)
        q_vs_time = np.apply_along_axis(
            lambda p: np.interp(
                q,
                self._quantiles,
                p,
                left=np.nan,
                right=np.nan),
            quantile_axis,
            self._positions[:, :, position_axis])

        return self._time.copy(), q_vs_time

    def time_cfrac(self, position, position_axis=0):
        """Cumulative fraction of particles with time past a given position

        Parameters
        ----------
        position : float
            Fractional position of simulation domain
        position_axis : {0, 1, 2}, optional
            Position axis of simulation domain. 0 for x, 1 for y,
            2 for z. The default is 0.

        Returns
        -------
        numpy.ndarray, numpy.ndarray
            time, cumulative_fraction

        """

        quantile_axis = 0

        q_through_time = np.apply_along_axis(
            lambda pq: np.interp(
                position,
                pq,
                self._quantiles,
                left=0,
                right=1),
            quantile_axis,
            self._positions[:, :, position_axis])
        cfrac = 1 - q_through_time

        return self._time.copy(), cfrac


def load_from_hdf(file_path):
    with h5py.File(file_path, 'r') as f:
        set_type = f.get('type')[()]
        if set_type == "FullResults":
            return FullResults.from_hdf(file_path)
        elif set_type == "QuantileResults":
            return QuantileResults.from_hdf(file_path)
        else:
            raise ValueError("Invalid results file type: {}".format(set_type))