def get_dict_func_getter(d, label=''):
    def func(key):
        try:
            if callable(key):
                return key
            return d[key]
        except KeyError:
            msg = "The value '%s' is not a valid %s type."
            raise KeyError(msg % (label, key))
    return func def get_bond_type(element1, element2, dist):
    """
    Get the bond type between two elements based on their distance.

    If there is no bond, return None.

    Parameters
    ----------
    element1 : str
        The element of the first atom
    element2 : str
        The element of the second atom
    dist : float
        The distance between the two atoms
    Returns
    -------
    key : str
        The type of the bond
    """
    bad_eles = [x for x in (element1, element2) if x not in BOND_LENGTHS]
    if len(bad_eles):
        msg = "The following elements are not in BOND_LENGTHS: %s" % bad_eles
        warnings.warn(msg)
        return

    for key in TYPE_ORDER[::-1]:
        try:
            cutoff = BOND_LENGTHS[element1][key] + BOND_LENGTHS[element2][key]
            if dist < cutoff:
                return key
        except KeyError:
            continue def get_connections(elements1, coords1, elements2=None, coords2=None):
    """
    Return a dictionary edge list

    If two sets of elements and coordinates are given, then they
    will be treated as two disjoint sets of atoms.

    Each value is is a tuple of the index of the connecting atom and the bond
    order as a string. Where the bond order is one of ['1', 'Ar', '2', '3'].

    Note: If two sets are given, this returns only the connections from the
    first set to the second. This is in contrast to returning connections from
    both directions.

    Parameters
    ----------
    elements1 : list
        All the elements in set 1.

    coords1 : array, shape=(n_atoms, 3)
        The coordinates of the atoms in set 1.

    elements2 : list, default=None
        All the elements in set 2.

    coords2 : array, shape=(n_atoms, 3), default=None
        The coordinates of the atoms in set 2.

    Returns
    -------
    connections : dict, int->dict
        Contains all atoms that are connected to each atom and bond type.
    """
    disjoint = True
    if elements2 is None or coords2 is None:
        disjoint = False
        elements2 = elements1
        coords2 = coords1

    dist_mat = cdist(coords1, coords2)
    connections = {i: {} for i in range(len(elements1))}
    for i, element1 in enumerate(elements1):
        for j, element2 in enumerate(elements2):
            if not disjoint and i >= j:
                continue

            dist = dist_mat[i, j]
            bond_type = get_bond_type(element1, element2, dist)
            if not bond_type:
                continue

            connections[i][j] = bond_type
            if not disjoint:
                connections[j][i] = bond_type
    return connections def get_depth_threshold_mask_connections(connections, min_depth=0,
                                         max_depth=numpy.inf):
    """
    Get the depth threshold mask from connections.

    Parameters
    ----------
    connections : dict, index->list of indices
        A dictionary that contains lists of all connected atoms.


    min_depth : int, default=0
        The minimum depth to allow in the masking

    max_depth : int, default=numpy.inf
        The maximum depth to allow in the masking

    Returns
    -------
    mask : numpy.array, shape=(len(connections), len(connections))
        A mask of all the atoms that are less than or equal to `max_depth`
        away.
    """
    if max_depth < 1:
        max_depth = numpy.inf
    if not min_depth and max_depth is numpy.inf:
        V = len(connections)
        return numpy.ones((V, V)).astype(bool)
    dist = get_graph_distance(connections)
    return (min_depth <= dist) & (dist <= max_depth) def get_graph_distance(connections):
    """
    Compute the graph distance between all pairs of atoms using Floyd-Warshall

    Parameters
    ----------
    connections : dict, index->list of indices
        A dictionary that contains lists of all connected atoms.

    Returns
    -------
    dist : numpy.array, shape=(len(connections), len(connections))
        The graph distance between all pairs of atoms
    """
    # Floyd-Warshall
    V = len(connections)
    dist = numpy.ones((V, V)) * numpy.inf
    numpy.fill_diagonal(dist, numpy.zeros(V))
    for key, values in connections.items():
        for val in values:
            dist[key, val] = 1

    for k in range(V):
        for i in range(V):
            for j in range(V):
                temp = dist[i, k] + dist[k, j]
                if dist[i, j] > temp:
                    dist[i, j] = temp
    return dist def lerp_smooth(x):
def multi_beta(f):
def get_dict_func_getter(d, label=''):
    def func(key):
def get_bond_type(element1, element2, dist):
def get_connections(elements1, coords1, elements2=None, coords2=None):
def get_graph_distance(connections):
def get_depth_threshold_mask_connections(connections, min_depth=0,
                                         max_depth=numpy.inf):
    def __init__(self, connections=None, coords=None, numbers=None,
                 elements=None, unit_cell=None):
    def _none_check(self, x):
    def fill_in_crystal(self, radius=None, units=None):
    def _expand_connections(self, offsets):
    def connections(self):
    def unit_cell(self):
    def coords(self):
    def numbers(self):
    def elements(self):
def get_coulomb_matrix(numbers, coords, alpha=1, use_decay=False):
def get_element_pairs(elements):
def deslugify(string):
def cosine_decay(R, r_cut=6.):
    def __init__(self, values, depth, add_unknown=False,
                 use_comb_idxs=False):
    def is_valid(self, values):
    def __len__(self):
    def __iter__(self):
    def __getitem__(self, key):
    def _inner(self, key, idxs):
    def get_idx_iter(self, key, other=None):
    def get_value_order(self):
    def _get_form_indices(value_length, depth):
    def get_index_mapping(values, depth, idx_groups):
def needs_reversal(chain):
def sort_chain(chain):
def get_angles(coords):
def _load_transformer(data):
def load_json(f):
def _radial_iterator(X, r_max):
def _unit_iterator(X, unit_max):
SMOOTHING_FUNCTIONS = {
    "norm_cdf": multi_beta(scipy.stats.norm.cdf),
    "zero_one": lambda x, beta: (beta * x > 0.).astype(float),
    "expit": multi_beta(expit),
    "tanh": lambda x, beta: (numpy.tanh(beta * x) + 1) / 2,
    "norm": multi_beta(scipy.stats.norm.pdf),
    "circ": scipy.stats.vonmises.pdf,
    "expit_pdf": multi_beta(scipy.stats.logistic.pdf),
    "spike": lambda x, beta: (numpy.abs(beta * x) < 1.).astype(float),
    "lerp": multi_beta(lerp_smooth),
}
SPACING_FUNCTIONS = {
    "log": lambda x: numpy.log(x),
    "inverse": lambda x: 1 / x,
    "linear": lambda x: x,
}
    V = len(connections)
        V = len(connections)
    X = numpy.array(X)
    X = numpy.array(X)
class LazyValues(object):
class IndexMap(object): def get_coulomb_matrix(numbers, coords, alpha=1, use_decay=False):
    r"""
    Return the coulomb matrix for the given coords and numbers.

    .. math::

        C_{ij} = \begin{cases}
            \frac{Z_i Z_j}{\| r_i - r_j \|^\alpha} & i \neq j\\
            \frac{1}{2} Z_i^{2.4} & i = j
        \end{cases}

    Parameters
    ----------
    numbers : array-like, shape=(n_atoms, )
        The atomic numbers of all the atoms

    coords : array-like, shape=(n_atoms, 3)
        The xyz coordinates of all the atoms (in angstroms)

    alpha : number, default=6
        Some value to exponentiate the distance in the coulomb matrix.

    use_decay : bool, default=False
        This setting defines an extra decay for the values as they get futher
        away from the "central atom". This is to alleviate issues the arise as
        atoms enter or leave the cutoff radius.

    Returns
    -------
    top : array, shape=(n_atoms, n_atoms)
        The coulomb matrix
    """
    top = numpy.outer(numbers, numbers).astype(numpy.float64)
    r = cdist(coords, coords)
    if use_decay:
        other = cdist([coords[0]], coords).reshape(-1)
        r += numpy.add.outer(other, other)

    r **= alpha

    with numpy.errstate(divide='ignore', invalid='ignore'):
        numpy.divide(top, r, top)
    numpy.fill_diagonal(top, 0.5 * numpy.array(numbers) ** 2.4)
    top[top == numpy.Infinity] = 0
    top[numpy.isnan(top)] = 0
    return top def get_element_pairs(elements):
    """
    Extract all the element pairs in a molecule.

    Parameters
    ----------
    elements : list
        All the elements in the molecule

    Returns
    -------
    value : list
        All the element pairs in the molecule
    """
    # This is like computing set(combinations(sorted(elements), 2))
    # We do this because it scales with elements instead of atoms.
    counts = Counter(elements)
    pairs = {}
    order = sorted(counts)
    for i, x in enumerate(order):
        for j, y in enumerate(order):
            if i > j:
                continue
            if x == y and counts[x] < 2:
                continue
            pairs[x, y] = 1
    return list(pairs.keys()) def deslugify(string):
    """
    Convert a string to a feature name and its parameters.

    Parameters
    ----------
        string : str
            The slug string to extract values from.

    Returns
    -------
        name : str
            The name of the class corresponding to the string.

        final_params : dict
            A dictionary of the feature parameters.
    """
    parts = string.split('__')
    name = parts[0]
    params = parts[1:]
    swap = {
        'None': None,
        'True': True,
        'False': False,
    }
    final_params = dict()
    for param in params:
        arg, value = param.split('=')
        try:
            value = int(value)
        except ValueError:
            try:
                value = float(value)
            except ValueError:
                pass
        if value in swap:
            value = swap[value]
        final_params[arg] = value
    return name, final_params def sort_chain(chain):
    """
    Sort a chain from the inside out.

    Parameters
    ----------
    chain : tuple
        A tuple of elements to treat as a chain

    Returns
    -------
    chain : tuple
        The sorted chain
    """
    if needs_reversal(chain):
        return chain[::-1]
    return chain def needs_reversal(chain):
    """
    Determine if the chain needs to be reversed.

    This is to set the chains such that they are in a canonical ordering

    Parameters
    ----------
    chain : tuple
        A tuple of elements to treat as a chain

    Returns
    -------
    needs_flip : bool
        Whether or not the chain needs to be reversed
    """
    x = len(chain)
    if x == 1:
        first = 0
        second = 0
    else:
        q, r = divmod(x, 2)
        first = q - 1
        second = q + r

    while first >= 0 and second < len(chain):
        if chain[first] > chain[second]:
            # Case where order reversal is needed
            return True
        elif chain[first] == chain[second]:
            # Indeterminate case
            first -= 1
            second += 1
        else:
            # Case already in the correct order
            return False
    return False def load_json(f):
    """
    Load the model data from a json file

    Parameters
    ----------
    f : str or file descriptor
        The path to save the data or a file descriptor to save it to.

    Returns
    -------
    obj : Transformer
        The transformer object.
    """
    try:
        data = json.load(f)
    except AttributeError:
        with open(f, 'r') as in_file:
            data = json.load(in_file)
    return _load_transformer(data) class IndexMap(object):
    '''
    An object to handle dynamic mapping of groups to indices.

    The intention of the class is to allow for dynamic subselection from lists
    to give new mapping groups.

    For example, with the a group of values that have length three might be
    reduced as follows:
        ('A', 'B', 'C') -> ('A', 'C')
    using some predefined index selection. Then, this new reduced value is
    used to find the index for this new value just like a dict that maps
    the new shorter values to an int index.

    This class also allows handling of groups that are not in the map.

    Parameters
    ----------
    values : list of tuples
        A collection of values overwhich the mapping will be done.

    depth : int
        The number of values that are retained in the subselection.

    add_unknown : bool, default=False
        Whether or not to allocate an UNKNOWN index.

    use_comb_idxs : bool, default=False
        Whether or not to use all combinations of indices when doing the
        subselection. If this is False, it will use the old style of trying to
        select indices from the center of the value outward.
    '''
    def __init__(self, values, depth, add_unknown=False,
                 use_comb_idxs=False):
        self.values = values
        self.depth = depth
        self.add_unknown = add_unknown
        self.use_comb_idxs = use_comb_idxs
        length = len(values[0])
        if use_comb_idxs:
            idx_values = list(range(length))
            self.idx_groups = list(combinations(idx_values, self.depth))
        else:
            self.idx_groups = IndexMap._get_form_indices(length, depth)
        self._mapping = IndexMap.get_index_mapping(values, depth,
                                                   self.idx_groups)

    def is_valid(self, values):
        return self.values == values

    def __len__(self):
        return len(self._mapping) + int(self.add_unknown)

    def __iter__(self):
        for x in self.get_value_order():
            yield x

    def __getitem__(self, key):
        return self._inner(key, self.idx_groups[0])

    def _inner(self, key, idxs):
        key = sort_chain(tuple(key[i] for i in idxs))
        if key not in self._mapping and self.add_unknown:
            return -1
        return self._mapping[key]

    def get_idx_iter(self, key, other=None):
        if other is None:
            other = tuple()
        for idxs in self.idx_groups:
            try:
                yield other + (self._inner(key, idxs), )
            except KeyError:
                yield None

    def get_value_order(self):
        base = [None for x in range(len(self))]
        for key, val in self._mapping.items():
            base[val] = key
        if self.add_unknown:
            base[-1] = (UNKNOWN, )
        return base

    @staticmethod
    def _get_form_indices(value_length, depth):
        """
        Determine the indices to select from a value for a given form.

        One can view this function as a way to get the indices to extract all
        the unique subchains from a single chain starting from the center of a
        given size.

        Ex:
            (A, B, C) -> [0, 1, 2]  # Depth 3
            (A, B, C) -> [0, 2]     # Depth 2
            (A, B, C) -> [1]        # Depth 1
            (A, B, C) -> []         # Depth 0

        Parameters
        ----------
        value_length : int
            The length of a value to select from.

        depth : int
            The number of elements to use from the value.

        Returns
        -------
        final : list of tuples of ints
            A list of list of indices to select.
        """
        if depth < 1:
            return [tuple()]

        if value_length < 1:
            raise ValueError("No values to use.")

        if depth >= value_length:
            return [tuple(range(value_length))]

        middle_idx = value_length // 2
        even = not (value_length % 2)
        both = even and depth % 2
        half_depth = depth // 2
        start = middle_idx - half_depth - both
        end = middle_idx + half_depth + (not even)
        res = list(range(start, end))
        if not even and not (depth % 2):
            res.remove(middle_idx)
        final = [tuple(res)]
        if both:
            final.append(tuple(i + 1 for i in res))
        return final

    @staticmethod
    def get_index_mapping(values, depth, idx_groups):
        """
        Determine the ordering and mapping of feature groups.

        Parameters
        ----------
        values : list
            A list of possible values.

        depth : int
            The number of elements to use from each values value.

        idx_groups : list of list of int
            A list of list of indices to select.

        Returns
        -------
        mapping : dict(key)->int
            A dict that gives the mapping index for a given key.
        """
        if depth < 1:
            # Just a constant value
            return {tuple(): 0}
        temp_values = [[tuple(x[i] for i in idxs) for x in values]
                       for idxs in idx_groups]
        new_values = set(sort_chain(x) for x in sum(temp_values, []))
        mapping = {key: i for i, key in enumerate(sorted(new_values))}
        return mapping UNKNOWN = 'UNKNOWN'
import unittest
import os
import json
import warnings
import numpy
from molml.utils import get_dict_func_getter
from molml.utils import get_bond_type
from molml.utils import get_connections, get_depth_threshold_mask_connections
from molml.utils import get_graph_distance
from molml.utils import LazyValues, get_smoothing_function
from molml.utils import get_coulomb_matrix, get_element_pairs
from molml.utils import deslugify
from molml.utils import sort_chain, needs_reversal
from molml.utils import load_json, IndexMap
from molml.constants import UNKNOWN




DATA_PATH = os.path.join(os.path.dirname(__file__), "data")
ELEMENTS = ['C', 'H', 'H', 'H', 'H']
NUMBERS = [6, 1, 1, 1, 1]
COORDS = [
    [0.99826008, -0.00246000, -0.00436000],
    [2.09021016, -0.00243000, 0.00414000],
    [0.63379005, 1.02686007, 0.00414000],
    [0.62704006, -0.52773003, 0.87811010],
    [0.64136006, -0.50747003, -0.90540005],
]
CONNECTIONS = {
    0: {1: "1", 2: "1", 3: "1", 4: "1"},
    1: {0: "1"},
    2: {0: "1"},
    3: {0: "1"},
    4: {0: "1"},
}
UNIT_CELL = [
    [2., .1, 0.],
    [0., 1., 0.],
    [0., 0., 1.]
]


class UtilsTest(unittest.TestCase):

    def test_smoothing_functions(self):
        # These are mostly sanity checks for the functions
        expected = {
            "norm_cdf": numpy.array([0., 0.158655, 0.308538, 0.5, 0.691462,
                                     0.841345, 1.]),
            "zero_one": numpy.array([0., 0., 0., 0., 1., 1., 1.]),
            "expit": numpy.array([0., 0.268941, 0.377541, 0.5, 0.622459,
                                  0.731059, 1.]),
            "tanh": numpy.array([0., 0.11920292, 0.26894142, 0.5, 0.73105858,
                                 0.88079708, 1.]),
            "norm": numpy.array([0., 0.241971, 0.352065, 0.398942, 0.352065,
                                 0.241971, 0.]),
            "circ": numpy.array([0.220598, 0.215781, 0.302338, 0.34171,
                                 0.302338, 0.215781, 0.220598]),
            "expit_pdf": numpy.array([0., 0.196612, 0.235004, 0.25, 0.235004,
                                      0.196612, 0.]),
            "spike": numpy.array([0., 0., 1., 1., 1., 0., 0.]),
        }

        values = numpy.array([-1000., -1., -0.5, 0, 0.5, 1., 1000.])
        for key, expected_value in expected.items():
            with numpy.warnings.catch_warnings():
                numpy.warnings.filterwarnings('ignore', 'overflow')

                res = get_smoothing_function(key)(values, 1.)