amd.compare

Complexity

Total Complexity

60

Size/Duplication

Total Lines	464
Duplicated Lines	0 %

Importance

Changes

0

13 Functions

Rating	Name	Duplication	Size	Complexity
A	_amd_pdist_nns()	0	10	2
C	SDD_EMD()	0	74	10
A	_amd_pdist_nns_low_memory()	0	12	4
B	PDD_cdist_AMD_filter()	0	64	5
C	PDD_cdist()	0	48	9
B	AMD_pdist()	0	46	5
A	_amd_cdist_nns_low_memory()	0	10	4
A	_emd_nns_from_nn_inds()	0	16	5
A	set_verbose()	0	5	1
B	PDD_pdist()	0	40	6
A	EMD()	0	42	2
A	emd()	0	8	1
B	AMD_cdist()	0	48	6

"""Functions for comparing AMDs and PDDs of crystals.
"""

from turtle import update

The name update does not seem to exist in module turtle.

Unused update imported from turtle

from typing import List, Tuple, Optional, Union
import warnings

import numpy as np
import scipy.spatial    # cdist, pdist, squareform
import scipy.optimize   # linear_sum_assignment

from ._network_simplex import network_simplex
from .utils import ETA


def set_verbose(setting, update_rate=100):

Missing function or method docstring

    global _VERBOSE

Global variable '_VERBOSE' undefined at the module level

    global _VERBOSE_UPDATE_RATE

Global variable '_VERBOSE_UPDATE_RATE' undefined at the module level

    _VERBOSE = setting
    _VERBOSE_UPDATE_RATE = update_rate
    

Trailing whitespace

set_verbose(False)

def EMD(
        pdd: np.ndarray,
        pdd_: np.ndarray,
        metric: Optional[str] = 'chebyshev',
        return_transport: Optional[bool] = False,
        **kwargs):
    r"""Earth mover's distance (EMD) between two PDDs, also known as
    the Wasserstein metric.

    Parameters
    ----------
    pdd : ndarray
        PDD of a crystal.
    pdd\_ : ndarray
        PDD of a crystal.
    metric : str or callable, optional
        EMD between PDDs requires defining a distance between rows of two PDDs.
        By default, Chebyshev/l-infinity distance is chosen as with AMDs.
        Can take any metric + ``kwargs`` accepted by
        ``scipy.spatial.distance.cdist``.
    return_transport: bool, optional
        Return a tuple (distance, transport_plan) with the optimal transport.

    Returns
    -------
    float
        Earth mover's distance between PDDs.

    Raises
    ------
    ValueError
        Thrown if the two PDDs do not have the
        same number of columns (``k`` value).
    """

    dm = scipy.spatial.distance.cdist(pdd[:, 1:], pdd_[:, 1:], metric=metric, **kwargs)
    emd_dist, transport_plan = network_simplex(pdd[:, 0], pdd_[:, 0], dm)

    if return_transport:
        return emd_dist, transport_plan.reshape(dm.shape)

    return emd_dist


def SDD_EMD(sdd, sdd_, return_transport: Optional[bool] = False):
    r"""Earth mover's distance (EMD) between two SDDs.

    Parameters
    ----------
    sdd : tuple of ndarrays
        SDD of a crystal.
    sdd\_ : tuple of ndarrays
        SDD of a crystal.
    return_transport: bool, optional
        Return a tuple (distance, transport_plan) with the optimal transport.

    Returns
    -------
    float
        Earth mover's distance between SDDs.

    Raises
    ------
    ValueError
        Thrown if the two SDDs are not of the same order or do not have the
        same number of columns (``k`` value).
    """

    dists, dists_ = sdd[2], sdd_[2]

    # first order SDD, equivalent to PDD
    if dists.ndim == 2 and dists_.ndim == 2:
        dm = scipy.spatial.distance.cdist(dists, dists_, metric='chebyshev')
        emd_dist, transport_plan = network_simplex(sdd[0], sdd_[0], dm)

        if return_transport:
            return emd_dist, transport_plan.reshape(dm.shape)

        return emd_dist

    order = dists.shape[-1]
    n, m = len(sdd[0]), len(sdd_[0])

    dist_cdist = None
    if order == 2:
        dist_cdist = np.abs(sdd[1][:, None] - sdd_[1])
    else:
        dist, dist_ = sdd[1], sdd_[1]

        # take EMDs between finite PDDs in dist column
        weights = np.full((order, ), 1 / order)
        dist_cdist = np.empty((n, m), dtype=np.float64)
        for i in range(n):
            for j in range(m):
                finite_pdd_dm = scipy.spatial.distance.cdist(dist[i], dist_[j], metric='chebyshev')
                dists_emd, _ = network_simplex(weights, weights, finite_pdd_dm)
                dist_cdist[i, j] = dists_emd

        # flatten and compare by linf
        # flat_dist = dist.reshape((n, order * (order - 1)))
        # flat_dist_ = dist_.reshape((m, order * (order - 1)))
        # flat_dist = np.sort(flat_dist, axis=-1)
        # flat_dist_ = np.sort(flat_dist_, axis=-1)
        # dist_cdist = scipy.spatial.distance.cdist(flat_dist, flat_dist_, metric='chebyshev')

    dm = np.empty((n, m), dtype=np.float64)
    for i in range(n):
        for j in range(m):
            cost_matrix = scipy.spatial.distance.cdist(dists[i], dists_[j], metric='chebyshev')
            row_ind, col_ind = scipy.optimize.linear_sum_assignment(cost_matrix)
            dm[i, j] = max(np.amax(cost_matrix[row_ind, col_ind]), dist_cdist[i, j])

    emd_dist, transport_plan = network_simplex(sdd[0], sdd_[0], dm)

    if return_transport:
        return emd_dist, transport_plan.reshape(dm.shape)

    return emd_dist


def AMD_cdist(
        amds: Union[np.ndarray, List[np.ndarray]],
        amds_: Union[np.ndarray, List[np.ndarray]],
        metric: str = 'chebyshev',
        low_memory: bool = False,
        **kwargs
) -> np.ndarray:
    r"""Compare two sets of AMDs with each other, returning a distance matrix.

    Parameters
    ----------
    amds : array_like
        A list of AMDs.
    amds\_ : array_like
        A list of AMDs.
    metric : str or callable, optional
        Usually AMDs are compared with the Chebyshev/l-infinity distance.
        Can take any metric + kwargs accepted by ``scipy.spatial.distance.cdist``.
    low_memory : bool, optional
        Use a slower but more memory efficient method for
        large collections of AMDs (Chebyshev/l-inf distance only).

    Returns
    -------
    ndarray
        A distance matrix shape ``(len(amds), len(amds_))``.
        The :math:`ij` th entry is the distance between ``amds[i]``
        and ``amds[j]`` given by the ``metric``.
    """

    amds, amds_ = np.asarray(amds), np.asarray(amds_)

    if len(amds.shape) == 1:
        amds = np.array([amds])
    if len(amds_.shape) == 1:
        amds_ = np.array([amds_])

    if low_memory:
        if metric != 'chebyshev':
            warnings.warn("Using only allowed metric 'chebyshev' for low_memory", UserWarning)

        dm = np.empty((len(amds), len(amds_)))
        for i, amd_vec in enumerate(amds):
            dm[i] = np.amax(np.abs(amds_ - amd_vec), axis=-1)
    else:
        dm = scipy.spatial.distance.cdist(amds, amds_, metric=metric, **kwargs)

    return dm


def AMD_pdist(
        amds: Union[np.ndarray, List[np.ndarray]],
        metric: str = 'chebyshev',
        low_memory: bool = False,
        **kwargs
) -> np.ndarray:
    """Compare a set of AMDs pairwise, returning a condensed distance matrix.

    Parameters
    ----------
    amds : array_like
        An array/list of AMDs.
    metric : str or callable, optional
        Usually AMDs are compared with the Chebyshev/l-infinity distance.
        Can take any metric + kwargs accepted by ``scipy.spatial.distance.cdist``.
    low_memory : bool, optional
        Optionally use a slightly slower but more memory efficient method for
        large collections of AMDs (Chebyshev/l-inf distance only).

    Returns
    -------
    ndarray
        Returns a condensed distance matrix. Collapses a square distance
        matrix into a vector just keeping the upper half. Use
        ``scipy.spatial.distance.squareform`` to convert to a square distance matrix.
    """

    amds = np.asarray(amds)

    if len(amds.shape) == 1:
        amds = np.array([amds])

    if low_memory:
        m = len(amds)
        if metric != 'chebyshev':
            warnings.warn("Using only allowed metric 'chebyshev' for low_memory", UserWarning)
        cdm = np.empty((m * (m - 1)) // 2, dtype=np.double)
        ind = 0
        for i in range(m):
            ind_ = ind + m - i - 1
            cdm[ind:ind_] = np.amax(np.abs(amds[i+1:] - amds[i]), axis=-1)
            ind = ind_
    else:
        cdm = scipy.spatial.distance.pdist(amds, metric=metric, **kwargs)

    return cdm


def PDD_cdist(
        pdds: List[np.ndarray],
        pdds_: List[np.ndarray],
        metric: str = 'chebyshev',
        **kwargs
) -> np.ndarray:
    r"""Compare two sets of PDDs with each other, returning a distance matrix.

    Parameters
    ----------
    pdds : list of ndarrays
        A list of PDDs.
    pdds\_ : list of ndarrays
        A list of PDDs.
    metric : str or callable, optional
        Usually PDD rows are compared with the Chebyshev/l-infinity distance.
        Can take any metric + kwargs accepted by
        ``scipy.spatial.distance.cdist``.

    Returns
    -------
    ndarray
        Returns a distance matrix shape ``(len(pdds), len(pdds_))``.
        The :math:`ij` th entry is the distance between ``pdds[i]``
        and ``pdds_[j]`` given by Earth mover's distance.
    """

    if isinstance(pdds, np.ndarray):
        if len(pdds.shape) == 2:
            pdds = [pdds]

    if isinstance(pdds_, np.ndarray):
        if len(pdds_.shape) == 2:
            pdds_ = [pdds_]

    n, m = len(pdds), len(pdds_)
    dm = np.empty((n, m))
    if _VERBOSE:

The variable _VERBOSE does not seem to be defined.

        eta = ETA(n * m, update_rate=_VERBOSE_UPDATE_RATE)

The variable _VERBOSE_UPDATE_RATE does not seem to be defined.

    for i in range(n):
        pdd = pdds[i]
        for j in range(m):
            dm[i, j] = EMD(pdd, pdds_[j], metric=metric, **kwargs)
            if _VERBOSE:
                eta.update()

The variable eta does not seem to be defined in case _VERBOSE on line 279 is False. Are you sure this can never be the case?

    return dm


def PDD_pdist(
        pdds: List[np.ndarray],
        metric: str = 'chebyshev',
        **kwargs
) -> np.ndarray:
    """Compare a set of PDDs pairwise, returning a condensed distance matrix.

    Parameters
    ----------
    pdds : list of ndarrays
        A list of PDDs.
    metric : str or callable, optional
        Usually PDD rows are compared with the Chebyshev/l-infinity distance.
        Can take any metric + kwargs accepted by ``scipy.spatial.distance.cdist``.

    Returns
    -------
    ndarray
        Returns a condensed distance matrix. Collapses a square
        distance matrix into a vector just keeping the upper half. Use
        ``scipy.spatial.distance.squareform`` to convert to a square
        distance matrix.
    """

    if isinstance(pdds, np.ndarray):
        if len(pdds.shape) == 2:
            pdds = [pdds]

    m = len(pdds)
    cdm = np.empty((m * (m - 1)) // 2, dtype=np.double)
    if _VERBOSE:

The variable _VERBOSE does not seem to be defined.

        eta = ETA((m * (m - 1)) // 2, update_rate=_VERBOSE_UPDATE_RATE)

The variable _VERBOSE_UPDATE_RATE does not seem to be defined.

    inds = ((i, j) for i in range(0, m - 1) for j in range(i + 1, m))

The variable j does not seem to be defined.

The variable i does not seem to be defined.

    for r, (i, j) in enumerate(inds):
        cdm[r] = EMD(pdds[i], pdds[j], metric=metric, **kwargs)
        if _VERBOSE:
            eta.update()

The variable eta does not seem to be defined in case _VERBOSE on line 322 is False. Are you sure this can never be the case?

    return cdm


def emd(
        pdd: np.ndarray,
        pdd_: np.ndarray,
        metric: Optional[str] = 'chebyshev',
        return_transport: Optional[bool] = False, 

Trailing whitespace

        **kwargs):
    """Alias for amd.emd()."""
    return EMD(pdd, pdd_, metric=metric, return_transport=return_transport, **kwargs)


def PDD_cdist_AMD_filter(
        n: int,
        pdds: List[np.ndarray],
        pdds_: Optional[List[np.ndarray]] = None,
        low_memory: bool = False,
        metric: str = 'chebyshev',
        **kwargs
) -> Tuple[np.ndarray, np.ndarray]:
    r"""For each item in ``pdds``, get the ``n`` nearest items in ``pdds_`` by AMD,
    then compare references to these nearest items with PDDs.
    Tries to comprimise between the speed of AMDs and the accuracy of PDDs.

    If ``pdds_`` is :const:`None`, this essentially sets ``pdds_ = pdds``, i.e.
    do an 'AMD neighbourhood graph' for one set whose weights are PDD distances.

    Parameters
    ----------
    n : int
        Number of nearest neighbours to find.
    pdds : list of ndarrays
        A list of PDDs.
    pdds\_ : list of ndarrays, optional
        A list of PDDs.
    low_memory : bool, optional
        Optionally use a slightly slower but more memory efficient method for
        large collections of AMDs (Chebyshev/l-inf distance only).
    metric : str or callable, optional
        Usually PDD rows are compared with the Chebyshev/l-infinity distance.
        Can take any metric + kwargs accepted by ``scipy.spatial.distance.cdist``.

    Returns
    -------
    tuple of ndarrays (distance_matrix, indices)
        For the :math:`i` th item in reference and some :math:`j<n`,
        ``distance_matrix[i][j]`` is the distance from reference i to its j-th
        nearest neighbour in comparison (after the AMD filter).
        ``indices[i][j]`` is the index of said neighbour in ``pdds_``.
    """

    kwargs = {'metric': metric, **kwargs}
    amds = np.array([np.average(pdd[:, 1:], weights=pdd[:, 0], axis=0)
                     for pdd in pdds])

    if low_memory:
        if metric != 'chebyshev':
            warnings.warn(
                "Using only allowed metric 'chebyshev' for low_memory",
                UserWarning)
        if pdds_ is None:
            inds = _amd_pdist_nns_low_memory(amds, n)
            pdds_ = pdds
        else:
            amds_ = np.array([np.average(pdd[:, 1:], weights=pdd[:, 0], axis=0)
                              for pdd in pdds_])
            inds = _amd_cdist_nns_low_memory(amds, amds_, n) 

Trailing whitespace

    else:
        if pdds_ is None:
            inds = _amd_pdist_nns(amds, n, **kwargs)
        else:
            amd_dm = AMD_cdist(amds, amds_, low_memory=low_memory, **kwargs)

The variable amds_ does not seem to be defined.

            inds = np.array([np.argpartition(row, n)[:n] for row in amd_dm])

    dm, inds = _emd_nns_from_nn_inds(pdds, pdds_, inds, **kwargs)
    return dm, inds


def _amd_pdist_nns(amds, n, **kwargs):
    amd_cdm = AMD_pdist(amds, **kwargs)
    amd_dm = scipy.spatial.distance.squareform(amd_cdm)
    inds = []
    for i, row in enumerate(amd_dm):
        inds_row = np.argpartition(row, n+1)[:n+1]
        inds_row = inds_row[inds_row != i][:n]
        inds.append(inds_row)
    inds = np.array(inds)
    return inds


def _amd_pdist_nns_low_memory(amds, n):
    inds = []
    if _VERBOSE:

The variable _VERBOSE does not seem to be defined.

        eta = ETA(len(amds), update_rate=_VERBOSE_UPDATE_RATE)

The variable _VERBOSE_UPDATE_RATE does not seem to be defined.

    for i, amd_vec in enumerate(amds):
        dists = np.amax(np.abs(amds - amd_vec), axis=-1)
        inds_row = np.argpartition(dists, n+1)[:n+1]
        inds_row = inds_row[inds_row != i][:n]
        inds.append(inds_row)
        if _VERBOSE:
            eta.update()

The variable eta does not seem to be defined in case _VERBOSE on line 424 is False. Are you sure this can never be the case?

    return np.array(inds)


def _amd_cdist_nns_low_memory(amds, amds_, n):
    inds = []
    if _VERBOSE:

The variable _VERBOSE does not seem to be defined.

        eta = ETA(len(amds) * len(amds_), update_rate=_VERBOSE_UPDATE_RATE)

The variable _VERBOSE_UPDATE_RATE does not seem to be defined.

    for amd_vec in amds:
        row = np.amax(np.abs(amds_ - amd_vec), axis=-1)
        inds.append(np.argpartition(row, n)[:n])
        if _VERBOSE:
            eta.update()

The variable eta does not seem to be defined in case _VERBOSE on line 438 is False. Are you sure this can never be the case?

    return np.array(inds)


def _emd_nns_from_nn_inds(pdds, pdds_, inds, **kwargs):
    dm = np.empty(inds.shape)
    if _VERBOSE:

The variable _VERBOSE does not seem to be defined.

        eta = ETA(inds.shape[0] * inds.shape[1], update_rate=_VERBOSE_UPDATE_RATE)

The variable _VERBOSE_UPDATE_RATE does not seem to be defined.

    for i, row in enumerate(inds):
        for i_, j in enumerate(row):
            dm[i, i_] = EMD(pdds[i], pdds_[j], **kwargs)
            if _VERBOSE:
                eta.update()

The variable eta does not seem to be defined in case _VERBOSE on line 450 is False. Are you sure this can never be the case?

    sorted_inds = np.argsort(dm, axis=-1)
    inds = np.take_along_axis(inds, sorted_inds, axis=-1)
    dm = np.take_along_axis(dm, sorted_inds, axis=-1)
    

Trailing whitespace

    return dm, inds