allow passing rotation center for tilt

import multiprocessing

from imars3d.backend.util.functions import clamp_max_workers

import numpy as np

from typing import Tuple

from typing import Tuple, Union, Optional

from functools import partial

from scipy.optimize import minimize_scalar

from scipy.optimize import OptimizeResult

    tilt: float,

    image0: np.ndarray,

    image1: np.ndarray,

    center: Optional[Tuple[Union[float, int], Union[float, int]]] = None,

) -> float:

    """Calculate the dissimilarity between two images with given tilt.

    image1:

        The second image for comparison, which is often the radiograph taken at

        omega + 180 deg

    center:

        The center of the rotation axis, default is None, which means the center

        of the image. This will be passed to the rotation function from skimage.

    Returns

    -------

        resize=True,

        preserve_range=True,

        order=1,  # use default bi-linear interpolation for rotation

        center=center,

    )

    # since 180 is flipped, tilting back -2 deg of the original img180 means tilting +2 deg

    # of the flipped one

        resize=True,

        preserve_range=True,

        order=1,  # use default bi-linear interpolation for rotation

        center=center,

    )

    # p-norm

    image180: np.ndarray,

    low_bound: float = -5.0,

    high_bound: float = 5.0,

    center: Optional[Tuple[Union[float, int], Union[float, int]]] = None,

) -> OptimizeResult:

    """

    Use optimization to find the in-plane tilt angle.

        The lower bound of the tilt angle search space

    high_bound:

        The upper bound of the tilt angle search space

    center:

        The center of the rotation axis, default is None, which means the center

        of the image. This will be passed to the rotation function from skimage.

    Returns

    -------

        The optimization results from scipy.optimize.minimize_scalar

    """

    # make the error function

    err_func = partial(calculate_dissimilarity, image0=image0, image1=image180)

    err_func = partial(calculate_dissimilarity, image0=image0, image1=image180, center=center)

    # use bounded uni-variable optimizer to locate the tilt angle that minimize

    # the dissimilarity of the 180 deg pair

    res = minimize_scalar(

    cut_off_angle_deg: float

        The angle in degrees to cut off the rotation axis tilt correction, i.e.

        skip applying tilt correction for tilt angles that are too small.

    center: Any

        The center of the rotation axis, default is None, which means the center

        of the image. This will be passed to the rotation function from skimage.

    max_workers:

        Number of cores to use for parallel median filtering, default is 0,

        which means using all available cores.

        default=2.0,

        doc="The angle in degrees to cut off the rotation axis tilt correction, i.e. skip applying tilt correction for tilt angles that are too small.",

    )

    center = param.Parameter(

        default=None,

        doc="The center of the rotation axis, default is None, which means the center of the image. This will be passed to the rotation function from skimage.",

    )

    # NOTE:

    # The front and backend are sharing the same computing unit, therefore we can

    # set a hard cap on the max_workers.

                    calculate_tilt,

                    low_bound=params.low_bound,

                    high_bound=params.high_bound,

                    center=params.center,

                ),

                [shm_arrays[il] for il in idx_lowrange],

                [shm_arrays[ih] for ih in idx_highrange],

            corrected_array = apply_tilt_correction(

                arrays=params.arrays,

                tilt=tilt,

                center=params.center,

                max_workers=self.max_workers,

            )

        return corrected_array

        The array for tilt correction

    tilt: float

        The rotation axis tilt angle in degrees

    center: Any

        The center of the rotation axis, default is None, which means the center

    max_workers: int

        Number of cores to use for parallel median filtering, default is 0, which means using all available cores.

    tqdm_class: panel.widgets.Tqdm

    arrays = param.Array(doc="The array for tilt correction", default=None)

    tilt = param.Number(doc="The rotation axis tilt angle in degrees", default=None)

    center = param.Parameter(

        default=None,

        doc="The center of the rotation axis, default is None, which means the center of the image. This will be passed to the rotation function from skimage.",

    )

    # NOTE:

    # The front and backend are sharing the same computing unit, therefore we can

    # set a hard cap on the max_workers.

        # dimensionality check

        if params.arrays.ndim == 2:

            logger.info(f"2D image detected, applying tilt correction with tilt = {params.tilt:.3f} deg")

            corrected_array = rotate(params.arrays, -params.tilt, resize=False, preserve_range=True)

            corrected_array = rotate(

                params.arrays, -params.tilt, resize=False, preserve_range=True, center=params.center

            )

        elif params.arrays.ndim == 3:

            logger.info(f"3D array detected, applying tilt correction with tilt = {params.tilt:.3f} deg")

            with SharedMemoryManager() as smm:

                if params.tqdm_class:

                    kwargs["tqdm_class"] = params.tqdm_class

                rst = process_map(

                    partial(rotate, angle=-params.tilt, resize=False, preserve_range=True),

                    partial(rotate, angle=-params.tilt, resize=False, preserve_range=True, center=params.center),

                    [shm_arrays[idx] for idx in range(params.arrays.shape[0])],

                    **kwargs,

                )