Prepare weights for ilp unwrapping

[the-mysh]

Hi @yoyololicon, I think I have figured out a way of passing the weights to the ILP unwrapping.

Please review it; in particular, I was wondering if the rescaling approach and calculating the mean to go from weight-per-phase to weight-per edge is correct.

If it's all right, I think it should resolve #6 which I opened earlier.

[the-mysh]

I've now added a choice of methods to compute an edge weight from a pair of phase weights: choose from mean, max, and min.

[yoyolicoris]

prepare_weights looks nice, but unwrap_dimensional should be the one to call it instead of directly calling calculate_k/m. The weights to the latter are edge weights and should have a shape of (M,), not (N,).

I suggest the following:

def unwrap_dimensional(
    x: np.ndarray,
    start_pixel: Optional[Union[Tuple[int, int], Tuple[int, int, int]]] = None,
    use_edgelist: bool = False,
    cyclical_axis: Union[int, Tuple[int, int]] = (),
    merging_method: str = 'mean',
    weights = None,
    **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.
    cyclical_axis : int or (int, int)
        The axis that is cyclical.
        Default to ().
    weights: 2-D or 3-D np.ndarray
        array of weights for each pixel.
    kwargs : dict
        Other arguments besides `weights` 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, cyclical_axis=cyclical_axis
        )
    elif x.ndim == 3:
        edges, simplices = get_3d_edges_and_simplices(
            x.shape, cyclical_axis=cyclical_axis
        )
    else:
        raise ValueError("x must be 2D or 3D")
    psi = x.ravel()

    if weights is not None:
        weights = prepare_weights(weights, edges, merging_method)

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

    return result.reshape(x.shape)

[the-mysh]

Hi @yoyololicon , thanks for your review. I made changes according to your suggestions. I saw though that in case you use puma, the weights argument should not be present - so I opted for extracting it from kwargs and adding the prepared weights only if any were originally defined.
Please let me know what you think!

[yoyolicoris]

Just some minor issue.
Btw, can you attach some unwrapping examples to show how this vertex-based weighting works? Just curious. It's optional, though.

[yoyolicoris]

@the-mysh please format the code using black.
See the failing workflow for details.

[the-mysh]

Hi @yoyololicon , sorry for the radio silence, got distracted by other work projects. I've implemented your change requests, please have a look when you get the chance.

[yoyolicoris]

There are a few lines I would like to change, but overall looks good!
I'll add a few commits to address them, and then we can merge.

[yoyolicoris]

Thank @the-mysh, for the contribution! This will be included in the next release v0.1.2.