impy-array 2.2.0__py3-none-any.whl → 2.2.2__py3-none-any.whl

impy/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.2.0"
+__version__ = "2.2.2"
 __author__ = "Hanjin Liu"
 __email__ = "liuhanjin-sc@g.ecc.u-tokyo.ac.jp"
 
@@ -6,7 +6,7 @@ import logging
 
 from ._const import Const, SetConst, use
 
-from .collections import *
+from .collections import DataList, DataDict
 from .core import *
 from .binder import bind
 from .viewer import gui
@@ -15,19 +15,17 @@ from .arrays import ImgArray, LazyImgArray, BigImgArray, Label  # for typing
 from . import random, io
 from .axes import slicer
 
-r"""
-Inheritance
------------
+# Inheritance
+# -----------
+#
+# AxesMixin -----> LazyImgArray ------> BigImgArray
+#     |
+#     +-----> MetaArray --+--> LabeledArray --+--> ImgArray
+#     |                   |                   |
+# np.ndarray              +--> Label          +--> PhaseArray
+#                         |
+#                         +--> PropArray
 
-AxesMixin -----> LazyImgArray ------> BigImgArray
-    |
-    +-----> MetaArray --+--> LabeledArray --+--> ImgArray
-    |                   |                   |
-np.ndarray              +--> Label          +--> PhaseArray
-                        |
-                        +--> PropArray
-
-"""
 
 logging.getLogger("skimage").setLevel(logging.ERROR)
 logging.getLogger("tifffile").setLevel(logging.ERROR)

impy/array_api.py

@@ -72,11 +72,13 @@ class XP:
         self.round = np.round
         self.gradient = np.gradient
         self.tensordot = np.tensordot
+        self.concatenate = np.concatenate
         self.stack = np.stack
         self.unravel_index = np.unravel_index
         self.argmax = np.argmax
         self.argmin = np.argmin
         self.pad = np.pad
+        self.isnan = np.isnan
         
         self.state = "numpy"
         from ._const import Const
@@ -141,11 +143,13 @@ class XP:
         except AttributeError:
             self.gradient = _gradient
         self.tensordot = cp.tensordot
+        self.concatenate = cp.concatenate
         self.stack = cp.stack
         self.unravel_index = cp.unravel_index
         self.argmax = cp.argmax
         self.argmin = cp.argmin
         self.pad = cp.pad
+        self.isnan = cp.isnan
         self.state = "cupy"
         
         from ._const import Const

impy/arrays/__init__.py

@@ -16,3 +16,13 @@ try:
 
 except Exception:
     pass
+
+__all__ = [
+    "ImgArray",
+    "LabeledArray",
+    "Label",
+    "PhaseArray",
+    "PropArray",
+    "LazyImgArray",
+    "BigImgArray",
+]

impy/arrays/_utils/_deconv.py

@@ -1,7 +1,13 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 from functools import partial
 import numpy as np
-from ...array_api import xp
+from impy.array_api import xp
 
+if TYPE_CHECKING:
+    from impy.arrays import ImgArray
+    
 try:
     gradient = xp.gradient
 except AttributeError:
@@ -11,11 +17,7 @@ except AttributeError:
         out = numpy.gradient(a.get(), axis=axis)
         return xp.asarray(out)
 
-__all__ = ["wiener", 
-           "lucy",
-           "lucy_tv",
-           "check_psf",
-           ]
+__all__ = ["wiener", "lucy", "lucy_tv", "check_psf"]
 
 def wiener(obs, psf_ft, psf_ft_conj, lmd):
     obs = xp.asarray(obs)
@@ -67,20 +69,23 @@ def richardson_lucy_tv(obs, psf_ft, psf_ft_conj, max_iter, lmd, tol, eps):
     return xp.fft.fftshift(est_new)
 
 
-def _safe_div(a, b, eps=1e-8):
+def _safe_div(a: np.ndarray, b: np.ndarray, eps: float = 1e-8):
     out = xp.zeros(a.shape, dtype=np.float32)
     mask = b > eps
-    out[mask] = a[mask]/b[mask]
+    out[mask] = a[mask] / b[mask]
     return out
 
 
-def check_psf(img, psf, dims):
+def check_psf(img_shape: tuple[int, ...], img_scale: tuple[float, ...], psf):
+    if callable(psf):
+        psf = psf(img_shape, img_scale)
     psf = xp.asarray(psf, dtype=np.float32)
     psf /= xp.sum(psf)
     
-    if img.sizesof(dims) != psf.shape:
-        raise ValueError("observation and PSF have different shape: "
-                        f"{img.sizesof(dims)} and {psf.shape}")
+    if img_shape != psf.shape:
+        raise ValueError(
+            f"observation and PSF have different shape: {img_shape} and {psf.shape}"
+        )
     psf_ft = xp.fft.rfftn(psf)
     psf_ft_conj = xp.conjugate(psf_ft)
-    return psf_ft, psf_ft_conj
+    return psf_ft, psf_ft_conj

impy/arrays/_utils/_linalg.py

@@ -3,11 +3,11 @@ import numpy as np
 from typing import TYPE_CHECKING
 from skimage.feature.corner import _symmetric_image
 from ._skimage import skfeat
-from ...array_api import xp
+from impy.array_api import xp
 
 if TYPE_CHECKING:
-    from ..imgarray import ImgArray
-    from ...axes import Axes
+    from impy.arrays.imgarray import ImgArray
+    from impy.axes import Axes
 
 def eigh(a: np.ndarray):
     a = xp.asarray(a, dtype=a.dtype)
@@ -20,33 +20,35 @@ def eigvalsh(a: np.ndarray):
     return val
 
 def structure_tensor_eigval(img: np.ndarray, sigma: float, pxsize: float):
-    tensor_elements = skfeat.structure_tensor(img, sigma, order="xy",
-                                              mode="reflect")
+    tensor_elements = skfeat.structure_tensor(
+        img, sigma, mode="reflect"
+    )
     return _solve_hermitian_eigval(tensor_elements, pxsize)
 
 def structure_tensor_eigh(img: np.ndarray, sigma: float, pxsize: float):
-    tensor_elements = skfeat.structure_tensor(img, sigma, order="xy",
-                                              mode="reflect")
+    tensor_elements = skfeat.structure_tensor(
+        img, sigma, mode="reflect"
+    )
     return _solve_hermitian_eigs(tensor_elements, pxsize)
 
 def hessian_eigval(img: np.ndarray, sigma: float, pxsize: float):
     hessian_elements = skfeat.hessian_matrix(
-        img, sigma=sigma, order="xy", mode="reflect"
+        img, sigma=sigma, mode="reflect", use_gaussian_derivatives=False,
     )
     return _solve_hermitian_eigval(hessian_elements, pxsize)
 
 def hessian_eigh(img: np.ndarray, sigma: float, pxsize: float):
     hessian_elements = skfeat.hessian_matrix(
-        img, sigma=sigma, order="xy", mode="reflect"
+        img, sigma=sigma, mode="reflect", use_gaussian_derivatives=False,
     )
     return _solve_hermitian_eigs(hessian_elements, pxsize)
 
 
-def _solve_hermitian_eigval(elems, pxsize):
+def _solve_hermitian_eigval(elems, pxsize) -> np.ndarray:
     # Correct for scale
     pxsize = np.asarray(pxsize)
     hessian = _symmetric_image(elems)
-    hessian *= (pxsize.reshape(-1,1) * pxsize.reshape(1,-1))
+    hessian *= (pxsize.reshape(-1, 1) * pxsize.reshape(1, -1))
     eigval = eigvalsh(hessian)
     return eigval
 

impy/arrays/_utils/_transform.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import Iterable, NamedTuple, Sequence, TYPE_CHECKING
+import math
 import numpy as np
 
 from ._skimage import sktrans
@@ -149,7 +150,7 @@ def check_matrix(matrices: list[np.ndarray | float]):
 def polar2d(
     img: xp.ndarray,
     rmax: int,
-    dtheta: float = 0.1,
+    dtheta: float = 0.1,  # radian
     order: int = 1,
     mode: str = "constant",
     cval: float = 0,
@@ -307,8 +308,8 @@ def iradon(
     # Reconstruct image by interpolation
     reconstructed = np.zeros(output_shape, dtype=dtype)
     xpr, ypr = np.indices(output_shape)
-    xpr -= output_shape[0] // 2
-    ypr -= output_shape[1] // 2
+    xpr = xpr - output_shape[0] // 2
+    ypr = ypr - output_shape[1] // 2
     
     from scipy.interpolate import interp1d
     x = np.arange(img_shape) - img_shape // 2  # NOTE: use CPU!
@@ -322,12 +323,13 @@ def iradon(
 # This function is almost ported from `skimage.transform`.
 def get_fourier_filter(size: int, filter_name: str):
     n = xp.concatenate(
-        [xp.arange(1, size / 2 + 1, 2, dtype=int),
-         xp.arange(size / 2 - 1, 0, -2, dtype=int)]
+        [xp.arange(1, size // 2 + 1, 2, dtype=np.float32),
+         xp.arange(size // 2 - 1, 0, -2, dtype=np.float32)]
     )
     f = xp.zeros(size)
     f[0] = 0.25
     f[1::2] = -1 / (np.pi * n[:len(f[1::2])]) ** 2
+
     fourier_filter = 2 * xp.real(xp.fft.fft(f))  # ramp filter
     if filter_name == "ramp":
         pass

impy/arrays/axesmixin.py

@@ -200,6 +200,9 @@ class AxesMixin:
                 {a: s for a, s in other.scale.items() if a in self.axes}  # type: ignore
             )
         
+        elif isinstance(other, (list, tuple)):
+            for a, s in zip(self.axes, other):
+                a.scale = s
         else:
             raise TypeError(
                 f"'other' must be str or axes supported object, but got {type(other)}"

impy/arrays/bases/metaarray.py

@@ -27,7 +27,7 @@ SupportSlicing = Union[
     Mapping[str, SupportOneSlicing],
 ]
 
-class MetaArray(AxesMixin, np.ndarray):
+class MetaArray(AxesMixin, np.ndarray[Any, np.dtype[np.number]]):
     additional_props = ["_source", "_metadata", "_name"]
     NP_DISPATCH = {}
     _name: str
@@ -747,7 +747,7 @@ class MetaArray(AxesMixin, np.ndarray):
         except ImageAxesError:
             self.axes = None
         
-        return None
+        return self
     
     def _process_output(self, func, args, kwargs):
         # find the largest MetaArray. Largest because of broadcasting.

impy/arrays/bases/overload.py

@@ -181,6 +181,29 @@ def _(img: MetaArray, axis1: int | AxisLike, axis2: int | AxisLike):
     out._set_info(img, new_axes=axes_list)
     return out
 
+# argmax does not support axis as tuple
+@MetaArray.implements(np.argmax)
+def _(img: MetaArray, axis: int | AxisLike | None = None):
+    if isinstance(axis, (str, Axis)):
+        axis = img.axisof(axis)
+    new_axes = img.axes.drop(axis)
+    out = np.argmax(img.value, axis=axis)
+    if isinstance(out, np.ndarray):
+        out = out.view(img.__class__)
+        out._set_info(img, new_axes=new_axes)
+    return out
+
+@MetaArray.implements(np.argmin)
+def _(img: MetaArray, axis: int | AxisLike | None = None):
+    if isinstance(axis, (str, Axis)):
+        axis = img.axisof(axis)
+    new_axes = img.axes.drop(axis)
+    out = np.argmin(img.value, axis=axis)
+    if isinstance(out, np.ndarray):
+        out = out.view(img.__class__)
+        out._set_info(img, new_axes=new_axes)
+    return out
+
 @MetaArray.implements(np.cross)
 def _(
     img: MetaArray,
@@ -206,6 +229,41 @@ def _(
     out._set_info(img, new_axes=img.axes)
     return out
 
+@MetaArray.implements(np.diff)
+def _(
+    img: MetaArray,
+    n: int = 1,
+    axis: int | AxisLike = -1,
+    prepend=None,
+    append=None,
+):
+    if isinstance(axis, (str, Axis)):
+        axis = img.axisof(axis)
+    kwargs = dict(n=n, axis=axis)
+    if prepend is not None:
+        kwargs["prepend"] = prepend
+    if append is not None:
+        kwargs["append"] = append
+    out = np.diff(img.value, **kwargs)
+    out = out.view(img.__class__)
+    out._set_info(img, new_axes=img.axes)
+    return out
+
+@MetaArray.implements(np.gradient)
+def _(img: MetaArray, *varargs, axis=None, edge_order=1):
+    if isinstance(axis, (str, Axis)):
+        axis = img.axisof(axis)
+    elif isinstance(axis, Axes):
+        axis = tuple(img.axisof(ax) for ax in axis)
+    out = np.gradient(img.value, *varargs, axis=axis, edge_order=edge_order)
+    if isinstance(out, np.ndarray):
+        out = out.view(img.__class__)._set_info(img, new_axes=img.axes)
+    else:
+        out = [arr.view(img.__class__)._set_info(img, new_axes=img.axes) for arr in out]
+        out = DataList(out)
+    return out
+
+
 # This function is ported from numpy.core.numeric.normalize_axis_tuple
 def np_normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False):
     # Optimization to speed-up the most common cases.

impy/arrays/imgarray.py

@@ -13,25 +13,24 @@ from .specials import PropArray
 from ._utils._skimage import skexp, skfeat, skfil, skimage, skmes, skres, skseg, sktrans
 from ._utils import _filters, _linalg, _deconv, _misc, _glcm, _docs, _transform, _structures, _corr
 
-from ..utils.axesop import add_axes, switch_slice, complement_axes, find_first_appeared
-from ..utils.deco import check_input_and_output, dims_to_spatial_axes, same_dtype
-from ..utils.gauss import GaussianBackground, GaussianParticle
-from ..utils.misc import check_nd, largest_zeros
-from ..utils.slicer import solve_slicer
-
-from ..collections import DataDict
-from ..axes import AxisLike, slicer, Axes, Axis
-from .._types import nDInt, nDFloat, Dims, Coords, AxesTargetedSlicer, PaddingMode
-from .._const import Const
-from ..array_api import xp, cupy_dispatcher
+from impy.utils.axesop import add_axes, switch_slice, complement_axes, find_first_appeared
+from impy.utils.deco import check_input_and_output, dims_to_spatial_axes, same_dtype
+from impy.utils.gauss import GaussianBackground, GaussianParticle
+from impy.utils.misc import check_nd, largest_zeros
+from impy.utils.slicer import solve_slicer
+
+from impy.collections import DataDict
+from impy.axes import AxisLike, slicer, Axes, Axis
+from impy._types import nDInt, nDFloat, Dims, Coords, AxesTargetedSlicer, PaddingMode
+from impy._const import Const
+from impy.array_api import xp, cupy_dispatcher
 
 if TYPE_CHECKING:
     from ..frame import MarkerFrame
     from typing import Literal, Union
-    ThreasholdMethod = Union[
-        Literal["isodata"], Literal["li"], Literal["local"], Literal["mean"], Literal["min"],
-        Literal["minimum"], Literal["niblack"], Literal["otsu"], Literal["sauvola"], 
-        Literal["triangle"], Literal["yen"]
+    ThreasholdMethod = Literal[
+        "isodata", "li", "local", "mean", "min", "minimum", "niblack", "otsu", "sauvola",
+        "triangle", "yen"
     ]
     FftShape = Literal["same", "square"]
 
@@ -220,7 +219,7 @@ class ImgArray(LabeledArray):
     @same_dtype(asfloat=True)
     def rescale(
         self,
-        scale: float = 1/16,
+        scale: nDFloat = 1/16,
         *,
         order: int = None,
         mode: PaddingMode = "reflect",
@@ -228,31 +227,45 @@ class ImgArray(LabeledArray):
         dims: Dims = None,
     ) -> ImgArray:
         """
-        Rescale image.
+        Rescale image by applying the diagonal component of Affine matrix.
 
         Parameters
         ----------
         scale : float, optional
-            scale of the new image.
+            Relative scale of the new image. scale 1/2 means that the shape of
+            the output image will be (N/2, ...).
         {order}{mode}{cval}{dims}
         
         Returns
         -------
         ImgArray
             Rescaled image.
-        """        
-        outshape = switch_slice(
-            dims, self.axes, ifin=np.array(self.shape)*scale, ifnot=self.shape
-        )
-        gb = np.prod(outshape) * 4/1e9
+        """
+        scale_is_seq = hasattr(scale, "__iter__")
+        
+        # Check if output is too large.
+        gb = -1
+        if not scale_is_seq and scale > 1:
+            gb = np.prod(self.shape) * (scale ** len(dims)) / 1e9
+        elif scale_is_seq and np.prod(list(scale)) > 1:
+            gb = np.prod(self.shape) * np.prod(list(scale)) / 1e9
         if gb > Const["MAX_GB"]:
-            raise MemoryError(f"Too large: {gb} GB")
-        scale_ = [scale if a in dims else 1 for a in self.axes]
+            raise MemoryError(f"Output image is too large: {gb} GB")
+        
+        if scale_is_seq:
+            scale = list(scale)
+            if len(scale) != len(dims):
+                raise ValueError(
+                    "scale must have the same length as the spatial dimensions."
+                )
+            it = iter(scale)
+            scale_ = [next(it) if a in dims else 1 for a in self.axes]
+        else:
+            scale_ = [scale if a in dims else 1 for a in self.axes]
         out = sktrans.rescale(
             self.value, scale_, order=order, mode=mode, cval=cval, anti_aliasing=False
         )
-        out: ImgArray = out.view(self.__class__)
-        out._set_info(self)
+        out = out.view(self.__class__)._set_info(self)
         out.axes = self.axes.copy()
         out.set_scale({a: self.scale[a]/scale for a, scale in zip(self.axes, scale_)})
         return out
@@ -303,11 +316,11 @@ class ImgArray(LabeledArray):
         )
         
         reshaped_img = img_to_reshape.reshape(shape)
-        axes_to_reduce = tuple(i*2+1 for i in range(self.ndim))
+        axes_to_reduce = tuple(i * 2 + 1 for i in range(self.ndim))
         out: np.ndarray = binfunc(reshaped_img, axis=axes_to_reduce)
         out: ImgArray = out.view(self.__class__)
         out._set_info(self)
-        out.axes = str(self.axes) # _set_info does not pass copy so new axes must be defined here.
+        out.axes = self.axes.copy()  # _set_info does not pass copy so new axes must be defined here.
         out.set_scale(
             {a: self.scale[a]/scale for a, scale in zip(self.axes, scale_)}
         )
@@ -318,7 +331,7 @@ class ImgArray(LabeledArray):
     def radial_profile(
         self,
         nbin: int = 32,
-        center: Iterable[float] = None,
+        center: Iterable[float] | None = None,
         r_max: float = None, 
         *, 
         method: str = "mean",
@@ -3111,8 +3124,7 @@ class ImgArray(LabeledArray):
         if out.ndim == 3:
             out = np.moveaxis(out, 0, 1)
         out = out[::-1]
-        out._set_info(self, new_axes)
-        return out
+        return out._set_info(self, new_axes)
 
     @check_input_and_output
     def threshold(
@@ -4195,7 +4207,7 @@ class ImgArray(LabeledArray):
     @check_input_and_output
     def wiener(
         self,
-        psf: np.ndarray,
+        psf: np.ndarray | Callable[[tuple[int, ...]], np.ndarray],
         lmd: float = 0.1,
         *, 
         dims: Dims = None, 
@@ -4213,8 +4225,9 @@ class ImgArray(LabeledArray):
 
         Parameters
         ----------
-        psf : np.ndarray
-            Point spread function.
+        psf : ndarray or callable
+            Point spread function. If a function is given, `psf(shape)` will be
+            called to generate the PSF.
         lmd : float, default is 0.1
             Constant value used in the deconvolution. See Formulation below.
         {dims}{update}
@@ -4233,22 +4246,23 @@ class ImgArray(LabeledArray):
         if lmd <= 0:
             raise ValueError(f"lmd must be positive, but got: {lmd}")
         
-        psf_ft, psf_ft_conj = _deconv.check_psf(self, psf, dims)
+        psf_ft, psf_ft_conj = _deconv.check_psf(
+            self.sizesof(dims), tuple(self.scale[axis] for axis in dims), psf,
+        )
         
         return self._apply_dask(
             _deconv.wiener, 
             c_axes=complement_axes(dims, self.axes),
             args=(psf_ft, psf_ft_conj, lmd)
-        )
-        
-    
+        )    
+
     @_docs.write_docs
     @dims_to_spatial_axes
     @same_dtype(asfloat=True)
     @check_input_and_output
     def lucy(
         self,
-        psf: np.ndarray,
+        psf: np.ndarray | Callable[[tuple[int, ...]], np.ndarray],
         niter: int = 50,
         eps: float = 1e-5,
         *, 
@@ -4260,14 +4274,15 @@ class ImgArray(LabeledArray):
         
         Parameters
         ----------
-        psf : np.ndarray
-            Point spread function.
+        psf : ndarray or callable
+            Point spread function. If a function is given, `psf(shape)` will be
+            called to generate the PSF.
         niter : int, default is 50.
             Number of iterations.
         eps : float, default is 1e-5
-            During deconvolution, division by small values in the convolve image of estimation and 
-            PSF may cause divergence. Therefore, division by values under `eps` is substituted
-            to zero.
+            During deconvolution, division by small values in the convolve image 
+            of estimation and PSF may cause divergence. Therefore, division by 
+            values under `eps` is substituted to zero.
         {dims}{update}
         
         Returns
@@ -4281,7 +4296,9 @@ class ImgArray(LabeledArray):
         wiener
         """
         
-        psf_ft, psf_ft_conj = _deconv.check_psf(self, psf, dims)
+        psf_ft, psf_ft_conj = _deconv.check_psf(
+            self.sizesof(dims), tuple(self.scale[axis] for axis in dims), psf,
+        )
 
         return self._apply_dask(
             _deconv.richardson_lucy, 
@@ -4295,7 +4312,7 @@ class ImgArray(LabeledArray):
     @check_input_and_output
     def lucy_tv(
         self,
-        psf: np.ndarray,
+        psf: np.ndarray | Callable[[tuple[int, ...]], np.ndarray],
         max_iter: int = 50,
         lmd: float = 1e-3,
         tol: float = 1e-3,
@@ -4319,8 +4336,9 @@ class ImgArray(LabeledArray):
         
         Parameters
         ----------
-        psf : np.ndarray
-            Point spread function.
+        psf : ndarray or callable
+            Point spread function. If a function is given, `psf(shape)` will be
+            called to generate the PSF.
         max_iter : int, default is 50.
             Maximum number of iterations.
         lmd : float, default is 1e-3
@@ -4361,7 +4379,9 @@ class ImgArray(LabeledArray):
                 "In Richadson-Lucy with total-variance-regularization, "
                 "parameter `lmd` must be positive."
             )
-        psf_ft, psf_ft_conj = _deconv.check_psf(self, psf, dims)
+        psf_ft, psf_ft_conj = _deconv.check_psf(
+            self.sizesof(dims), tuple(self.scale[axis] for axis in dims), psf
+        )
 
         return self._apply_dask(
             _deconv.richardson_lucy_tv, 
@@ -4369,6 +4389,39 @@ class ImgArray(LabeledArray):
             args=(psf_ft, psf_ft_conj, max_iter, lmd, tol, eps)
         )
 
+    @_docs.write_docs
+    @dims_to_spatial_axes
+    @same_dtype(asfloat=True)
+    @check_input_and_output
+    def tiled_lucy(
+        self,
+        psf: np.ndarray | Callable[[tuple[int, ...]], np.ndarray],
+        niter: int = 50,
+        eps: float = 1e-5,
+        chunks = "auto",
+        overlap: int = 32,
+        *, 
+        dims: Dims = None, 
+        update: bool = False
+    ) -> ImgArray:
+        from dask import array as da
+
+        input = da.from_array(self.value, chunks=chunks)
+        depth = switch_slice(dims, self.axes, overlap, 0)
+        scale = tuple(self.scale[axis] for axis in dims)
+        
+        def func(arr: np.ndarray):
+            psf_ft, psf_ft_conj = _deconv.check_psf(arr.shape, scale, psf)
+            return _deconv.richardson_lucy(arr, psf_ft, psf_ft_conj, niter, eps)
+            
+        out = da.map_overlap(
+            func, 
+            input,
+            depth=depth,
+            boundary="reflect",
+            dtype=self.dtype,
+        ).compute()
+        return out
 
 def _check_coordinates(coords, img: ImgArray, dims: Dims = None):
     from impy.frame import MarkerFrame