Prepare weights for ilp unwrapping

kamui/__init__.py

@@ -1,5 +1,5 @@
 import numpy as np
-from typing import Tuple, Optional, Iterable, Union
+from typing import Tuple, Optional, Iterable, Union, Any
 
 from .core import *
 from .utils import *
@@ -27,8 +27,9 @@ def unwrap_dimensional(
     start_pixel: Optional[Union[Tuple[int, int], Tuple[int, int, int]]] = None,
     use_edgelist: bool = False,
     cyclical_axis: Union[int, Tuple[int, int]] = (),
-    **kwargs,
-) -> np.ndarray:
+    merging_method: str = 'mean',
+    **kwargs: Any,
+) -> Optional[np.ndarray]:
     """
     Unwrap the phase of a 2-D or 3-D array.
 
@@ -45,6 +46,7 @@ def unwrap_dimensional(
     cyclical_axis : int or (int, int)
         The axis that is cyclical.
         Default to ().
+    merging_method  : Way of combining two phase weights into a single edge weight.
     kwargs : dict
         Other arguments passed to `kamui.unwrap_arbitrary`.
 
@@ -74,6 +76,10 @@ def unwrap_dimensional(
         raise ValueError("x must be 2D or 3D")
     psi = x.ravel()
 
+    weights = kwargs.pop('weights', None)
+    if weights is not None:
+        kwargs['weights'] = prepare_weights(weights, edges=edges, merging_method=merging_method)
+
     result = unwrap_arbitrary(
         psi, edges, None if use_edgelist else simplices, start_i=start_i, **kwargs
     )

kamui/core.py

@@ -4,6 +4,7 @@
 import numpy as np
 from typing import Optional, Iterable
 
+
 try:
     import maxflow
 except ImportError:

kamui/utils.py

@@ -1,7 +1,8 @@
 import numpy as np
-from typing import Tuple, Optional, Iterable, Union
+import numpy.typing as npt
+from typing import Tuple, Iterable, Union
 
-__all__ = ["get_2d_edges_and_simplices", "get_3d_edges_and_simplices"]
+__all__ = ["get_2d_edges_and_simplices", "get_3d_edges_and_simplices", "prepare_weights"]
 
 
 def get_2d_edges_and_simplices(
@@ -194,3 +195,54 @@ def get_3d_edges_and_simplices(
         ).tolist()
 
     return edges, simplices
+
+
+def prepare_weights(weights: npt.NDArray, edges: npt.NDArray[np.int_], smoothing: float = 0.1,
+                    merging_method: str = 'mean') -> npt.NDArray[np.float_]:
+    """Prepare weights for `calculate_m` and `calculate_k` functions.
+
+    Assume the weights are the same shape as the phases to be unwrapped.
+
+    Scale the weights from 0 to 1. Pick the weights corresponding to the phase pairs connected by the edges.
+    Compute the mean/max/min (depending on the `merging_method`) of each of those pairs to give a weight for each edge.
+
+    Args:
+        weights         :   Array of weights of shapr corresponding to the original phases array shape.
+        edges           :   Edges connecting the phases. Shape: (M, 2), where M is the number of edges.
+        smoothing       :   A positive value in range [0, 1). This is the minimal value of the rescaled weights
+                            where they are defined. If smoothing > 0, the value of 0 is reserved for places where
+                            the weights are originally NaN. If smoothing == 0, 0 will be used for both NaN weights
+                            and smallest non-NaN ones.
+        merging_method  :   Way of combining two phase weights into a single edge weight.
+
+    Returns:
+        Array of weights for the edges, shape: (M,). Rescaled to [0, 1].
+    """
+
+    if not 0 <= smoothing < 1:
+        raise ValueError(
+            "`smoothing` should be a value between 0 (inclusive) and 1 (non inclusive); got " + str(smoothing))
+
+    # scale the weights from 0 to 1
+    weights = weights - np.nanmin(weights)
+    current_max = np.nanmax(weights)
+    if not current_max:
+        # current maximum is 0, which means all weights originally had the same value, now 0; replace everything with 1
+        weights += 1
+    else:
+        weights /= current_max
+        weights *= (1 - smoothing)
+        weights += smoothing
+
+    # pick the weights corresponding to the phases connected by the edges
+    # and use `merging_method` to get one weight for each edge
+    allowed_merging_methods = ['min', 'max', 'mean']
+    if merging_method not in allowed_merging_methods:
+        raise ValueError(
+            "`merging_method` should be one of: " + ', '.join(merging_method) + '; got ' + str(merging_method))
+    weights_for_edges = getattr(np, merging_method)(weights.ravel()[edges], axis=1)
+
+    # make sure there are no NaNs in the weights; replace any with 0s
+    weights_for_edges[np.isnan(weights_for_edges)] = 0
+
+    return weights_for_edges