import numpy as np
from typing import Tuple, Optional, Iterable, Union

from .core import *
from .utils import *

__all__ = [
    "unwrap_dimensional",
    "unwrap_arbitrary",
]


def wrap_difference(x: np.ndarray, period: float = 2 * np.pi) -> np.ndarray:
    """
    Wrap a 1D array.
    Args:
        x: (N,) array
        period: interval of the wrapped axis
    Returns:
        (N,) array
    """
    return np.mod(x + period / 2, period) - period / 2


def unwrap_dimensional(
    x: np.ndarray,
    start_pixel: Optional[Union[Tuple[int, int], Tuple[int, int, int]]] = None,
    use_edgelist: bool = False,
    **kwargs,
) -> np.ndarray:
    """
    Unwrap the phase of a 2-D or 3-D array.

    Parameters
    ----------
    x : 2-D or 3-D np.ndarray
        The phase to be unwrapped.
    start_pixel : (2,) or (3,) tuple
        the reference pixel to start unwrapping.
        Default to (0, 0) for 2-D data and (0, 0, 0) for 3-D data.
    use_edgelist : bool
        Whether to use the edgelist method.
        Default to False.
    kwargs : dict
        Other arguments passed to `kamui.unwrap_arbitrary`.

    Returns
    -------
    np.ndarray
        The unwrapped phase of the same shape as x.
    """
    if start_pixel is None:
        start_pixel = (0,) * x.ndim
    assert x.ndim == len(start_pixel), "start_pixel must have the same dimension as x"

    start_i = 0
    for i, s in enumerate(start_pixel):
        start_i *= x.shape[i]
        start_i += s

    if x.ndim == 2:
        edges, simplices = get_2d_edges_and_simplices(x.shape)
    elif x.ndim == 3:
        edges, simplices = get_3d_edges_and_simplices(x.shape)
    else:
        raise ValueError("x must be 2D or 3D")
    psi = x.ravel()

    result = unwrap_arbitrary(
        psi, edges, None if use_edgelist else simplices, start_i=start_i, **kwargs
    )
    if result is None:
        return None

    return result.reshape(x.shape)


def unwrap_arbitrary(
    psi: np.ndarray,
    edges: np.ndarray,
    simplices: Iterable[Iterable[int]] = None,
    method: str = "ilp",
    period: float = 2 * np.pi,
    start_i: int = 0,
    **kwargs,
) -> Optional[np.ndarray]:
    """
    Unwrap the phase of arbitrary data.

    Parameters
    ----------
    psi : 1D np.ndarray of shape (P,)
        The phase (vertices) to be unwrapped.
    edges : 2-D np.ndarray of shape (M, 2)
        The edges of the graph.
    simplices : Iterable[Iterable[int]] of length (N,)
        Each element is a list of vertices that form a simplex (a.k.a elementary cycle).
        The connections should be consistent with the edges.
        This is also used to compute automatic weights for each edge.
        If not provided and method is "ilp", an edgelist-based ILP solver will be used without weighting.
    method : str
        The method to be used. Valid options are "ilp" and "gc", where "gc" correponds to PUMA.
        Default to "ilp".
    period : float
        The period of the phase.
        Default to 2 * np.pi.
    start_i : int
        The index of the reference vertex to start unwrapping.
        Default to 0.
    kwargs : dict
        Other arguments passed to the solver.

    Returns
    -------
    np.ndarray
        The unwrapped phase of the same shape as psi.
    """
    if method == "gc":
        m = puma(psi / period, edges, **kwargs)
        m -= m[start_i]
        result = m * period + psi
    elif method == "ilp":
        if simplices is None:
            m = calculate_m(
                edges,
                np.round((psi[edges[:, 1]] - psi[edges[:, 0]]) / period).astype(
                    np.int64
                ),
                **kwargs,
            )
            if m is None:
                return None
            m -= m[start_i]
            result = m * period + psi
        else:
            diff = wrap_difference(psi[edges[:, 1]] - psi[edges[:, 0]], period)
            k = calculate_k(edges, simplices, diff / period, **kwargs)
            correct_diff = diff + k * period

            result = (
                integrate(
                    np.concatenate((edges, np.flip(edges, 1)), axis=0),
                    np.concatenate((correct_diff, -correct_diff), axis=0),
                    start_i=start_i,
                )
                + psi[start_i]
            )
    else:
        raise ValueError("method must be 'gc' or 'ilp'")
    return result

kamui/__init__.py
yoyololicon/kamui
from scipy.sparse import csgraph as csg
import scipy.sparse as sp
from scipy.optimize import linprog
import numpy as np
from typing import Optional, Iterable

try:
    import maxflow
except ImportError:
    print("PyMaxflow not found, some functions will not be available.")

__all__ = ["integrate", "calculate_k", "calculate_m", "puma"]


def integrate(edges: np.ndarray, weights: np.ndarray, start_i: int = 0):
    """
    Args:
        edges: (N, 2) array of edges
        weights: (N,) array of weights
        start_i: starting index
    Returns:
        (N,) array of integrated weights
    """
    G = sp.csr_matrix((weights, (edges[:, 0], edges[:, 1])))
    N = max(G.shape)
    result = np.zeros(N, dtype=weights.dtype)

    nodes = csg.depth_first_order(G, start_i, directed=True, return_predecessors=False)

    pairs = np.stack([nodes[:-1], nodes[1:]], axis=1)
    for u, v in pairs:
        result[v] = result[u] + G[u, v]

    return result


def calculate_k(
    edges: np.ndarray,
    simplices: Iterable[Iterable[int]],
    differences: np.ndarray,
    weights: Optional[np.ndarray] = None,
    adaptive_weighting: bool = True,
) -> Optional[np.ndarray]:
    """
    Args:
        edges: (M, 2) array of edges
        simplices: (N,) iterable of simplices
        differences: (M,) array of differences, could be float or int
        weights: (M,) array of weights
        adaptive_weighting: whether to use adaptive weighting
    """

    M, N = edges.shape[0], len(simplices)

    edge_dict = {tuple(x): i for i, x in enumerate(edges)}
    if len(edge_dict) != M:
        raise ValueError("edges must be unique")
    rows = []
    cols = []
    vals = []
    for i, simplex in enumerate(simplices):
        u = simplex[-1]
        for v in simplex:
            key = (u, v)
            rows.append(i)
            if key in edge_dict:
                cols.append(edge_dict[key])
                vals.append(1)
            else:
                try:
                    cols.append(edge_dict[(v, u)])
                except KeyError:
                    raise ValueError("simplices contain invalid edges")
                vals.append(-1)
            u = v

    rows = np.array(rows)
    cols = np.array(cols)
    vals = np.array(vals)

    V = sp.csr_matrix((vals, (rows, cols)), shape=(N, M))
    y = V @ differences
    y = np.round(y).astype(np.int64)

    A_eq = sp.csr_matrix(
        (
            np.concatenate((vals, -vals)),
            (np.tile(rows, 2), np.concatenate((cols, cols + M))),
        ),
        shape=(N, M * 2),
    )
    b_eq = -y

    if weights is None:
        if adaptive_weighting:
            nonzero_simplices = np.minimum(np.abs(b_eq), 1)
            W = np.abs(A_eq)
            num_nonzero_simplices = nonzero_simplices @ W
            num_simplices = W.sum(0).A1
            c = num_simplices - num_nonzero_simplices
        else:
            c = np.ones((M * 2,), dtype=np.int64)
    else:
        c = np.tile(c, 2)

    res = linprog(c, A_eq=A_eq, b_eq=b_eq, integrality=1)
    k = res.x[:M] - res.x[M:]
    k = k.astype(np.int64)
    cost = np.abs(V @ k + y).sum()
    if cost > 0:
        print(f"Warning: no solution found.")
        return None
    return k


def calculate_m(
    edges: np.ndarray,
    differences: np.ndarray,
    weights: Optional[np.ndarray] = None,
) -> Optional[np.ndarray]:
    """
    Args:
        edges: (M, 2) array of edges
        differences: (M,) quantised array of differences, must be int
        weights: (M,) array of weights
    """
    assert differences.dtype == np.int64, "differences must be int"
    M = edges.shape[0]
    N = np.max(edges) + 1

    vals = np.concatenate(
        (np.ones((M,), dtype=np.int64), -np.ones((M,), dtype=np.int64))
    )
    rows = np.tile(np.arange(M), 2)
    cols = np.concatenate((edges[:, 0], edges[:, 1]))

    A_eq = sp.csr_matrix(
        (
            np.concatenate((vals, np.ones(M), -np.ones(M))).astype(np.int64),
            (
                np.tile(rows, 2),
                np.concatenate((cols, np.arange(2 * M) + N)),
            ),
        ),
        shape=(M, N + 2 * M),
    )
    if weights is None:
        weights = np.ones((M,), dtype=np.int64)

    c = np.concatenate((np.zeros(N, dtype=np.int64), weights, weights))

    b_eq = differences

    res = linprog(c, A_eq=A_eq, b_eq=b_eq, integrality=1)
    if res.x is None:
        return None
    m = res.x[:N]
    return m.astype(np.int64)


def puma(psi: np.ndarray, edges: np.ndarray, max_jump: int = 1, p: float = 1):
    """
    Args:
        psi: (N,) array
        edges: (M, 2) array of edges
        max_jump: maximum jump step
        p: p-norm
    Returns:
        (N,) array
    """
    if max_jump > 1:
        jump_steps = list(range(1, max_jump + 1)) * 2
    else:
        jump_steps = [max_jump]

    total_nodes = psi.size

    def V(x):
        return np.abs(x) ** p

    K = np.zeros_like(psi)

    def cal_Ek(K, psi, i, j):
        return np.sum(V(K[j] - K[i] - psi[i] + psi[j]))

    prev_Ek = cal_Ek(K, psi, edges[:, 0], edges[:, 1])

    energy_list = []

    for step in jump_steps:
        while 1:
            energy_list.append(prev_Ek)
            G = maxflow.Graph[float]()
            G.add_nodes(total_nodes)

            i, j = edges[:, 0], edges[:, 1]
            psi_diff = psi[i] - psi[j]
            a = (K[j] - K[i]) - psi_diff
            e00 = e11 = V(a)
            e01 = V(a - step)
            e10 = V(a + step)
            weight = np.maximum(0, e10 + e01 - e00 - e11)

            G.add_edges(edges[:, 0], edges[:, 1], weight, np.zeros_like(weight))

            tmp_st_weight = np.zeros((2, total_nodes))

            for i in range(edges.shape[0]):
                u, v = edges[i]
                tmp_st_weight[0, u] += max(0, e10[i] - e00[i])
                tmp_st_weight[0, v] += max(0, e11[i] - e10[i])
                tmp_st_weight[1, u] -= min(0, e10[i] - e00[i])
                tmp_st_weight[1, v] -= min(0, e11[i] - e10[i])

            for i in range(total_nodes):
                G.add_tedge(i, tmp_st_weight[0, i], tmp_st_weight[1, i])

            G.maxflow()

            partition = G.get_grid_segments(np.arange(total_nodes))
            K[~partition] += step

            energy = cal_Ek(K, psi, edges[:, 0], edges[:, 1])

            if energy < prev_Ek:
                prev_Ek = energy
            else:
                K[~partition] -= step
                break

    return K

kamui/core.py
yoyololicon/kamui
import numpy as np
from typing import Tuple, Optional, Iterable

__all__ = ["get_2d_edges_and_simplices", "get_3d_edges_and_simplices"]


def get_2d_edges_and_simplices(
    shape: Tuple[int, int]
) -> Tuple[np.ndarray, Iterable[Iterable[int]]]:
    nodes = np.arange(np.prod(shape)).reshape(shape)
    edges = np.concatenate(
        (
            np.stack([nodes[:, :-1].ravel(), nodes[:, 1:].ravel()], axis=1),
            np.stack([nodes[:-1, :].ravel(), nodes[1:, :].ravel()], axis=1),
        ),
        axis=0,
    )
    simplices = np.stack(
        (
            nodes[:-1, :-1].ravel(),
            nodes[1:, :-1].ravel(),
            nodes[1:, 1:].ravel(),
            nodes[:-1, 1:].ravel(),
        ),
        axis=1,
    ).tolist()
    return edges, simplices


def get_3d_edges_and_simplices(
    shape: Tuple[int, int, int]
) -> Tuple[np.ndarray, Iterable[Iterable[int]]]:
    nodes = np.arange(np.prod(shape)).reshape(shape)
    edges = np.concatenate(
        (
            np.stack([nodes[:, :-1, :].ravel(), nodes[:, 1:, :].ravel()], axis=1),
            np.stack([nodes[:-1, :, :].ravel(), nodes[1:, :, :].ravel()], axis=1),
            np.stack([nodes[:, :, :-1].ravel(), nodes[:, :, 1:].ravel()], axis=1),
        ),
        axis=0,
    )
    simplices = np.concatenate(
        (
            np.stack(
                (
                    nodes[:-1, :-1, :-1].ravel(),
                    nodes[1:, :-1, :-1].ravel(),
                    nodes[1:, 1:, :-1].ravel(),
                    nodes[:-1, 1:, :-1].ravel(),
                ),
                axis=1,
            ),
            np.stack(
                (
                    nodes[:-1, :-1, :-1].ravel(),
                    nodes[:-1, 1:, :-1].ravel(),
                    nodes[:-1, 1:, 1:].ravel(),
                    nodes[:-1, :-1, 1:].ravel(),
                ),
                axis=1,
            ),
            np.stack(
                (
                    nodes[:-1, :-1, :-1].ravel(),
                    nodes[:-1, :-1, 1:].ravel(),
                    nodes[1:, :-1, 1:].ravel(),
                    nodes[1:, :-1, :-1].ravel(),
                ),
                axis=1,
            ),
        ),
        axis=0,
    ).tolist()
    return edges, simplices

kamui/utils.py
yoyololicon/kamui
import setuptools
from pathlib import Path

root_path = Path(__file__).parent
version_file = root_path / "VERSION.txt"

with open("README.md", "r") as fh:
    long_description = fh.read()

setuptools.setup(
    name="kamui",
    author="Chin-Yun Yu",
    author_email="chin-yun.yu@qmul.ac.uk",
    description="A Python package for phase unwrapping",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/yoyololicon/kamui",
    packages=["kamui"],
    install_requires=["numpy", "scipy"],
    extras_require={
        "extra": ["PyMaxflow"],
    },
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],
    python_requires=">=3.6",
    setuptools_git_versioning={
        "enabled": True,
        "version_file": version_file,
        "count_commits_from_version_file": True,
        "dev_template": "{tag}.{branch}{ccount}",
        "dirty_template": "{tag}.{branch}{ccount}",
    },
    setup_requires=["setuptools-git-versioning<2"],
)

setup.py
yoyololicon/kamui