Source code for dcma.tools

"""
Basic general toolbox.
"""

import inspect
import os
import pkg_resources

import numpy as np


def path_relative_to_module(relative_path): # type (object) -> object
    """Get file from a path that is relative to caller's module.
    Returns:    absolute path as string"""
    caller = inspect.stack()[1]
    mod = inspect.getmodule(caller[0])
    return os.path.normpath(pkg_resources.resource_filename(
        mod.__name__, relative_path))


def bin_edges_to_centers(edges):
    centers = [0.5 * (edges[i + 1] + edges[i])
               for i in range(len(edges) - 1)]
    return np.array(centers)


def bin_edges_to_widths(edges):
    widths = [abs(edges[i + 1] - edges[i])
              for i in range(len(edges) - 1)]
    return np.array(widths)


def are_edges_consistent(edges_list, weights=None):
    """
    Check if the bin edges for different replicas are consistent.
    Args:
        edges_list: a list of numpy arrays, where each array holds a set of edges
        weights: weights for the replicas

    Returns:
        A boolean

    """
    if weights is None:
        weights = [1.0 for _ in edges_list]
    edges = np.average(edges_list, axis=0, weights=weights)
    for e in edges_list:
        if not np.all(
            np.absolute((e - edges) / (edges + 1e-4)) < 0.1
        ):
            return False
    return True


class UnitConverter(object):
    """
    In the c++ part of the program, all diffusion (and free energy profiles)
    are normalized, i.e. they are agnostic of the coordinates. This converter
    class helps to convert between units.

        - Permeabilities are given in cm/s.
        - Diffusion profiles are given in [10^-5 cm^2/s].
        - Free energies are given in [k_B T].
    """

    def __init__(self, edges, weights=None, test_bin_consistency=True):
        if not type(edges) is list:
            edges = [edges]
        assert all(type(e) is np.ndarray for e in edges)
        if weights is None:
            weights = [1.0 for _ in edges]
        self._edges = np.average(edges, axis=0, weights=weights)
        self._bin_widths = bin_edges_to_widths(self._edges)
        self._bin_centers = bin_edges_to_centers(self._edges)
        if test_bin_consistency:
            assert are_edges_consistent(edges, weights)

    @staticmethod
    def from_transitions(transition_matrices):
        return UnitConverter(
            [tmat.edges for tmat in transition_matrices],
            [tmat.weight for tmat in transition_matrices]
        )

    @property
    def bin_widths(self):
        """ bin width in Angstrom"""
        return self._bin_widths

    @property
    def n_bins(self):
        """ number of bins """
        return len(self._bin_widths)

    @property
    def edges(self):
        return self._edges

    @property
    def bin_centers(self):
        """ bin centers in Angstrom"""
        return self._bin_centers

    def get_generalized_centers(self):
        """
        Returns an array of size (n_bins + 2) that contains the bin centers plus "virtual" centers to the left and
        right. The latter are required to properly define the effect of biasing potentials at the edges.
        """
        generalized_centers = self.bin_centers
        d = (self.bin_widths[0] + self.bin_widths[-1]) / 2.0
        beyond_left = generalized_centers[0] - d
        beyond_right = generalized_centers[-1] + d
        generalized_centers = np.insert(generalized_centers, 0, beyond_left)
        generalized_centers = np.append(generalized_centers, [beyond_right])
        return generalized_centers

    def normalize_diffusion(self, diff, apply_log=False):
        """
        Args:
            diff: diffusion profile in 10^-5 cm^2/s
            apply_log:  Whether to apply a natural logarithm
                        to the dimensionless diffusion.
                        (log D is used internally by the c++ code
                        and constant diffusion profiles (initial profiles)
                        are given as log D)

        Returns:
            Dimensionless diffusion profile
        """
        d = diff / self._bin_widths**2 / 10
        # (factor 10 converts from 10e-5 cm^2/s to A^2/ps)
        if apply_log:
            d = np.log(d)
        return d

    def denormalize_diffusion(self, diff):
        """ The reverse transformation of normalize_diffusion."""
        return diff * self._bin_widths**2 * 10


def comma_string_to_list(some_string, converter=float):
    """
    Converts "1,1,1" to [1.0, 1.0, 1.0].
    """
    some_list = some_string.strip().split(',')
    return [converter(l) for l in some_list]