From 45fcf2e924ade756f369872415fd8639d546294e Mon Sep 17 00:00:00 2001
From: hanjinliu <liuhanjin-sc@g.ecc.u-tokyo.ac.jp>
Date: Thu, 24 Feb 2022 21:12:50 +0900
Subject: [PATCH 01/13] fix GPU

---
 impy/arrays/_utils/_corr.py      |  2 +-
 impy/arrays/_utils/_transform.py | 19 ++++++++++---------
 impy/correlation.py              |  2 +-
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/impy/arrays/_utils/_corr.py b/impy/arrays/_utils/_corr.py
index 243b166f..3ebc379e 100644
--- a/impy/arrays/_utils/_corr.py
+++ b/impy/arrays/_utils/_corr.py
@@ -15,7 +15,7 @@ def subpixel_pcc(
     cross_correlation = xp_fft.ifftn(product)
     power = abs2(cross_correlation)
     if max_shifts is not None:
-        max_shifts = np.asarray(max_shifts)
+        max_shifts = xp.asarray(max_shifts)
         power = crop_by_max_shifts(power, max_shifts, max_shifts)
         
     maxima = xp.unravel_index(xp.argmax(power), power.shape)
diff --git a/impy/arrays/_utils/_transform.py b/impy/arrays/_utils/_transform.py
index c578356c..a87ec580 100644
--- a/impy/arrays/_utils/_transform.py
+++ b/impy/arrays/_utils/_transform.py
@@ -1,5 +1,6 @@
 import numpy as np
 import scipy
+from scipy import ndimage as ndi
 from collections import namedtuple
 from ._skimage import sktrans
 from ..._cupy import xp, xp_ndi, xp_ndarray
@@ -123,7 +124,7 @@ def check_matrix(matrices):
     return mtx
 
 def polar2d(
-    img: xp_ndarray,
+    img: np.ndarray,
     rmax: int,
     dtheta: float = 0.1,
     order=1,
@@ -131,14 +132,14 @@ def polar2d(
     cval=0
 ) -> np.ndarray:
     centers = np.array(img.shape)/2 - 0.5
-    r = xp.arange(rmax) + 0.5
-    theta = xp.arange(0, 2*np.pi, dtheta)
-    r, theta = xp.meshgrid(r, theta)
-    y = r * xp.sin(theta) + centers[0]
-    x = r * xp.cos(theta) + centers[1]
-    coords = xp.stack([y, x], axis=0)
-    img = xp.asarray(img, dtype=img.dtype)
-    out = xp_ndi.map_coordinates(img, coords, order=order, mode=mode, cval=cval, prefilter=order>1)
+    r = np.arange(rmax) + 0.5
+    theta = np.arange(0, 2*np.pi, dtheta)
+    r, theta = np.meshgrid(r, theta)
+    y = r * np.sin(theta) + centers[0]
+    x = r * np.cos(theta) + centers[1]
+    coords = np.stack([y, x], axis=0)
+    img = np.asarray(img, dtype=img.dtype)
+    out = ndi.map_coordinates(img, coords, order=order, mode=mode, cval=cval, prefilter=order>1)
     return out
 
 # def polar3d(
diff --git a/impy/correlation.py b/impy/correlation.py
index ce4f9cb7..5f575c1c 100644
--- a/impy/correlation.py
+++ b/impy/correlation.py
@@ -78,7 +78,7 @@ def radial_sum(arr):
     
         out = radial_sum(cov)/xp.sqrt(radial_sum(pw0) * radial_sum(pw1))
     freq = (np.arange(len(out)) + 0.5) * dfreq
-    return freq, out
+    return freq, asnumpy(out)
 
 # alias
 fourier_shell_correlation = fsc

From 1622fe1842a264349910555c6741b2ecfcb8f779 Mon Sep 17 00:00:00 2001
From: hanjinliu <hanjin.liu@bs.s.u-tokyo.ac.jp>
Date: Fri, 25 Feb 2022 01:41:41 +0900
Subject: [PATCH 02/13] refactor cupy

---
 impy/_cupy.py                     | 93 ++++++++++++++++++++-----------
 impy/arrays/_utils/_corr.py       | 24 ++++----
 impy/arrays/_utils/_deconv.py     | 22 ++++----
 impy/arrays/_utils/_filters.py    | 10 ++--
 impy/arrays/_utils/_linalg.py     |  6 +-
 impy/arrays/_utils/_structures.py |  1 -
 impy/arrays/_utils/_transform.py  | 18 +++---
 impy/arrays/imgarray.py           | 76 ++++++++++++++-----------
 impy/arrays/lazy.py               | 34 +++++------
 impy/correlation.py               | 14 ++---
 tests/test_fft.py                 |  6 +-
 11 files changed, 170 insertions(+), 134 deletions(-)

diff --git a/impy/_cupy.py b/impy/_cupy.py
index 1f3911f0..bfa4c888 100644
--- a/impy/_cupy.py
+++ b/impy/_cupy.py
@@ -1,40 +1,69 @@
-try:
-    import cupy as xp
-    GPU_AVAILABLE = True
-except ImportError:
-    import numpy as xp
-    GPU_AVAILABLE = False
-
-if GPU_AVAILABLE:
-    def asnumpy(arr, dtype=None):
-        out = xp.asnumpy(arr)
-        if dtype is None:
-            return out
-        return out.astype(dtype)
-    from cupyx.scipy import fft as xp_fft
-    from cupyx.scipy import ndimage as xp_ndi
-    from cupy import linalg as xp_linalg
-    from cupy import ndarray as xp_ndarray
-    import numpy as np
-    _as_cupy = lambda a: xp.asarray(a) if isinstance(a, np.ndarray) else a
-    
-else:
-    asnumpy = xp.asarray
-    try:
-        from scipy import fft as xp_fft
-    except ImportError:
-        from scipy import fftpack as xp_fft
-    from scipy import ndimage as xp_ndi
-    from numpy import linalg as xp_linalg
-    from numpy import ndarray as xp_ndarray
-    _as_cupy = lambda a: a
 
 from functools import wraps
+import numpy as np
+from numpy.typing import ArrayLike, DTypeLike
 
 def cupy_dispatcher(function):
     @wraps(function)
     def func(*args, **kwargs):
-        args = map(_as_cupy, args)
+        if xp.state == "cupy":
+            args = (xp.asarray(a) if isinstance(a, np.ndarray) else a for a in args)
         out = function(*args, **kwargs)
         return out
-    return func
\ No newline at end of file
+    return func
+
+from types import ModuleType
+from scipy import ndimage as scipy_ndi
+from typing import Callable
+
+class XP:
+    fft: ModuleType
+    linalg: ModuleType
+    random: ModuleType
+    ndi: ModuleType
+    asnumpy: Callable[[ArrayLike, DTypeLike], np.ndarray]
+    asarray: Callable[[ArrayLike], ArrayLike]
+    ndarray: type
+    state: str
+    
+    def __init__(self):
+        self.setNumpy()
+    
+    def __getattr__(self, key: str):
+        return getattr(self._module, key)
+    
+    def setNumpy(self) -> None:
+        self._module = np
+        self.fft = np.fft
+        self.linalg = np.linalg
+        self.random = np.random
+        self.ndi = scipy_ndi
+        self.asnumpy = np.asarray
+        self.asarray = np.asarray
+        self.ndarray = np.ndarray
+        self.state = "numpy"
+    
+    def setCupy(self) -> None:
+        import cupy as cp
+        def cp_asnumpy(arr, dtype=None):
+            out = cp.asnumpy(arr)
+            if dtype is None:
+                return out
+            return out.astype(dtype)
+        from cupyx.scipy import fft as cp_fft
+        from cupyx.scipy import ndimage as cp_ndi
+        from cupy import linalg as cp_linalg
+        from cupy import ndarray as cp_ndarray
+        
+        self._module = cp
+        self.fft = cp_fft
+        self.linalg = cp_linalg
+        self.random = cp.random
+        self.ndi = cp_ndi
+        self.asnumpy = cp_asnumpy
+        self.asarray = cp.asarray
+        self.ndarray = cp_ndarray
+        self.state = "cupy"
+
+xp = XP()
+    
\ No newline at end of file
diff --git a/impy/arrays/_utils/_corr.py b/impy/arrays/_utils/_corr.py
index 243b166f..b849acc1 100644
--- a/impy/arrays/_utils/_corr.py
+++ b/impy/arrays/_utils/_corr.py
@@ -1,18 +1,18 @@
 from __future__ import annotations
 import numpy as np
-from ..._cupy import xp, xp_fft, xp_ndarray
+from ..._cupy import xp
 
 # Modified from skimage/registration/_phase_cross_correlation.py
 # Compatible between numpy/cupy and supports maximum shifts.
 
 def subpixel_pcc(
-    f0: xp_ndarray, 
-    f1: xp_ndarray, 
+    f0: xp.ndarray, 
+    f1: xp.ndarray, 
     upsample_factor: int,
     max_shifts: tuple[float, ...] | None = None,
-) -> xp_ndarray:
+) -> xp.ndarray:
     product = f0 * f1.conj()
-    cross_correlation = xp_fft.ifftn(product)
+    cross_correlation = xp.fft.ifftn(product)
     power = abs2(cross_correlation)
     if max_shifts is not None:
         max_shifts = np.asarray(max_shifts)
@@ -52,11 +52,11 @@ def subpixel_pcc(
     return shifts
 
 def _upsampled_dft(
-    data: xp_ndarray, 
+    data: xp.ndarray, 
     upsampled_region_size: np.ndarray, 
     upsample_factor: int, 
-    axis_offsets: xp_ndarray
-) -> xp_ndarray:
+    axis_offsets: xp.ndarray
+) -> xp.ndarray:
     # if people pass in an integer, expand it to a list of equal-sized sections
     upsampled_region_size = [upsampled_region_size,] * data.ndim
 
@@ -70,14 +70,14 @@ def _upsampled_dft(
         data = xp.tensordot(kernel, data, axes=(1, -1))
     return data
 
-def abs2(a: xp_ndarray) -> xp_ndarray:
+def abs2(a: xp.ndarray) -> xp.ndarray:
     return a.real**2 + a.imag**2
 
-def crop_by_max_shifts(power: xp_ndarray, left, right):
-    shifted_power = xp_fft.fftshift(power)
+def crop_by_max_shifts(power: xp.ndarray, left, right):
+    shifted_power = xp.fft.fftshift(power)
     centers = tuple(s//2 for s in power.shape)
     slices = tuple(
         slice(max(c - int(shiftl), 0), min(c + int(shiftr) + 1, s), None) 
         for c, shiftl, shiftr, s in zip(centers, left, right, power.shape)
     )
-    return xp_fft.ifftshift(shifted_power[slices])
\ No newline at end of file
+    return xp.fft.ifftshift(shifted_power[slices])
\ No newline at end of file
diff --git a/impy/arrays/_utils/_deconv.py b/impy/arrays/_utils/_deconv.py
index 3ecbe2bf..42f77215 100644
--- a/impy/arrays/_utils/_deconv.py
+++ b/impy/arrays/_utils/_deconv.py
@@ -1,5 +1,5 @@
 from functools import partial
-from ..._cupy import xp, xp_fft
+from ..._cupy import xp
 
 try:
     gradient = xp.gradient
@@ -18,19 +18,19 @@ def gradient(a, axis=None):
 
 def wiener(obs, psf_ft, psf_ft_conj, lmd):
     obs = xp.asarray(obs)
-    fft = xp_fft.rfftn
-    ifft = partial(xp_fft.irfftn, s=obs.shape)
+    fft = xp.fft.rfftn
+    ifft = partial(xp.fft.irfftn, s=obs.shape)
     
     img_ft = fft(obs)
     
     estimated = xp.real(ifft(img_ft*psf_ft_conj / (psf_ft*psf_ft_conj + lmd)))
-    return xp_fft.fftshift(estimated)
+    return xp.fft.fftshift(estimated)
     
 def richardson_lucy(obs, psf_ft, psf_ft_conj, niter, eps):
     # Identical to the algorithm in Deconvolution.jl of Julia.
     obs = xp.asarray(obs)
-    fft = xp_fft.rfftn
-    ifft = partial(xp_fft.irfftn, s=obs.shape)
+    fft = xp.fft.rfftn
+    ifft = partial(xp.fft.irfftn, s=obs.shape)
     conv = factor = xp.empty(obs.shape, dtype=xp.float32) # placeholder
     estimated = xp.real(ifft(fft(obs) * psf_ft))   # initialization
     
@@ -39,12 +39,12 @@ def richardson_lucy(obs, psf_ft, psf_ft_conj, niter, eps):
         factor[:] = ifft(fft(_safe_div(obs, conv, eps=eps)) * psf_ft_conj).real
         estimated *= factor
         
-    return xp_fft.fftshift(estimated)
+    return xp.fft.fftshift(estimated)
 
 def richardson_lucy_tv(obs, psf_ft, psf_ft_conj, max_iter, lmd, tol, eps):
     obs = xp.asarray(obs)
-    fft = xp_fft.rfftn
-    ifft = partial(xp_fft.irfftn, s=obs.shape)
+    fft = xp.fft.rfftn
+    ifft = partial(xp.fft.irfftn, s=obs.shape)
     est_old = ifft(fft(obs) * psf_ft).real
     est_new = xp.empty(obs.shape, dtype=xp.float32)
     conv = factor = norm = gg = xp.empty(obs.shape, dtype=xp.float32) # placeholder
@@ -63,7 +63,7 @@ def richardson_lucy_tv(obs, psf_ft, psf_ft_conj, max_iter, lmd, tol, eps):
             break
         est_old[:] = est_new
         
-    return xp_fft.fftshift(est_new)
+    return xp.fft.fftshift(est_new)
 
 
 def _safe_div(a, b, eps=1e-8):
@@ -80,6 +80,6 @@ def check_psf(img, psf, dims):
     if img.sizesof(dims) != psf.shape:
         raise ValueError("observation and PSF have different shape: "
                         f"{img.sizesof(dims)} and {psf.shape}")
-    psf_ft = xp_fft.rfftn(psf)
+    psf_ft = xp.fft.rfftn(psf)
     psf_ft_conj = xp.conjugate(psf_ft)
     return psf_ft, psf_ft_conj
\ No newline at end of file
diff --git a/impy/arrays/_utils/_filters.py b/impy/arrays/_utils/_filters.py
index 9a3f1e95..76a9222c 100644
--- a/impy/arrays/_utils/_filters.py
+++ b/impy/arrays/_utils/_filters.py
@@ -30,12 +30,12 @@
            ]
 
 
-from ..._cupy import xp_ndi, xp, asnumpy, cupy_dispatcher
+from ..._cupy import xp, cupy_dispatcher
 from scipy import ndimage as scipy_ndi
 
 def get_func(function_name):
-    if hasattr(xp_ndi, function_name):
-        _func = getattr(xp_ndi, function_name)    
+    if hasattr(xp.ndi, function_name):
+        _func = getattr(xp.ndi, function_name)    
         func = cupy_dispatcher(_func)
     else:
         func = getattr(scipy_ndi, function_name)
@@ -90,14 +90,14 @@ def std_filter(data, selem):
     selem = selem / np.sum(selem)
     x1 = convolve(data, selem)
     x2 = convolve(data**2, selem)
-    std_img = _safe_sqrt(asnumpy(x2 - x1**2), fill=0)
+    std_img = _safe_sqrt(xp.asnumpy(x2 - x1**2), fill=0)
     return std_img
 
 def coef_filter(data, selem):
     selem = selem / np.sum(selem)
     x1 = convolve(data, selem)
     x2 = convolve(data**2, selem)
-    out = _safe_sqrt(asnumpy(x2 - x1**2), fill=0)/asnumpy(x1)
+    out = _safe_sqrt(xp.asnumpy(x2 - x1**2), fill=0)/xp.asnumpy(x1)
     return out
     
 def dog_filter(img, low_sigma, high_sigma):
diff --git a/impy/arrays/_utils/_linalg.py b/impy/arrays/_utils/_linalg.py
index 4d275329..4a414d8b 100644
--- a/impy/arrays/_utils/_linalg.py
+++ b/impy/arrays/_utils/_linalg.py
@@ -1,16 +1,16 @@
 import numpy as np
 from skimage.feature.corner import _symmetric_image
 from ._skimage import skfeat
-from ..._cupy import xp, xp_linalg
+from ..._cupy import xp
 
 def eigh(a):
     a = xp.asarray(a, dtype=a.dtype)
-    val, vec = xp_linalg.eigh(a)
+    val, vec = xp.linalg.eigh(a)
     return val, vec
 
 def eigvalsh(a):
     a = xp.asarray(a, dtype=a.dtype)
-    val = xp_linalg.eigvalsh(a)
+    val = xp.linalg.eigvalsh(a)
     return val
 
 def structure_tensor_eigval(img, sigma, pxsize):
diff --git a/impy/arrays/_utils/_structures.py b/impy/arrays/_utils/_structures.py
index 782e7d4a..b9111443 100644
--- a/impy/arrays/_utils/_structures.py
+++ b/impy/arrays/_utils/_structures.py
@@ -1,4 +1,3 @@
-import numpy as np
 from ..._cupy import xp
 
 def circle(radius, shape, dtype="bool"):
diff --git a/impy/arrays/_utils/_transform.py b/impy/arrays/_utils/_transform.py
index c578356c..4b1b13a9 100644
--- a/impy/arrays/_utils/_transform.py
+++ b/impy/arrays/_utils/_transform.py
@@ -2,7 +2,7 @@
 import scipy
 from collections import namedtuple
 from ._skimage import sktrans
-from ..._cupy import xp, xp_ndi, xp_ndarray
+from ..._cupy import xp
 
 __all__ = ["compose_affine_matrix", 
            "decompose_affine_matrix",
@@ -15,16 +15,16 @@
                                             field_names=["translation", "rotation", "scale", "shear"]
                                             )
 
-def warp(img: xp_ndarray, matrix: xp_ndarray, cval=0, mode="constant", output_shape=None, order=1):
+def warp(img: xp.ndarray, matrix: xp.ndarray, cval=0, mode="constant", output_shape=None, order=1):
     img = xp.asarray(img, dtype=img.dtype)
     matrix = xp.asarray(matrix)
-    out = xp_ndi.affine_transform(img, matrix, cval=cval, mode=mode, output_shape=output_shape, 
+    out = xp.ndi.affine_transform(img, matrix, cval=cval, mode=mode, output_shape=output_shape, 
                                   order=order, prefilter=order>1)
     return out
 
-def shift(img: xp_ndarray, shift: xp_ndarray, cval=0, mode="constant", order=1):
+def shift(img: xp.ndarray, shift: xp.ndarray, cval=0, mode="constant", order=1):
     img = xp.asarray(img, dtype=img.dtype)
-    out = xp_ndi.shift(img, shift, cval=cval, mode=mode, 
+    out = xp.ndi.shift(img, shift, cval=cval, mode=mode, 
                        order=order, prefilter=order>1)
     return out
 
@@ -123,7 +123,7 @@ def check_matrix(matrices):
     return mtx
 
 def polar2d(
-    img: xp_ndarray,
+    img: xp.ndarray,
     rmax: int,
     dtheta: float = 0.1,
     order=1,
@@ -138,11 +138,11 @@ def polar2d(
     x = r * xp.cos(theta) + centers[1]
     coords = xp.stack([y, x], axis=0)
     img = xp.asarray(img, dtype=img.dtype)
-    out = xp_ndi.map_coordinates(img, coords, order=order, mode=mode, cval=cval, prefilter=order>1)
+    out = xp.ndi.map_coordinates(img, coords, order=order, mode=mode, cval=cval, prefilter=order>1)
     return out
 
 # def polar3d(
-#     img: xp_ndarray, 
+#     img: xp.ndarray, 
 #     rmax: int,
 #     dtheta: float = 0.1,
 #     order=1,
@@ -160,5 +160,5 @@ def polar2d(
 #     x = _yx * xp.cos(theta) + centers[2]
 #     coords = xp.stack([z, y, x], axis=0)
 #     img = xp.asarray(img, dtype=img.dtype)
-#     out = xp_ndi.map_coordinates(img, coords, order=order, mode=mode, cval=cval, prefilter=order>1)
+#     out = xp.ndi.map_coordinates(img, coords, order=order, mode=mode, cval=cval, prefilter=order>1)
 #     return out
diff --git a/impy/arrays/imgarray.py b/impy/arrays/imgarray.py
index 35a4c970..03943f27 100644
--- a/impy/arrays/imgarray.py
+++ b/impy/arrays/imgarray.py
@@ -23,7 +23,7 @@
 from ..collections import DataDict
 from .._types import nDInt, nDFloat, Dims, Coords, Iterable, Callable
 from .._const import Const
-from .._cupy import xp, xp_ndi, xp_fft, asnumpy, cupy_dispatcher
+from .._cupy import xp, cupy_dispatcher
 
 if TYPE_CHECKING:
     from ..frame import MarkerFrame, PathFrame
@@ -301,7 +301,9 @@ def binning(self, binsize: int = 2, method = "mean", *, check_edges: bool = True
         if binsize == 1:
             return self
         with Progress("binning"):
-            img_to_reshape, shape, scale_ = _misc.adjust_bin(self.value, binsize, check_edges, dims, self.axes)
+            img_to_reshape, shape, scale_ = _misc.adjust_bin(
+                self.value, binsize, check_edges, dims, self.axes
+            )
             
             reshaped_img = img_to_reshape.reshape(shape)
             axes_to_reduce = tuple(i*2+1 for i in range(self.ndim))
@@ -340,13 +342,13 @@ def radial_profile(self, nbin: int = 32, center: Iterable[float] = None, r_max:
             Radial profile stored in x-axis by default. If input image has tzcyx-axes, then an array 
             with tcx-axes will be returned.
         """        
-        func = {"mean": xp_ndi.mean,
-                "sum": xp_ndi.sum_labels,
-                "median": xp_ndi.median,
-                "max": xp_ndi.maximum,
-                "min": xp_ndi.minimum,
-                "std": xp_ndi.standard_deviation,
-                "var": xp_ndi.variance}[method]
+        func = {"mean": xp.ndi.mean,
+                "sum": xp.ndi.sum_labels,
+                "median": xp.ndi.median,
+                "max": xp.ndi.maximum,
+                "min": xp.ndi.minimum,
+                "std": xp.ndi.standard_deviation,
+                "var": xp.ndi.variance}[method]
         
         spatial_shape = self.sizesof(dims)
         inds = xp.indices(spatial_shape)
@@ -373,11 +375,17 @@ def radial_profile(self, nbin: int = 32, center: Iterable[float] = None, r_max:
         
         c_axes = complement_axes(dims, self.axes)
         
-        out = PropArray(np.empty(self.sizesof(c_axes)+(int(labels.max()),)), dtype=np.float32, axes=c_axes+dims[-1], 
-                        dirpath=self.dirpath, metadata=self.metadata, propname="radial_profile")
+        out = PropArray(
+            np.empty(self.sizesof(c_axes)+(int(labels.max()),)),
+            dtype=np.float32,
+            axes=c_axes+dims[-1], 
+            dirpath=self.dirpath, 
+            metadata=self.metadata, 
+            propname="radial_profile"
+        )
         radial_func = partial(cupy_dispatcher(func), labels=labels, index=xp.arange(1, labels.max()+1))
         for sl, img in self.iter(c_axes, exclude=dims):
-            out[sl] = asnumpy(radial_func(img))
+            out[sl] = xp.asnumpy(radial_func(img))
         return out
     
     @record
@@ -754,9 +762,9 @@ def lowpass_filter(self, cutoff: nDFloat = 0.2, order: float = 2, *, dims: Dims
         input = xp.asarray(self)
         if len(dims) < self.ndim:
             weight = add_axes(self.axes, self.shape, weight, dims)
-        out = xp_fft.irfftn(weight*xp_fft.rfftn(input, axes=spatial_axes), 
+        out = xp.fft.irfftn(weight*xp.fft.rfftn(input, axes=spatial_axes), 
                             s=spatial_shape, axes=spatial_axes)
-        return asnumpy(out)
+        return xp.asnumpy(out)
     
     @_docs.write_docs
     @dims_to_spatial_axes
@@ -788,7 +796,7 @@ def lowpass_conv_filter(self, cutoff: nDFloat = 0.2, order: float = 2, *,
             return self
         spatial_shape = self.sizesof(dims)
         weight = _get_ND_butterworth_filter(spatial_shape, cutoff, order, False, True)
-        ker_all = asnumpy(xp_fft.irfftn(weight, s=spatial_shape))
+        ker_all = xp.asnumpy(xp.fft.irfftn(weight, s=spatial_shape))
         ker_all = np.fft.fftshift(ker_all)
         sl = []
         for s, c in zip(spatial_shape, cutoff):
@@ -840,9 +848,9 @@ def func(arr):
             arr = xp.asarray(arr)
             shape = arr.shape
             weight = _get_ND_butterworth_filter(shape, cutoff, order, False, True)
-            ft = weight * xp_fft.rfftn(arr)
-            ift = xp_fft.irfftn(ft, s=shape)
-            return asnumpy(ift)
+            ft = weight * xp.fft.rfftn(arr)
+            ift = xp.fft.irfftn(ft, s=shape)
+            return xp.asnumpy(ift)
         
         input = da.from_array(self.value, chunks=chunks)
         out = da.map_overlap(func, input, depth=depth, boundary="reflect", dtype=self.dtype
@@ -881,9 +889,9 @@ def highpass_filter(self, cutoff: nDFloat = 0.2, order: float = 2, *, dims: Dims
         input = xp.asarray(self)
         if len(dims) < self.ndim:
             weight = add_axes(self.axes, self.shape, weight, dims)
-        out = xp_fft.irfftn(weight*xp_fft.rfftn(input, axes=spatial_axes), 
+        out = xp.fft.irfftn(weight*xp.fft.rfftn(input, axes=spatial_axes), 
                             s=spatial_shape, axes=spatial_axes)
-        return asnumpy(out)
+        return xp.asnumpy(out)
     
     
     @_docs.write_docs
@@ -1958,12 +1966,12 @@ def peak_local_max(self, *, min_distance: float = 1.0, percentile: float = None,
         for sl, img in self.iter(c_axes, israw=True, exclude=dims):
             # skfeat.peak_local_max overwrite something so we need to give copy of img.
             if use_labels and hasattr(img, "labels"):
-                labels = asnumpy(img.labels)
+                labels = xp.asnumpy(img.labels)
             else:
                 labels = None
             
-            indices = skfeat.peak_local_max(asnumpy(img),
-                                            footprint=asnumpy(_structures.ball_like(min_distance, ndim)),
+            indices = skfeat.peak_local_max(xp.asnumpy(img),
+                                            footprint=xp.asnumpy(_structures.ball_like(min_distance, ndim)),
                                             threshold_abs=thr,
                                             num_peaks=topn,
                                             num_peaks_per_label=topn_per_label,
@@ -2029,12 +2037,12 @@ def corner_peaks(self, *, min_distance:int=1, percentile:float=None,
         for sl, img in self.iter(c_axes, israw=True, exclude=dims):
             # skfeat.corner_peaks overwrite something so we need to give copy of img.
             if use_labels and hasattr(img, "labels"):
-                labels = asnumpy(img.labels)
+                labels = xp.asnumpy(img.labels)
             else:
                 labels = None
             
-            indices = skfeat.corner_peaks(asnumpy(img),
-                                          footprint=asnumpy(_structures.ball_like(min_distance, ndim)),
+            indices = skfeat.corner_peaks(xp.asnumpy(img),
+                                          footprint=xp.asnumpy(_structures.ball_like(min_distance, ndim)),
                                           threshold_abs=thr,
                                           num_peaks=topn,
                                           num_peaks_per_label=topn_per_label,
@@ -2761,10 +2769,10 @@ def fft(self, *, shape: int | Iterable[int] | str = "same", shift: bool = True,
         else:
             shape = check_nd(shape, len(dims))
         dtype = np.float64 if double_precision else np.float32
-        freq = xp_fft.fftn(xp.asarray(self.value, dtype=dtype), s=shape, axes=axes)
+        freq = xp.fft.fftn(xp.asarray(self.value, dtype=dtype), s=shape, axes=axes)
         if shift:
             freq[:] = np.fft.fftshift(freq, axes=axes)
-        return asnumpy(freq, dtype=np.complex64)
+        return xp.asnumpy(freq, dtype=np.complex64)
     
     @_docs.write_docs
     @dims_to_spatial_axes
@@ -2906,13 +2914,13 @@ def ifft(self, real: bool = True, *, shift: bool = True, double_precision = Fals
         else:
             freq = self.value
         dtype = np.complex128 if double_precision else np.complex64
-        out = xp_fft.ifftn(xp.asarray(freq, dtype=dtype), 
+        out = xp.fft.ifftn(xp.asarray(freq, dtype=dtype), 
                            axes=axes
                            ).astype(np.complex64)
         
         if real:
             out = np.real(out)
-        return asnumpy(out)
+        return xp.asnumpy(out)
     
     @_docs.write_docs
     @dims_to_spatial_axes
@@ -3249,7 +3257,7 @@ def skeletonize(self, radius: float = 0, *, dims: Dims = None, update=False) ->
             Skeletonized image.
         """        
         if radius >= 1:
-            selem = asnumpy(_structures.ball_like(radius, len(dims)))
+            selem = xp.asnumpy(_structures.ball_like(radius, len(dims)))
         else:
             selem = None
         
@@ -3907,7 +3915,7 @@ def track_drift(self, along: str = None, show_drift: bool = False,
         for i, (_, img) in enumerate(img_fft.iter(along)):
             img = xp.asarray(img)
             if last_img is not None:
-                result[i] = asnumpy(_corr.subpixel_pcc(last_img, img, upsample_factor=upsample_factor))
+                result[i] = xp.asnumpy(_corr.subpixel_pcc(last_img, img, upsample_factor=upsample_factor))
                 last_img = img
             else:
                 last_img = img
@@ -4116,14 +4124,14 @@ def pad_defocus(self, kernel, *, depth: int = 3, width: int = 6, bg: float = Non
         
         if kernel.ndim <= 1:
             def filter_func(img):
-                return asnumpy(_filters.gaussian_filter(img, kernel, mode="constant", cval=bg))
+                return xp.asnumpy(_filters.gaussian_filter(img, kernel, mode="constant", cval=bg))
             dz, dy, dx = kernel*3 # 3-sigma
             
         elif kernel.ndim == 3:
             kernel = kernel.astype(np.float32)
             kernel = kernel / np.sum(kernel)
             def filter_func(img):
-                return asnumpy(_filters.convolve(img, kernel, mode="constant", cval=bg))
+                return xp.asnumpy(_filters.convolve(img, kernel, mode="constant", cval=bg))
             dz, dy, dx = np.array(kernel.shape)//2
         else:
             raise ValueError("`kernel` only take 0, 1, 3 dimensional array as an input.")
diff --git a/impy/arrays/lazy.py b/impy/arrays/lazy.py
index f355807e..ed2b403c 100644
--- a/impy/arrays/lazy.py
+++ b/impy/arrays/lazy.py
@@ -25,7 +25,7 @@
 from .._types import nDFloat, Coords, Iterable, Dims
 from ..axes import ImageAxesError
 from .._const import Const
-from .._cupy import xp, xp_ndi, xp_fft, asnumpy
+from .._cupy import xp
 
 if TYPE_CHECKING:
     from dask import array as da
@@ -89,7 +89,7 @@ def gb(self):
     def __array__(self):
         # Should not be `self.compute` because in napari Viewer this function is called every time
         # sliders are moved.
-        return asnumpy(self.value.compute())
+        return xp.asnumpy(self.value.compute())
     
     def __getitem__(self, key):
         if isinstance(key, str):
@@ -261,7 +261,7 @@ def compute(self, ignore_limit: bool = False) -> ImgArray:
         with Progress("Converting to ImgArray"):
             arr = self.value.compute()
             if arr.ndim > 0:
-                img = asnumpy(arr).view(ImgArray)
+                img = xp.asnumpy(arr).view(ImgArray)
                 for attr in ["name", "dirpath", "axes", "metadata", "history"]:
                     setattr(img, attr, getattr(self, attr, None))
             else:
@@ -534,7 +534,7 @@ def _apply_dask_filter(self,
         #     it = self.axisof(a)
         #     output_chunks[it] = all_coords.shape[i+1]
         #
-        # cropped_img = self._apply_function(xp_ndi.map_coordinates, 
+        # cropped_img = self._apply_function(xp.ndi.map_coordinates, 
         #                          c_axes=complement_axes(dims, self.axes), 
         #                          dtype=self.dtype,
         #                          rechunk_to="max",
@@ -550,7 +550,7 @@ def erosion(self, radius: float = 1, *, dims: Dims = None, update: bool = False
                 ) -> LazyImgArray:
         disk = _structures.ball_like(radius, len(dims))
         c_axes = complement_axes(dims, self.axes)
-        filter_func = xp_ndi.grey_erosion if self.dtype != bool else xp_ndi.binary_erosion
+        filter_func = xp.ndi.grey_erosion if self.dtype != bool else xp.ndi.binary_erosion
         
         return self._apply_map_overlap(
             filter_func, 
@@ -566,7 +566,7 @@ def dilation(self, radius: float = 1, *, dims: Dims = None, update: bool = False
                  ) -> LazyImgArray:
         disk = _structures.ball_like(radius, len(dims))
         c_axes = complement_axes(dims, self.axes)
-        filter_func = xp_ndi.grey_dilation if self.dtype != bool else xp_ndi.binary_dilation
+        filter_func = xp.ndi.grey_dilation if self.dtype != bool else xp.ndi.binary_dilation
         
         return self._apply_map_overlap(
             filter_func, 
@@ -582,7 +582,7 @@ def opening(self, radius: float = 1, *, dims: Dims = None, update: bool = False
                 ) -> LazyImgArray:
         disk = _structures.ball_like(radius, len(dims))
         c_axes = complement_axes(dims, self.axes)
-        filter_func = xp_ndi.grey_opening if self.dtype != bool else xp_ndi.binary_opening
+        filter_func = xp.ndi.grey_opening if self.dtype != bool else xp.ndi.binary_opening
         
         return self._apply_map_overlap(
             filter_func, 
@@ -598,7 +598,7 @@ def closing(self, radius: float = 1, *, dims: Dims = None, update: bool = False
                 ) -> LazyImgArray:
         disk = _structures.ball_like(radius, len(dims))
         c_axes = complement_axes(dims, self.axes)
-        filter_func = xp_ndi.grey_closing if self.dtype != bool else xp_ndi.binary_closing
+        filter_func = xp.ndi.grey_closing if self.dtype != bool else xp.ndi.binary_closing
         
         return self._apply_map_overlap(
             filter_func, 
@@ -616,7 +616,7 @@ def gaussian_filter(self, sigma: nDFloat = 1.0, *, dims: Dims = None, update: bo
         c_axes = complement_axes(dims, self.axes)
         depth = _ceilint(sigma*4)
         return self._apply_map_overlap(
-            xp_ndi.gaussian_filter, 
+            xp.ndi.gaussian_filter, 
             c_axes=c_axes,
             depth=depth,
             kwargs=dict(sigma=sigma),
@@ -630,7 +630,7 @@ def median_filter(self, radius: float = 1, *, dims: Dims = None, update: bool =
                       ) -> LazyImgArray:
         disk = _structures.ball_like(radius, len(dims))
         return self._apply_map_overlap(
-            xp_ndi.median_filter,
+            xp.ndi.median_filter,
             depth=_ceilint(radius),
             c_axes=complement_axes(dims, self.axes),
             kwargs=dict(footprint=disk)
@@ -645,7 +645,7 @@ def mean_filter(self, radius: float = 1, *, dims: Dims = None, update: bool = Fa
         disk = _structures.ball_like(radius, len(dims))
         kernel = (disk/np.sum(disk)).astype(np.float32)
         return self._apply_map_overlap(
-            xp_ndi.convolve,
+            xp.ndi.convolve,
             depth=_ceilint(radius),
             c_axes=complement_axes(dims, self.axes),
             kwargs=dict(weights=kernel),
@@ -664,7 +664,7 @@ def convolve(self, kernel, *, mode: str = "reflect", cval: float = 0, dims: Dims
         depth = tuple(half_size)
         c_axes = complement_axes(dims, self.axes)
         return self._apply_map_overlap(
-            xp_ndi.convolve, 
+            xp.ndi.convolve, 
             c_axes=c_axes,
             depth=depth,
             kwargs=dict(weights=kernel, mode=mode, cval=cval),
@@ -701,7 +701,7 @@ def laplacian_filter(self, radius: int = 1, *, dims: Dims = None, update: bool =
         ndim = len(dims)
         _, laplace_op = skres.uft.laplacian(ndim, (2*radius+1,) * ndim)
         return self._apply_map_overlap(
-            xp_ndi.convolve,
+            xp.ndi.convolve,
             depth=_ceilint(radius),
             c_axes=complement_axes(dims, self.axes),
             args=(laplace_op,),
@@ -850,8 +850,8 @@ def func(arr):
             arr = xp.asarray(arr)
             shape = arr.shape
             weight = _get_ND_butterworth_filter(shape, cutoff, order, False, True)
-            ft = weight * xp_fft.rfftn(arr)
-            ift = xp_fft.irfftn(ft, s=shape)
+            ft = weight * xp.fft.rfftn(arr)
+            ift = xp.fft.irfftn(ft, s=shape)
             return ift
         
         out = self._apply_map_overlap(func, c_axes=c_axes, depth=depth, boundary="reflect")
@@ -921,7 +921,7 @@ def pcc(x):
                 return np.array([0]*ndim, dtype=np.float32).reshape(*each_shape)
             x = xp.asarray(x)
             result = _corr.subpixel_pcc(x[0], x[1], upsample_factor=upsample_factor)
-            return asnumpy(result[slice_out])
+            return xp.asnumpy(result[slice_out])
         
         from dask import array as da
 
@@ -994,7 +994,7 @@ def warp(arr, shift, block_info=None):
             mx = xp.eye(ndim+1, dtype=np.float32)
             loc = block_info[None]["array-location"][0]
             mx[:-1, -1] = -xp.asarray(shift[loc[t_index]])
-            return asnumpy(
+            return xp.asnumpy(
                 _transform.warp(arr[slice_in], mx, **affine_kwargs)[slice_out]
                 )
         
diff --git a/impy/correlation.py b/impy/correlation.py
index ce4f9cb7..4f5282e4 100644
--- a/impy/correlation.py
+++ b/impy/correlation.py
@@ -9,7 +9,7 @@
 from .utils.axesop import complement_axes, add_axes
 from .utils.utilcls import Progress
 from .utils.deco import dims_to_spatial_axes
-from ._cupy import xp, asnumpy, xp_ndi
+from ._cupy import xp
 from ._types import Dims
 
 __all__ = ["fsc", "fourier_shell_correlation", "ncc", "zncc", "fourier_ncc", "fourier_zncc",
@@ -62,12 +62,12 @@ def fsc(
     with Progress("fsc"):
         # make radially separated labels
         labels = (r/dfreq).astype(np.uint16)
-        nlabels = int(asnumpy(labels.max()))
+        nlabels = int(xp.asnumpy(labels.max()))
         
         out = xp.empty(nlabels, dtype=xp.float32)
         def radial_sum(arr):
             arr = xp.asarray(arr)
-            return xp_ndi.sum_labels(arr, labels=labels, index=xp.arange(1, nlabels+1))
+            return xp.ndi.sum_labels(arr, labels=labels, index=xp.arange(1, nlabels+1))
 
         f0 = img0.fft(dims=dims)
         f1 = img1.fft(dims=dims)
@@ -93,7 +93,7 @@ def _ncc(img0: ImgArray, img1: ImgArray, dims: Dims):
     img1 = xp.asarray(img1)
     corr = xp.sum(img0 * img1, axis=dims) / (
         xp.std(img0, axis=dims)*xp.std(img1, axis=dims)) / n
-    return asnumpy(corr)
+    return xp.asnumpy(corr)
 
 
 def _masked_ncc(img0: ImgArray, img1: ImgArray, dims: Dims, mask: ImgArray):
@@ -116,7 +116,7 @@ def _zncc(img0: ImgArray, img1: ImgArray, dims: Dims):
     img1 = xp.asarray(img1)
     corr = xp.sum(img0 * img1, axis=dims) / (
         xp.sqrt(xp.sum(img0**2, axis=dims)*xp.sum(img1**2, axis=dims)))
-    return asnumpy(corr)
+    return xp.asnumpy(corr)
 
 
 def _masked_zncc(img0: ImgArray, img1: ImgArray, dims: Dims, mask: ImgArray):
@@ -375,7 +375,7 @@ def pcc_maximum(
             upsample_factor, 
             max_shifts=max_shifts
         )
-    return asnumpy(shift)
+    return xp.asnumpy(shift)
 
 @_docs.write_docs
 def ft_pcc_maximum(
@@ -418,7 +418,7 @@ def ft_pcc_maximum(
             upsample_factor,
             max_shifts=max_shifts,
         )
-    return asnumpy(shift)
+    return xp.asnumpy(shift)
 
 @_docs.write_docs
 def polar_pcc_maximum(
diff --git a/tests/test_fft.py b/tests/test_fft.py
index ded0c215..4a5a8ba4 100644
--- a/tests/test_fft.py
+++ b/tests/test_fft.py
@@ -1,7 +1,7 @@
 import impy as ip
 from pathlib import Path
 import numpy as np
-from impy._cupy import xp, xp_fft, asnumpy
+from impy._cupy import xp
 from numpy.testing import assert_allclose
 ip.Const["SHOW_PROGRESS"] = False
 
@@ -10,11 +10,11 @@ def test_precision():
     img = ip.imread(path)["c=1;t=0"].as_float()
     
     assert_allclose(img.fft(shift=False),
-                    asnumpy(xp_fft.fftn(xp.asarray(img.value)))
+                    xp.asnumpy(xp.fft.fftn(xp.asarray(img.value)))
                     )
     
     assert_allclose(img.fft(shift=False, double_precision=True),
-                    asnumpy(xp_fft.fftn(xp.asarray(img.value.astype(np.float64))).astype(np.complex64))
+                    xp.asnumpy(xp.fft.fftn(xp.asarray(img.value.astype(np.float64))).astype(np.complex64))
                     )
     
     assert_allclose(img.fft().ifft(), img, rtol=1e-6)

From 9c494790843db42b5bb3f616f2b031e5007f2836 Mon Sep 17 00:00:00 2001
From: hanjinliu <hanjin.liu@bs.s.u-tokyo.ac.jp>
Date: Fri, 25 Feb 2022 01:53:09 +0900
Subject: [PATCH 03/13] refactor cupy

---
 impy/__init__.py               |  8 ------
 impy/_const.py                 | 51 +++++++++++++++++++++-------------
 impy/arrays/bases/metaarray.py |  8 +++---
 impy/utils/deco.py             |  6 ++--
 4 files changed, 39 insertions(+), 34 deletions(-)

diff --git a/impy/__init__.py b/impy/__init__.py
index 1a3b652b..e6a13126 100644
--- a/impy/__init__.py
+++ b/impy/__init__.py
@@ -7,14 +7,6 @@
 
 from ._const import Const, SetConst
 
-from ._cupy import GPU_AVAILABLE
-if GPU_AVAILABLE:
-    Const._setitem_("RESOURCE", "cupy")
-    Const["SCHEDULER"] = "single-threaded"
-else:
-    Const._setitem_("RESOURCE", "numpy")
-del GPU_AVAILABLE
-
 from .collections import *
 from .core import *
 from .binder import bind
diff --git a/impy/_const.py b/impy/_const.py
index d51055e9..3696f84c 100644
--- a/impy/_const.py
+++ b/impy/_const.py
@@ -1,17 +1,26 @@
 import dask
 import psutil
+from typing import Any, MutableMapping
 
 memory = psutil.virtual_memory()
 
 MAX_GB_LIMIT = memory.total / 2 * 1e-9
 
-class GlobalConstant(dict):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+class GlobalConstant(MutableMapping[str, Any]):
+    __const: dict[str, Any]
+    
+    def __init__(self, **kwargs):
+        object.__setattr__(self, "__const", dict(**kwargs))
+    
+    def __len__(self) -> int:
+        return len(self.__const)
+    
+    def __iter__(self):
+        raise StopIteration
 
     def __getitem__(self, k):
         try:
-            return super().__getitem__(k)
+            return self.__const[k]
         except KeyError:
             raise KeyError(f"Global constants: {', '.join(self.keys())}")
 
@@ -33,18 +42,21 @@ def __setitem__(self, k, v):
                 raise ValueError("ID_AXIS must be single character.")
         elif k == "FONT_SIZE_FACTOR":
             if not isinstance(v, (int, float)):
-                raise TypeError("MAX_GB must be float.")
+                raise TypeError("FONT_SIZE_FACTOR must be float.")
         elif k == "RESOURCE":
-            raise RuntimeError("Cannot set RESOURCE.")
+            from ._cupy import xp
+            if v == "numpy":
+                xp.setNumpy()
+            elif v == "cupy":
+                xp.setCupy()
+            else:
+                raise ValueError("RESOURCES must be either 'numpy' or 'cupy'.")
         elif k == "SCHEDULER":
             dask.config.set(scheduler=v)
         else:
             raise RuntimeError("Cannot set new keys.")
         
-        super().__setitem__(k, v)
-
-    def _setitem_(self, k, v):
-        super().__setitem__(k, v)
+        self.__const[k] = v
     
     __getattr__ = __getitem__
     __setattr__ = __setitem__
@@ -53,15 +65,16 @@ def __delitem__(self, v):
         raise RuntimeError("Cannot delete any items.")
     
     def __repr__(self):
-        return \
-        f"""
-              MAX_GB    : {self.MAX_GB:.2f} GB
-          SHOW_PROGRESS : {self.SHOW_PROGRESS}
-             ID_AXIS    : {self.ID_AXIS}
-        FONT_SIZE_FACTOR: {self.FONT_SIZE_FACTOR}
-             RESOURCE   : {self.RESOURCE}
-            SCHEDULER   : {self.SCHEDULER}
-        """
+        return (
+            f"""
+                MAX_GB    : {self.MAX_GB:.2f} GB
+            SHOW_PROGRESS : {self.SHOW_PROGRESS}
+                ID_AXIS    : {self.ID_AXIS}
+            FONT_SIZE_FACTOR: {self.FONT_SIZE_FACTOR}
+                RESOURCE   : {self.RESOURCE}
+                SCHEDULER   : {self.SCHEDULER}
+            """
+        )
 
 Const = GlobalConstant(
     MAX_GB = MAX_GB_LIMIT/2,
diff --git a/impy/arrays/bases/metaarray.py b/impy/arrays/bases/metaarray.py
index 141363c3..c1f48ddc 100644
--- a/impy/arrays/bases/metaarray.py
+++ b/impy/arrays/bases/metaarray.py
@@ -3,7 +3,7 @@
 from ..axesmixin import AxesMixin, get_axes_tuple
 from ..._types import *
 from ...axes import ImageAxesError
-from ..._cupy import xp, xp_ndarray, asnumpy
+from ..._cupy import xp
 from ...utils.axesop import *
 from ...utils.slicer import *
 from ...collections import DataList
@@ -306,7 +306,7 @@ def apply_dask(self,
         
         if len(c_axes) == 0:
             # Do not construct dask tasks if it is not needed.
-            out = asnumpy(func(self.value, *args, **kwargs), dtype=dtype)
+            out = xp.asnumpy(func(self.value, *args, **kwargs), dtype=dtype)
         else:
             from dask import array as da
             new_axis = _list_of_axes(self, new_axis)
@@ -336,7 +336,7 @@ def apply_dask(self,
             img_idx = []
             args = []
             for i, arg in enumerate(all_args):
-                if isinstance(arg, (np.ndarray, xp_ndarray)) and arg.shape == self.shape:
+                if isinstance(arg, (np.ndarray, xp.ndarray)) and arg.shape == self.shape:
                     args.append(da.from_array(arg, chunks=chunks))
                     img_idx.append(i)
                 else:
@@ -349,7 +349,7 @@ def _func(*args, **kwargs):
                         continue
                     args[i] = args[i][slice_in]
                 out = func(*args, **kwargs)
-                return asnumpy(out[slice_out])
+                return xp.asnumpy(out[slice_out])
             
             out = da.map_blocks(_func, 
                                 *args, 
diff --git a/impy/utils/deco.py b/impy/utils/deco.py
index faae6a15..88f76c42 100644
--- a/impy/utils/deco.py
+++ b/impy/utils/deco.py
@@ -3,7 +3,7 @@
 import inspect
 import re
 from .utilcls import Progress
-from .._cupy import xp_ndarray, asnumpy
+from .._cupy import xp
 
 __all__ = ["record",
            "record_lazy",
@@ -37,8 +37,8 @@ def _record(self, *args, **kwargs):
             
             temp = getattr(out, "temp", None)
             
-            if type(out) in (np.ndarray, xp_ndarray):
-                out = asnumpy(out).view(self.__class__)
+            if type(out) in (np.ndarray, xp.ndarray):
+                out = xp.asnumpy(out).view(self.__class__)
             
             # record history and update if needed
             ifupdate = kwargs.pop("update", False)

From f59329aa8c804b2ab3ede09c78a1cf571f9a8f88 Mon Sep 17 00:00:00 2001
From: hanjinliu <hanjin.liu@bs.s.u-tokyo.ac.jp>
Date: Fri, 25 Feb 2022 01:59:37 +0900
Subject: [PATCH 04/13] cons bug fix

---
 impy/_const.py | 27 ++++++++++++---------------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/impy/_const.py b/impy/_const.py
index 3696f84c..68b75e6c 100644
--- a/impy/_const.py
+++ b/impy/_const.py
@@ -1,4 +1,3 @@
-import dask
 import psutil
 from typing import Any, MutableMapping
 
@@ -7,22 +6,19 @@
 MAX_GB_LIMIT = memory.total / 2 * 1e-9
 
 class GlobalConstant(MutableMapping[str, Any]):
-    __const: dict[str, Any]
+    _const: dict[str, Any]
     
     def __init__(self, **kwargs):
-        object.__setattr__(self, "__const", dict(**kwargs))
+        object.__setattr__(self, "_const", dict(**kwargs))
     
     def __len__(self) -> int:
-        return len(self.__const)
+        return len(self._const)
     
     def __iter__(self):
         raise StopIteration
 
     def __getitem__(self, k):
-        try:
-            return self.__const[k]
-        except KeyError:
-            raise KeyError(f"Global constants: {', '.join(self.keys())}")
+        return self._const[k]
 
     def __setitem__(self, k, v):
         if k == "MAX_GB":
@@ -52,11 +48,12 @@ def __setitem__(self, k, v):
             else:
                 raise ValueError("RESOURCES must be either 'numpy' or 'cupy'.")
         elif k == "SCHEDULER":
+            import dask
             dask.config.set(scheduler=v)
         else:
             raise RuntimeError("Cannot set new keys.")
         
-        self.__const[k] = v
+        self._const[k] = v
     
     __getattr__ = __getitem__
     __setattr__ = __setitem__
@@ -67,12 +64,12 @@ def __delitem__(self, v):
     def __repr__(self):
         return (
             f"""
-                MAX_GB    : {self.MAX_GB:.2f} GB
-            SHOW_PROGRESS : {self.SHOW_PROGRESS}
-                ID_AXIS    : {self.ID_AXIS}
-            FONT_SIZE_FACTOR: {self.FONT_SIZE_FACTOR}
-                RESOURCE   : {self.RESOURCE}
-                SCHEDULER   : {self.SCHEDULER}
+                  MAX_GB    : {self['MAX_GB']:.2f} GB
+              SHOW_PROGRESS : {self['SHOW_PROGRESS']}
+                 ID_AXIS    : {self['ID_AXIS']}
+            FONT_SIZE_FACTOR: {self['FONT_SIZE_FACTOR']}
+                 RESOURCE   : {self['RESOURCE']}
+                SCHEDULER   : {self['SCHEDULER']}
             """
         )
 

From cd0b8a1c9e2e4dc03f0a4e3c03258fb522d5cd13 Mon Sep 17 00:00:00 2001
From: hanjinliu <hanjin.liu@bs.s.u-tokyo.ac.jp>
Date: Fri, 25 Feb 2022 02:01:43 +0900
Subject: [PATCH 05/13] set scheduler

---
 impy/_cupy.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/impy/_cupy.py b/impy/_cupy.py
index bfa4c888..0d5b42d7 100644
--- a/impy/_cupy.py
+++ b/impy/_cupy.py
@@ -42,6 +42,8 @@ def setNumpy(self) -> None:
         self.asarray = np.asarray
         self.ndarray = np.ndarray
         self.state = "numpy"
+        from ._const import Const
+        Const["SCHEDULER"] = "threads"
     
     def setCupy(self) -> None:
         import cupy as cp
@@ -64,6 +66,9 @@ def cp_asnumpy(arr, dtype=None):
         self.asarray = cp.asarray
         self.ndarray = cp_ndarray
         self.state = "cupy"
+        
+        from ._const import Const
+        Const["SCHEDULER"] = "single-threaded"
 
 xp = XP()
     
\ No newline at end of file

From 8e2a79db91476ed4daaf9575b9ddf06124f121a3 Mon Sep 17 00:00:00 2001
From: hanjinliu <hanjin.liu@bs.s.u-tokyo.ac.jp>
Date: Sun, 27 Feb 2022 11:26:50 +0900
Subject: [PATCH 06/13] update array API

---
 impy/_const.py                    |  2 +-
 impy/{_cupy.py => array_api.py}   | 69 +++++++++++++++++++++++++++++++
 impy/arrays/_utils/_corr.py       |  2 +-
 impy/arrays/_utils/_deconv.py     |  9 ++--
 impy/arrays/_utils/_filters.py    |  4 +-
 impy/arrays/_utils/_linalg.py     |  2 +-
 impy/arrays/_utils/_misc.py       |  2 +-
 impy/arrays/_utils/_skimage.py    |  2 +-
 impy/arrays/_utils/_structures.py |  9 ++--
 impy/arrays/_utils/_transform.py  |  2 +-
 impy/arrays/bases/metaarray.py    |  2 +-
 impy/arrays/imgarray.py           |  4 +-
 impy/arrays/lazy.py               |  2 +-
 impy/correlation.py               |  4 +-
 impy/utils/deco.py                |  2 +-
 impy/utils/io.py                  |  2 +-
 tests/test_fft.py                 |  2 +-
 17 files changed, 96 insertions(+), 25 deletions(-)
 rename impy/{_cupy.py => array_api.py} (50%)

diff --git a/impy/_const.py b/impy/_const.py
index 68b75e6c..b2615d81 100644
--- a/impy/_const.py
+++ b/impy/_const.py
@@ -40,7 +40,7 @@ def __setitem__(self, k, v):
             if not isinstance(v, (int, float)):
                 raise TypeError("FONT_SIZE_FACTOR must be float.")
         elif k == "RESOURCE":
-            from ._cupy import xp
+            from .array_api import xp
             if v == "numpy":
                 xp.setNumpy()
             elif v == "cupy":
diff --git a/impy/_cupy.py b/impy/array_api.py
similarity index 50%
rename from impy/_cupy.py
rename to impy/array_api.py
index 0d5b42d7..ec77c062 100644
--- a/impy/_cupy.py
+++ b/impy/array_api.py
@@ -16,6 +16,11 @@ def func(*args, **kwargs):
 from scipy import ndimage as scipy_ndi
 from typing import Callable
 
+# CUDA <= ver.8 does not have gradient    
+def _gradient(a, axis=None):
+    out = np.gradient(a.get(), axis=axis)
+    return xp.asarray(out)
+
 class XP:
     fft: ModuleType
     linalg: ModuleType
@@ -41,6 +46,37 @@ def setNumpy(self) -> None:
         self.asnumpy = np.asarray
         self.asarray = np.asarray
         self.ndarray = np.ndarray
+        self.empty = np.empty
+        self.zeros = np.zeros
+        self.ones = np.ones
+        self.array = np.array
+        self.exp = np.exp
+        self.sin = np.sin
+        self.cos = np.cos
+        self.tan = np.tan
+        self.sqrt = np.sqrt
+        self.mean = np.mean
+        self.max = np.max
+        self.min = np.min
+        self.median = np.median
+        self.sum = np.sum
+        self.prod = np.prod
+        self.std = np.std
+        self.meshgrid = np.meshgrid
+        self.indices = np.indices
+        self.arange = np.arange
+        self.linspace = np.linspace
+        self.real = np.real
+        self.imag = np.imag
+        self.conjugate = np.conjugate
+        self.angle = np.angle
+        self.abs = np.abs
+        self.mod = np.mod
+        self.fix = np.fix
+        self.gradient = np.gradient
+        self.tensordot = np.tensordot
+        self.stack = np.stack
+        
         self.state = "numpy"
         from ._const import Const
         Const["SCHEDULER"] = "threads"
@@ -65,6 +101,39 @@ def cp_asnumpy(arr, dtype=None):
         self.asnumpy = cp_asnumpy
         self.asarray = cp.asarray
         self.ndarray = cp_ndarray
+        self.empty = cp.empty
+        self.zeros = cp.zeros
+        self.ones = cp.ones
+        self.array = cp.array
+        self.exp = cp.exp
+        self.sin = cp.sin
+        self.cos = cp.cos
+        self.tan = cp.tan
+        self.sqrt = cp.sqrt
+        self.mean = cp.mean
+        self.max = cp.max
+        self.min = cp.min
+        self.median = cp.median
+        self.sum = cp.sum
+        self.prod = cp.prod
+        self.std = cp.std
+        self.meshgrid = cp.meshgrid
+        self.indices = cp.indices
+        self.arange = cp.arange
+        self.linspace = cp.linspace
+        self.real = cp.real
+        self.imag = cp.imag
+        self.conjugate = cp.conjugate
+        self.angle = cp.angle
+        self.abs = cp.abs
+        self.mod = cp.mod
+        self.fix = cp.fix
+        try:
+            self.gradient = cp.gradient
+        except AttributeError:
+            self.gradient = _gradient
+        self.tensordot = cp.tensordot
+        self.stack = cp.stack
         self.state = "cupy"
         
         from ._const import Const
diff --git a/impy/arrays/_utils/_corr.py b/impy/arrays/_utils/_corr.py
index b849acc1..2983d182 100644
--- a/impy/arrays/_utils/_corr.py
+++ b/impy/arrays/_utils/_corr.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import numpy as np
-from ..._cupy import xp
+from ...array_api import xp
 
 # Modified from skimage/registration/_phase_cross_correlation.py
 # Compatible between numpy/cupy and supports maximum shifts.
diff --git a/impy/arrays/_utils/_deconv.py b/impy/arrays/_utils/_deconv.py
index 42f77215..9ba75e7b 100644
--- a/impy/arrays/_utils/_deconv.py
+++ b/impy/arrays/_utils/_deconv.py
@@ -1,5 +1,6 @@
 from functools import partial
-from ..._cupy import xp
+import numpy as np
+from ...array_api import xp
 
 try:
     gradient = xp.gradient
@@ -47,7 +48,7 @@ def richardson_lucy_tv(obs, psf_ft, psf_ft_conj, max_iter, lmd, tol, eps):
     ifft = partial(xp.fft.irfftn, s=obs.shape)
     est_old = ifft(fft(obs) * psf_ft).real
     est_new = xp.empty(obs.shape, dtype=xp.float32)
-    conv = factor = norm = gg = xp.empty(obs.shape, dtype=xp.float32) # placeholder
+    conv = factor = norm = gg = xp.empty(obs.shape, dtype=np.float32) # placeholder
     
     for _ in range(max_iter):
         conv[:] = ifft(fft(est_old) * psf_ft).real
@@ -67,14 +68,14 @@ def richardson_lucy_tv(obs, psf_ft, psf_ft_conj, max_iter, lmd, tol, eps):
 
 
 def _safe_div(a, b, eps=1e-8):
-    out = xp.zeros(a.shape, dtype=xp.float32)
+    out = xp.zeros(a.shape, dtype=np.float32)
     mask = b > eps
     out[mask] = a[mask]/b[mask]
     return out
 
 
 def check_psf(img, psf, dims):
-    psf = xp.asarray(psf, dtype=xp.float32)
+    psf = xp.asarray(psf, dtype=np.float32)
     psf /= xp.sum(psf)
     
     if img.sizesof(dims) != psf.shape:
diff --git a/impy/arrays/_utils/_filters.py b/impy/arrays/_utils/_filters.py
index 76a9222c..42c352a9 100644
--- a/impy/arrays/_utils/_filters.py
+++ b/impy/arrays/_utils/_filters.py
@@ -30,7 +30,7 @@
            ]
 
 
-from ..._cupy import xp, cupy_dispatcher
+from ...array_api import xp, cupy_dispatcher
 from scipy import ndimage as scipy_ndi
 
 def get_func(function_name):
@@ -81,7 +81,7 @@ def mean_filter(img, selem):
     return convolve(img, selem/np.sum(selem))
 
 def phase_mean_filter(img, selem, a):
-    out = xp.empty(img.shape, dtype=xp.complex64)
+    out = xp.empty(img.shape, dtype=np.complex64)
     xp.exp(1j*a*img, out=out)
     convolve(out, selem, output=out)
     return xp.angle(out)/a
diff --git a/impy/arrays/_utils/_linalg.py b/impy/arrays/_utils/_linalg.py
index 4a414d8b..85ec7410 100644
--- a/impy/arrays/_utils/_linalg.py
+++ b/impy/arrays/_utils/_linalg.py
@@ -1,7 +1,7 @@
 import numpy as np
 from skimage.feature.corner import _symmetric_image
 from ._skimage import skfeat
-from ..._cupy import xp
+from ...array_api import xp
 
 def eigh(a):
     a = xp.asarray(a, dtype=a.dtype)
diff --git a/impy/arrays/_utils/_misc.py b/impy/arrays/_utils/_misc.py
index 653afabf..3f417337 100644
--- a/impy/arrays/_utils/_misc.py
+++ b/impy/arrays/_utils/_misc.py
@@ -1,5 +1,5 @@
 import numpy as np
-from ..._cupy import xp
+from ...array_api import xp
 
 def adjust_bin(img, binsize: int, check_edges: bool, dims: str, all_axes: str):
     shape = []
diff --git a/impy/arrays/_utils/_skimage.py b/impy/arrays/_utils/_skimage.py
index c3fb7f79..3ab930bb 100644
--- a/impy/arrays/_utils/_skimage.py
+++ b/impy/arrays/_utils/_skimage.py
@@ -13,7 +13,7 @@
 from skimage import util as skutil
 
 from functools import reduce, lru_cache
-from ..._cupy import xp
+from ...array_api import xp
 
 # same as the function in skimage.filters._fft_based (available in scikit-image >= 0.19)
 @lru_cache(maxsize=4)
diff --git a/impy/arrays/_utils/_structures.py b/impy/arrays/_utils/_structures.py
index b9111443..bedce31d 100644
--- a/impy/arrays/_utils/_structures.py
+++ b/impy/arrays/_utils/_structures.py
@@ -1,4 +1,5 @@
-from ..._cupy import xp
+import numpy as np
+from ...array_api import xp
 
 def circle(radius, shape, dtype="bool"):
     x = xp.arange(-(shape[0] - 1) / 2, (shape[0] - 1) / 2 + 1)
@@ -11,15 +12,15 @@ def ball_like(radius, ndim:int):
     L = xp.arange(-half_int, half_int + 1)
     
     if ndim == 1:
-        return xp.ones(int(radius)*2+1, dtype=xp.uint8)
+        return xp.ones(int(radius)*2+1, dtype=np.uint8)
     elif ndim == 2:
         X, Y = xp.meshgrid(L, L)
         s = X**2 + Y**2
-        return xp.array(s <= radius**2, dtype=xp.uint8)
+        return xp.array(s <= radius**2, dtype=np.uint8)
     elif ndim == 3:
         Z, Y, X = xp.meshgrid(L, L, L)
         s = X**2 + Y**2 + Z**2
-        return xp.array(s <= radius**2, dtype=xp.uint8)
+        return xp.array(s <= radius**2, dtype=np.uint8)
     else:
         raise ValueError(f"dims must be 1 - 3, but got {ndim}")
 
diff --git a/impy/arrays/_utils/_transform.py b/impy/arrays/_utils/_transform.py
index 4b1b13a9..30eb3a01 100644
--- a/impy/arrays/_utils/_transform.py
+++ b/impy/arrays/_utils/_transform.py
@@ -2,7 +2,7 @@
 import scipy
 from collections import namedtuple
 from ._skimage import sktrans
-from ..._cupy import xp
+from ...array_api import xp
 
 __all__ = ["compose_affine_matrix", 
            "decompose_affine_matrix",
diff --git a/impy/arrays/bases/metaarray.py b/impy/arrays/bases/metaarray.py
index c1f48ddc..8825b7dd 100644
--- a/impy/arrays/bases/metaarray.py
+++ b/impy/arrays/bases/metaarray.py
@@ -3,7 +3,7 @@
 from ..axesmixin import AxesMixin, get_axes_tuple
 from ..._types import *
 from ...axes import ImageAxesError
-from ..._cupy import xp
+from ...array_api import xp
 from ...utils.axesop import *
 from ...utils.slicer import *
 from ...collections import DataList
diff --git a/impy/arrays/imgarray.py b/impy/arrays/imgarray.py
index 03943f27..e482ea2c 100644
--- a/impy/arrays/imgarray.py
+++ b/impy/arrays/imgarray.py
@@ -23,7 +23,7 @@
 from ..collections import DataDict
 from .._types import nDInt, nDFloat, Dims, Coords, Iterable, Callable
 from .._const import Const
-from .._cupy import xp, cupy_dispatcher
+from ..array_api import xp, cupy_dispatcher
 
 if TYPE_CHECKING:
     from ..frame import MarkerFrame, PathFrame
@@ -4445,4 +4445,4 @@ def wave_num(sl: slice, s: int, uf: int) -> xp.ndarray:
             raise ValueError(f"Invalid value encountered in slice {sl}.")
         
     n = stop - start
-    return xp.linspace(start, stop, n*uf, endpoint=False)[:, xp.newaxis]
\ No newline at end of file
+    return xp.linspace(start, stop, n*uf, endpoint=False)[:, np.newaxis]
\ No newline at end of file
diff --git a/impy/arrays/lazy.py b/impy/arrays/lazy.py
index ed2b403c..a2a88ff5 100644
--- a/impy/arrays/lazy.py
+++ b/impy/arrays/lazy.py
@@ -25,7 +25,7 @@
 from .._types import nDFloat, Coords, Iterable, Dims
 from ..axes import ImageAxesError
 from .._const import Const
-from .._cupy import xp
+from ..array_api import xp
 
 if TYPE_CHECKING:
     from dask import array as da
diff --git a/impy/correlation.py b/impy/correlation.py
index 4f5282e4..6334568a 100644
--- a/impy/correlation.py
+++ b/impy/correlation.py
@@ -9,7 +9,7 @@
 from .utils.axesop import complement_axes, add_axes
 from .utils.utilcls import Progress
 from .utils.deco import dims_to_spatial_axes
-from ._cupy import xp
+from .array_api import xp
 from ._types import Dims
 
 __all__ = ["fsc", "fourier_shell_correlation", "ncc", "zncc", "fourier_ncc", "fourier_zncc",
@@ -64,7 +64,7 @@ def fsc(
         labels = (r/dfreq).astype(np.uint16)
         nlabels = int(xp.asnumpy(labels.max()))
         
-        out = xp.empty(nlabels, dtype=xp.float32)
+        out = xp.empty(nlabels, dtype=np.float32)
         def radial_sum(arr):
             arr = xp.asarray(arr)
             return xp.ndi.sum_labels(arr, labels=labels, index=xp.arange(1, nlabels+1))
diff --git a/impy/utils/deco.py b/impy/utils/deco.py
index 88f76c42..61d5560d 100644
--- a/impy/utils/deco.py
+++ b/impy/utils/deco.py
@@ -3,7 +3,7 @@
 import inspect
 import re
 from .utilcls import Progress
-from .._cupy import xp
+from ..array_api import xp
 
 __all__ = ["record",
            "record_lazy",
diff --git a/impy/utils/io.py b/impy/utils/io.py
index 9b4de895..b5951aa3 100644
--- a/impy/utils/io.py
+++ b/impy/utils/io.py
@@ -6,7 +6,7 @@
 import warnings
 import os
 import numpy as np
-from .._cupy import xp
+from ..array_api import xp
 
 from ..axes import ImageAxesError
 from .axesop import complement_axes
diff --git a/tests/test_fft.py b/tests/test_fft.py
index 4a5a8ba4..4c646b99 100644
--- a/tests/test_fft.py
+++ b/tests/test_fft.py
@@ -1,7 +1,7 @@
 import impy as ip
 from pathlib import Path
 import numpy as np
-from impy._cupy import xp
+from impy.array_api import xp
 from numpy.testing import assert_allclose
 ip.Const["SHOW_PROGRESS"] = False
 

From bd01d2acc65b2d34e6235216a4af340afc3aeeac Mon Sep 17 00:00:00 2001
From: hanjinliu <liuhanjin-sc@g.ecc.u-tokyo.ac.jp>
Date: Mon, 28 Feb 2022 14:59:06 +0900
Subject: [PATCH 07/13] change SetConst

---
 impy/_const.py | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/impy/_const.py b/impy/_const.py
index 68b75e6c..baed1339 100644
--- a/impy/_const.py
+++ b/impy/_const.py
@@ -1,3 +1,4 @@
+from __future__ import annotations
 import psutil
 from typing import Any, MutableMapping
 
@@ -84,19 +85,16 @@ def __repr__(self):
 
 class SetConst:
     n_ongoing = 0
-    def __init__(self, name=None, value=None, **kwargs):
-        if name is None and value is None and len(kwargs) == 1:
-            name, value = list(kwargs.items())[0]
-        elif name in Const.keys() and value is not None:
-            pass
-        else:
-            raise TypeError("Invalid input for SetConst")
-        self.name = name
-        self.value = value
+    def __init__(self, dict_: dict[str, Any] | None =None, **kwargs):
+        dict_ = dict_ or {}
+        dict_.update(kwargs)
+        self._kwargs = dict_
     
     def __enter__(self):
-        self.old_value = Const[self.name]
-        Const[self.name] = self.value
+        self._old_value = [(k, Const[k]) for k in self._kwargs.keys()]
+        for k, v in self._kwargs.items():
+            Const[k] = v
 
     def __exit__(self, exc_type, exc_value, traceback):
-        Const[self.name] = self.old_value
\ No newline at end of file
+        for k, v in self._old_value:
+            Const[k] = v

From 1b63aea2c91404ff6e38148a9f24a773dca45729 Mon Sep 17 00:00:00 2001
From: hanjinliu <liuhanjin-sc@g.ecc.u-tokyo.ac.jp>
Date: Mon, 28 Feb 2022 15:04:26 +0900
Subject: [PATCH 08/13] use cupy in random if available

---
 impy/random.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/impy/random.py b/impy/random.py
index d0c13b54..bb3989f5 100644
--- a/impy/random.py
+++ b/impy/random.py
@@ -2,15 +2,16 @@
 from functools import wraps
 import numpy as np
 from .arrays import ImgArray
+from .array_api import xp
 from .core import asarray
 
 def __getattr__(name: str):
-    npfunc = getattr(np.random, name)
-    @wraps(npfunc)
+    xpfunc = getattr(xp.random, name)
+    @wraps(xpfunc)
     def _func(*args, **kwargs) -> ImgArray:
-        name = kwargs.pop("name", npfunc.__name__)
+        name = kwargs.pop("name", xpfunc.__name__)
         axes = kwargs.pop("axes", None)
-        out = npfunc(*args, **kwargs)
+        out = xp.asnumpy(xpfunc(*args, **kwargs))
         return asarray(out, name=name, axes=axes)
     return _func
 
@@ -20,7 +21,7 @@ def random(size,
             name: str = None,
             axes: str = None) -> ImgArray:
     name = name or "random"
-    return asarray(np.random.random(size), name=name, axes=axes)
+    return asarray(xp.asnumpy(xp.random.random(size)), name=name, axes=axes)
 
 @wraps(np.random.normal)
 def normal(loc=0.0, 
@@ -30,7 +31,7 @@ def normal(loc=0.0,
            name: str = None, 
            axes: str = None) -> ImgArray:
     name = name or "normal"
-    return asarray(np.random.normal(loc, scale, size), name=name, axes=axes)
+    return asarray(xp.asnumpy(xp.random.normal(loc, scale, size)), name=name, axes=axes)
 
 def random_uint8(size: int | tuple[int], 
                  *, 
@@ -53,9 +54,9 @@ def random_uint8(size: int | tuple[int],
     ImgArray
         Random Image in dtype ``np.uint8``.
     """
-    arr = np.random.randint(0, 255, size, dtype=np.uint8)
+    arr = xp.random.randint(0, 255, size, dtype=np.uint8)
     name = name or "random_uint8"
-    return asarray(arr, name=name, axes=axes)
+    return asarray(xp.asnumpy(arr), name=name, axes=axes)
 
 def random_uint16(size,
                   *, 
@@ -78,7 +79,7 @@ def random_uint16(size,
     ImgArray
         Random Image in dtype ``np.uint16``.
     """
-    arr = np.random.randint(0, 65535, size, dtype=np.uint16)
+    arr = xp.random.randint(0, 65535, size, dtype=np.uint16)
     name = name or "random_uint16"
-    return asarray(arr, name=name, axes=axes)
+    return asarray(xp.asnumpy(arr), name=name, axes=axes)
 

From 1da8e3bfb59f97c7531ae6ceef493502072f2b66 Mon Sep 17 00:00:00 2001
From: hanjinliu <liuhanjin-sc@g.ecc.u-tokyo.ac.jp>
Date: Mon, 28 Feb 2022 15:23:47 +0900
Subject: [PATCH 09/13] better implementation for cupy

---
 impy/arrays/imgarray.py |  9 ++++++++-
 impy/correlation.py     | 29 ++++++++++++++++-------------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/impy/arrays/imgarray.py b/impy/arrays/imgarray.py
index e482ea2c..a29fced6 100644
--- a/impy/arrays/imgarray.py
+++ b/impy/arrays/imgarray.py
@@ -143,11 +143,18 @@ def affine(self, matrix=None, scale=None, rotation=None, shear=None, translation
                                            output_shape=output_shape)
                                )
     
+    @record
+    def map_coordinates(self, coords, *, mode="constant", cval: float = 0, order: int = 1):
+        ...
+        # TODO: Support multi-dimension
+        return xp.ndi.map_coordinates(self.values, coords, mode=mode, cval=cval, order=order, prefilter=order>1)
+        
+    
     @_docs.write_docs
     @dims_to_spatial_axes
     @same_dtype(asfloat=True)
     @record
-    def rotate(self, degree:float, center="center", *, mode="constant", cval: float=0, dims: Dims = 2,
+    def rotate(self, degree:float, center="center", *, mode="constant", cval: float = 0, dims: Dims = 2,
                order: int = 1, update: bool = False) -> ImgArray:
         """
         2D rotation of an image around a point. Outside will be padded with zero. For n-D images,
diff --git a/impy/correlation.py b/impy/correlation.py
index 6334568a..98f4ddbd 100644
--- a/impy/correlation.py
+++ b/impy/correlation.py
@@ -107,16 +107,12 @@ def _masked_ncc(img0: ImgArray, img1: ImgArray, dims: Dims, mask: ImgArray):
         np.ma.std(img0ma, axis=axis)*np.ma.std(img1ma, axis=axis)) / n
 
 
-def _zncc(img0: ImgArray, img1: ImgArray, dims: Dims):
+def _zncc(img0: xp.ndarray, img1: xp.ndarray, dims: Dims):
     # Basic Zero-Normalized Cross Correlation with batch processing.
     # Inputs must be already zero-normalized.
-    if isinstance(dims, str):
-        dims = tuple(img0.axisof(a) for a in dims)
-    img0 = xp.asarray(img0)
-    img1 = xp.asarray(img1)
     corr = xp.sum(img0 * img1, axis=dims) / (
         xp.sqrt(xp.sum(img0**2, axis=dims)*xp.sum(img1**2, axis=dims)))
-    return xp.asnumpy(corr)
+    return corr
 
 
 def _masked_zncc(img0: ImgArray, img1: ImgArray, dims: Dims, mask: ImgArray):
@@ -192,13 +188,16 @@ def zncc(
     """    
     with Progress("zncc"):
         img0, img1 = _check_inputs(img0, img1)
-        img0zn = img0 - np.mean(img0, axis=dims, keepdims=True)
-        img1zn = img1 - np.mean(img1, axis=dims, keepdims=True)
+        dims = tuple(img0.axisof(a) for a in dims)
+        img0 = xp.asarray(img0.value)
+        img1 = xp.asarray(img1.value)
+        img0zn = img0 - xp.mean(img0, axis=dims, keepdims=True)
+        img1zn = img1 - xp.mean(img1, axis=dims, keepdims=True)
         if mask is None:
             corr = _zncc(img0zn, img1zn, dims)
         else:
             corr = _masked_zncc(img0zn, img1zn, dims, mask)
-    return _make_corr_output(corr, img0, "zncc", squeeze, dims)
+    return _make_corr_output(xp.asnumpy(corr), img0, "zncc", squeeze, dims)
 
 # alias
 pearson_coloc = zncc
@@ -321,10 +320,14 @@ def fourier_zncc(
     """    
     with Progress("fourier_zncc"):
         img0, img1 = _check_inputs(img0, img1)
-        f0 = np.sqrt(img0.power_spectra(dims=dims, zero_norm=True))
-        f1 = np.sqrt(img1.power_spectra(dims=dims, zero_norm=True))
-        f0 -= np.mean(f0, axis=dims, keepdims=True)
-        f1 -= np.mean(f1, axis=dims, keepdims=True)
+        dims = tuple(img0.axisof(a) for a in dims)
+        pw0 = xp.asarray(img0.power_spectra(dims=dims, zero_norm=True).value)
+        pw1 = xp.asarray(img1.power_spectra(dims=dims, zero_norm=True).value)
+        
+        f0 = xp.sqrt(pw0)
+        f1 = xp.sqrt(pw1)
+        f0 -= xp.mean(f0, axis=dims, keepdims=True)
+        f1 -= xp.mean(f1, axis=dims, keepdims=True)
         if mask is None:
             corr = _zncc(f0, f1, dims)
         else:

From ca2bdf884ca54830c6667df77f2ef33da7d6141d Mon Sep 17 00:00:00 2001
From: hanjinliu <liuhanjin-sc@g.ecc.u-tokyo.ac.jp>
Date: Mon, 28 Feb 2022 15:39:18 +0900
Subject: [PATCH 10/13] map_coordinate

---
 impy/arrays/imgarray.py | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/impy/arrays/imgarray.py b/impy/arrays/imgarray.py
index a29fced6..0514ee66 100644
--- a/impy/arrays/imgarray.py
+++ b/impy/arrays/imgarray.py
@@ -143,11 +143,36 @@ def affine(self, matrix=None, scale=None, rotation=None, shear=None, translation
                                            output_shape=output_shape)
                                )
     
+    @_docs.write_docs
+    @same_dtype(asfloat=True)
     @record
-    def map_coordinates(self, coords, *, mode="constant", cval: float = 0, order: int = 1):
-        ...
+    def map_coordinates(self, coordinates, *, mode="constant", cval: float = 0, order: int = 1):
+        r"""
+        Coordinate mapping in the image. See ``scipy.ndimage.map_coordinates``.
+
+        Parameters
+        ----------
+        coordinates, mode, cval
+            Padding mode, constant value and the shape of output. See ``scipy.ndimage.map_coordinates``.
+            for details.
+        {order}
+
+        Returns
+        -------
+        ImgArray
+            Transformed image.
+        """
         # TODO: Support multi-dimension
-        return xp.ndi.map_coordinates(self.values, coords, mode=mode, cval=cval, order=order, prefilter=order>1)
+        return xp.asnumpy(
+            xp.ndi.map_coordinates(
+                xp.asarray(self.value), 
+                xp.asarray(coordinates), 
+                mode=mode,
+                cval=cval,
+                order=order,
+                prefilter=order>1
+            )
+        )
         
     
     @_docs.write_docs

From 3050d4a0881c6a19c3d254159c1fbcc7a91705f2 Mon Sep 17 00:00:00 2001
From: hanjinliu <liuhanjin-sc@g.ecc.u-tokyo.ac.jp>
Date: Mon, 28 Feb 2022 16:55:45 +0900
Subject: [PATCH 11/13] switch numpy and cupy in pytest

---
 impy/array_api.py              |  11 ++-
 impy/arrays/_utils/_corr.py    |  11 ++-
 impy/arrays/_utils/_filters.py |   9 +--
 impy/arrays/bases/metaarray.py |  21 ++---
 impy/arrays/imgarray.py        |  11 +--
 tests/conftest.py              |  10 +++
 tests/test_correlation.py      | 144 +++++++++++++++++----------------
 tests/test_fft.py              |  58 ++++++-------
 tests/test_lazy.py             |  34 ++++----
 tests/test_methods.py          |  11 +--
 10 files changed, 171 insertions(+), 149 deletions(-)
 create mode 100644 tests/conftest.py

diff --git a/impy/array_api.py b/impy/array_api.py
index ec77c062..14930b5b 100644
--- a/impy/array_api.py
+++ b/impy/array_api.py
@@ -9,7 +9,7 @@ def func(*args, **kwargs):
         if xp.state == "cupy":
             args = (xp.asarray(a) if isinstance(a, np.ndarray) else a for a in args)
         out = function(*args, **kwargs)
-        return out
+        return xp.asnumpy(out)
     return func
 
 from types import ModuleType
@@ -76,6 +76,9 @@ def setNumpy(self) -> None:
         self.gradient = np.gradient
         self.tensordot = np.tensordot
         self.stack = np.stack
+        self.unravel_index = np.unravel_index
+        self.argmax = np.argmax
+        self.argmin = np.argmin
         
         self.state = "numpy"
         from ._const import Const
@@ -91,7 +94,6 @@ def cp_asnumpy(arr, dtype=None):
         from cupyx.scipy import fft as cp_fft
         from cupyx.scipy import ndimage as cp_ndi
         from cupy import linalg as cp_linalg
-        from cupy import ndarray as cp_ndarray
         
         self._module = cp
         self.fft = cp_fft
@@ -100,7 +102,7 @@ def cp_asnumpy(arr, dtype=None):
         self.ndi = cp_ndi
         self.asnumpy = cp_asnumpy
         self.asarray = cp.asarray
-        self.ndarray = cp_ndarray
+        self.ndarray = cp.ndarray
         self.empty = cp.empty
         self.zeros = cp.zeros
         self.ones = cp.ones
@@ -134,6 +136,9 @@ def cp_asnumpy(arr, dtype=None):
             self.gradient = _gradient
         self.tensordot = cp.tensordot
         self.stack = cp.stack
+        self.unravel_index = cp.unravel_index
+        self.argmax = cp.argmax
+        self.argmin = cp.argmin
         self.state = "cupy"
         
         from ._const import Const
diff --git a/impy/arrays/_utils/_corr.py b/impy/arrays/_utils/_corr.py
index 2983d182..23ec8463 100644
--- a/impy/arrays/_utils/_corr.py
+++ b/impy/arrays/_utils/_corr.py
@@ -21,7 +21,7 @@ def subpixel_pcc(
     maxima = xp.unravel_index(xp.argmax(power), power.shape)
     midpoints = xp.array([np.fix(axis_size / 2) for axis_size in power.shape])
 
-    shifts = xp.asarray(maxima, dtype=xp.float32)
+    shifts = xp.asarray(maxima, dtype=np.float32)
     shifts[shifts > midpoints] -= xp.array(power.shape)[shifts > midpoints]
     # Initial shift estimate in upsampled grid
     shifts = xp.fix(shifts * upsample_factor) / upsample_factor
@@ -29,7 +29,7 @@ def subpixel_pcc(
         upsampled_region_size = np.ceil(upsample_factor * 1.5)
         # Center of output array at dftshift + 1
         dftshift = xp.fix(upsampled_region_size / 2.0)
-        upsample_factor = xp.array(upsample_factor, dtype=xp.float32)
+        upsample_factor = xp.array(upsample_factor, dtype=np.float32)
         # Matrix multiply DFT around the current shift estimate
         sample_region_offset = dftshift - shifts*upsample_factor
         cross_correlation = _upsampled_dft(product.conj(),
@@ -46,8 +46,7 @@ def subpixel_pcc(
             power = crop_by_max_shifts(power, _upsampled_left_shifts, _upsampled_right_shifts)
             
         maxima = xp.unravel_index(xp.argmax(power), power.shape)
-        maxima = xp.asarray(maxima, dtype=xp.float32) - dftshift
-
+        maxima = xp.asarray(maxima, dtype=np.float32) - dftshift
         shifts = shifts + maxima / upsample_factor
     return shifts
 
@@ -63,9 +62,9 @@ def _upsampled_dft(
     dim_properties = list(zip(data.shape, upsampled_region_size, axis_offsets))
 
     for (n_items, ups_size, ax_offset) in dim_properties[::-1]:
-        kernel = ((xp.arange(ups_size) - ax_offset)[:, xp.newaxis]
+        kernel = ((xp.arange(ups_size) - ax_offset)[:, np.newaxis]
                   * xp.fft.fftfreq(n_items, upsample_factor))
-        kernel = xp.exp(-2j * xp.pi * kernel)
+        kernel = xp.exp(-2j * np.pi * kernel)
 
         data = xp.tensordot(kernel, data, axes=(1, -1))
     return data
diff --git a/impy/arrays/_utils/_filters.py b/impy/arrays/_utils/_filters.py
index 42c352a9..0c6fb73a 100644
--- a/impy/arrays/_utils/_filters.py
+++ b/impy/arrays/_utils/_filters.py
@@ -34,11 +34,10 @@
 from scipy import ndimage as scipy_ndi
 
 def get_func(function_name):
-    if hasattr(xp.ndi, function_name):
-        _func = getattr(xp.ndi, function_name)    
-        func = cupy_dispatcher(_func)
-    else:
-        func = getattr(scipy_ndi, function_name)
+    def func(*args, **kwargs):
+        _f = getattr(xp.ndi, function_name, getattr(scipy_ndi, function_name))
+        return cupy_dispatcher(_f)(*args, **kwargs)
+    func.__name__ = function_name
     return func
 
 binary_erosion = get_func("binary_erosion")
diff --git a/impy/arrays/bases/metaarray.py b/impy/arrays/bases/metaarray.py
index 8825b7dd..214e093c 100644
--- a/impy/arrays/bases/metaarray.py
+++ b/impy/arrays/bases/metaarray.py
@@ -261,16 +261,17 @@ def sort_axes(self) -> MetaArray:
         return self.transpose(order)
 
     
-    def apply_dask(self, 
-                   func: Callable,
-                   c_axes: str | None = None,
-                   drop_axis: Iterable[int] = [], 
-                   new_axis: Iterable[int] = None, 
-                   dtype = np.float32, 
-                   out_chunks: tuple[int, ...] = None,
-                   args: tuple[Any] = None,
-                   kwargs: dict[str, Any] = None
-                   ) -> MetaArray:
+    def apply_dask(
+        self, 
+        func: Callable,
+        c_axes: str | None = None,
+        drop_axis: Iterable[int] = [], 
+        new_axis: Iterable[int] = None, 
+        dtype = np.float32, 
+        out_chunks: tuple[int, ...] = None,
+        args: tuple[Any] = None,
+        kwargs: dict[str, Any] = None
+    ) -> MetaArray:
         """
         Convert array into dask array and run a batch process in parallel. In many cases batch process 
         in this way is faster than `multiprocess` module.
diff --git a/impy/arrays/imgarray.py b/impy/arrays/imgarray.py
index 0514ee66..40fae6a0 100644
--- a/impy/arrays/imgarray.py
+++ b/impy/arrays/imgarray.py
@@ -147,7 +147,7 @@ def affine(self, matrix=None, scale=None, rotation=None, shear=None, translation
     @same_dtype(asfloat=True)
     @record
     def map_coordinates(self, coordinates, *, mode="constant", cval: float = 0, order: int = 1):
-        r"""
+        """
         Coordinate mapping in the image. See ``scipy.ndimage.map_coordinates``.
 
         Parameters
@@ -1741,10 +1741,11 @@ def rof_filter(self, lmd: float = 0.05, tol: float = 1e-4, max_iter: int = 50, *
         ImgArray
             Filtered image
         """        
-        return self.apply_dask(skres._denoise._denoise_tv_chambolle_nd, 
-                               c_axes=complement_axes(dims, self.axes),
-                               kwargs=dict(weight=lmd, eps=tol, n_iter_max=max_iter)
-                               )
+        return self.apply_dask(
+            skres._denoise._denoise_tv_chambolle_nd, 
+            c_axes=complement_axes(dims, self.axes),
+            kwargs=dict(weight=lmd, eps=tol, max_num_iter=max_iter)
+        )
         
     @_docs.write_docs
     @dims_to_spatial_axes
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..a042b69f
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,10 @@
+import pytest
+
+def pytest_addoption(parser):
+    parser.addoption("--resource", default="numpy")
+
+@pytest.fixture(scope="session")
+def resource(request):
+    import impy as ip
+    res = request.config.getoption('--resource')
+    return res
diff --git a/tests/test_correlation.py b/tests/test_correlation.py
index 442a06a2..9c59747d 100644
--- a/tests/test_correlation.py
+++ b/tests/test_correlation.py
@@ -3,83 +3,85 @@
 import numpy as np
 from numpy.testing import assert_allclose
 
-def test_pcc():
-    reference_image = ip.sample_image("camera")
-    shift = (-7, 12)
-    shifted_image = reference_image.affine(translation=shift)
+def test_pcc(resource):
+    with ip.SetConst(RESOURCE=resource):
+        reference_image = ip.sample_image("camera")
+        shift = (-7, 12)
+        shifted_image = reference_image.affine(translation=shift)
 
-    shift_sk, _, _ = phase_cross_correlation(shifted_image, reference_image)
-    shift_ip = ip.pcc_maximum(shifted_image, reference_image)
-    assert_allclose(shift_sk, shift_ip)
-    assert_allclose(shift_sk, (7, -12))
+        shift_sk, _, _ = phase_cross_correlation(shifted_image, reference_image)
+        shift_ip = ip.pcc_maximum(shifted_image, reference_image)
+        assert_allclose(shift_sk, shift_ip)
+        assert_allclose(shift_sk, (7, -12))
 
 
-def test_fourier():
-    reference_image = ip.sample_image("camera")
-    shift = (-7, 12)
-    shifted_image = reference_image.affine(translation=shift)
-    
-    reference_image_ft = reference_image.fft()
-    shifted_image_ft = shifted_image.fft()
-    
-    shift_sk, _, _ = phase_cross_correlation(shifted_image_ft, reference_image_ft, space="fourier")
-    shift_ip = ip.ft_pcc_maximum(shifted_image_ft, reference_image_ft)
-    assert_allclose(shift_sk, shift_ip)
-    assert_allclose(shift_sk, (7, -12))
+def test_fourier(resource):
+    with ip.SetConst(RESOURCE=resource):
+        reference_image = ip.sample_image("camera")
+        shift = (-7, 12)
+        shifted_image = reference_image.affine(translation=shift)
+        
+        reference_image_ft = reference_image.fft()
+        shifted_image_ft = shifted_image.fft()
+        
+        shift_sk, _, _ = phase_cross_correlation(shifted_image_ft, reference_image_ft, space="fourier")
+        shift_ip = ip.ft_pcc_maximum(shifted_image_ft, reference_image_ft)
+        assert_allclose(shift_sk, shift_ip)
+        assert_allclose(shift_sk, (7, -12))
 
-def test_max_shift():
-    # check shifts don't exceed max_shifts
-    for i in range(10):
-        np.random.seed(i)
-        ref = ip.random.random_uint16((128, 129))
-        img = ref.affine(translation=[30, -44])
-        shift = ip.pcc_maximum(img, ref, max_shifts=20)
-        assert all(shift <= 20)
-        shift = ip.pcc_maximum(img, ref, max_shifts=14.6)
-        assert all(shift <= 14.6)
-    
-    # check shifts are correct if max_shifts is large enough
-    reference_image = ip.sample_image("camera")
-    shift = (-7, 12)
-    shifted_image = reference_image.affine(translation=shift)
+def test_max_shift(resource):
+    with ip.SetConst(RESOURCE=resource):
+        # check shifts don't exceed max_shifts
+        for i in range(10):
+            np.random.seed(i)
+            ref = ip.random.random_uint16((128, 129))
+            img = ref.affine(translation=[30, -44])
+            shift = ip.pcc_maximum(img, ref, max_shifts=20)
+            assert all(shift <= 20)
+            shift = ip.pcc_maximum(img, ref, max_shifts=14.6)
+            assert all(shift <= 14.6)
+        
+        # check shifts are correct if max_shifts is large enough
+        reference_image = ip.sample_image("camera")
+        shift = (-7, 12)
+        shifted_image = reference_image.affine(translation=shift)
 
-    shift = ip.pcc_maximum(shifted_image, reference_image, max_shifts=15.7)
-    assert_allclose(shift, shift)
-    
-    # check shifts are correct even if at the edge of max_shifts
-    reference_image = ip.sample_image("camera")
-    shift = (-7.8, 6.6)
-    shifted_image = reference_image.affine(translation=shift)
+        shift = ip.pcc_maximum(shifted_image, reference_image, max_shifts=15.7)
+        assert_allclose(shift, shift)
+        
+        # check shifts are correct even if at the edge of max_shifts
+        reference_image = ip.sample_image("camera")
+        shift = (-7.8, 6.6)
+        shifted_image = reference_image.affine(translation=shift)
 
-    shift = ip.pcc_maximum(shifted_image, reference_image, max_shifts=[7.9, 6.7])
-    assert_allclose(shift, shift)
-    
-    # check sub-optimal shifts will be returned
-    ref = ip.zeros((128, 128))
-    ref[10, 10] = 1
-    ref[10, 20] = 1
-    ref0 = ip.zeros((128, 128))
-    ref0[10, 20] = 1
-    img = ip.zeros((128, 128))
-    img[12, 25] = 1
-    img[12, 35] = 1
-    
-    shift = ip.pcc_maximum(img, ref)
-    assert_allclose(shift, (2, 15))
-    shift0 = ip.pcc_maximum(img, ref0)
-    shift = ip.pcc_maximum(img, ref, max_shifts=[5.7, 10], upsample_factor=2)
-    assert_allclose(shift, shift0)
+        shift = ip.pcc_maximum(shifted_image, reference_image, max_shifts=[7.9, 6.7])
+        assert_allclose(shift, shift)
+        
+        # check sub-optimal shifts will be returned
+        ref = ip.zeros((128, 128))
+        ref[10, 10] = 1
+        ref[10, 20] = 1
+        img = ip.zeros((128, 128))
+        img[12, 25] = 1
+        img[12, 35] = 1
+        
+        shift = ip.pcc_maximum(img, ref)
+        assert_allclose(shift, (2, 15))
+        shift = ip.pcc_maximum(img, ref, max_shifts=[5.7, 10], upsample_factor=2)
+        assert_allclose(shift, [2, 5])
     
-def test_polar_pcc():
-    reference_image = ip.sample_image("camera")
-    deg = 21
-    rotated_image = reference_image.rotate(deg)
+def test_polar_pcc(resource):
+    with ip.SetConst(RESOURCE=resource):
+        reference_image = ip.sample_image("camera")
+        deg = 21
+        rotated_image = reference_image.rotate(deg)
 
-    rot = ip.polar_pcc_maximum(rotated_image, reference_image)
-    assert rot == deg
+        rot = ip.polar_pcc_maximum(rotated_image, reference_image)
+        assert rot == deg
 
-def test_fsc():
-    img0 = ip.random.random_uint16((80, 80))
-    img1 = ip.random.random_uint16((80, 80))
-    a, b = ip.fsc(img0, img1)
-    assert a.size == b.size
\ No newline at end of file
+def test_fsc(resource):
+    with ip.SetConst(RESOURCE=resource):
+        img0 = ip.random.random_uint16((80, 80))
+        img1 = ip.random.random_uint16((80, 80))
+        a, b = ip.fsc(img0, img1)
+        assert a.size == b.size
\ No newline at end of file
diff --git a/tests/test_fft.py b/tests/test_fft.py
index 4c646b99..6c2769c8 100644
--- a/tests/test_fft.py
+++ b/tests/test_fft.py
@@ -5,32 +5,34 @@
 from numpy.testing import assert_allclose
 ip.Const["SHOW_PROGRESS"] = False
 
-def test_precision():
-    path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
-    img = ip.imread(path)["c=1;t=0"].as_float()
-    
-    assert_allclose(img.fft(shift=False),
-                    xp.asnumpy(xp.fft.fftn(xp.asarray(img.value)))
-                    )
-    
-    assert_allclose(img.fft(shift=False, double_precision=True),
-                    xp.asnumpy(xp.fft.fftn(xp.asarray(img.value.astype(np.float64))).astype(np.complex64))
-                    )
-    
-    assert_allclose(img.fft().ifft(), img, rtol=1e-6)
-    assert_allclose(img.fft(double_precision=True).ifft(double_precision=True), img)
-    
-    assert_allclose(np.fft.fftshift(img.local_dft()).ifft(), img, rtol=1e-6)
-    assert_allclose(np.fft.fftshift(img.local_dft(double_precision=True)).ifft(double_precision=True), img)
+def test_precision(resource):
+    with ip.SetConst(RESOURCE=resource):
+        path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
+        img = ip.imread(path)["c=1;t=0"].as_float()
+        
+        assert_allclose(img.fft(shift=False),
+                        xp.asnumpy(xp.fft.fftn(xp.asarray(img.value)))
+                        )
+        
+        assert_allclose(img.fft(shift=False, double_precision=True),
+                        xp.asnumpy(xp.fft.fftn(xp.asarray(img.value.astype(np.float64))).astype(np.complex64))
+                        )
+        
+        assert_allclose(img.fft().ifft(), img, rtol=1e-6)
+        assert_allclose(img.fft(double_precision=True).ifft(double_precision=True), img)
+        
+        assert_allclose(np.fft.fftshift(img.local_dft()).ifft(), img, rtol=1e-6)
+        assert_allclose(np.fft.fftshift(img.local_dft(double_precision=True)).ifft(double_precision=True), img)
 
-def test_iteration():
-    path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
-    img = ip.imread(path)["c=1;t=0"].as_float()
-    
-    ft0 = img.fft(dims="zx")
-    ft1 = np.stack([img[f"y={i}"].fft(dims="zx") for i in range(img.shape.y)], axis="y")
-    assert_allclose(ft0, ft1)
-    
-    ft0 = img.ifft(dims="zx")
-    ft1 = np.stack([img[f"y={i}"].ifft(dims="zx") for i in range(img.shape.y)], axis="y")
-    assert_allclose(ft0, ft1)
\ No newline at end of file
+def test_iteration(resource):
+    with ip.SetConst(RESOURCE=resource):
+        path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
+        img = ip.imread(path)["c=1;t=0"].as_float()
+        
+        ft0 = img.fft(dims="zx")
+        ft1 = np.stack([img[f"y={i}"].fft(dims="zx") for i in range(img.shape.y)], axis="y")
+        assert_allclose(ft0, ft1)
+        
+        ft0 = img.ifft(dims="zx")
+        ft1 = np.stack([img[f"y={i}"].ifft(dims="zx") for i in range(img.shape.y)], axis="y")
+        assert_allclose(ft0, ft1)
\ No newline at end of file
diff --git a/tests/test_lazy.py b/tests/test_lazy.py
index 27b99ac7..8bf6b175 100644
--- a/tests/test_lazy.py
+++ b/tests/test_lazy.py
@@ -11,24 +11,26 @@
            "kalman_filter", "gaussian_filter"]
 
 
-def test_functions_and_slicing():
-    path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
-    img = ip.lazy_imread(path, chunks=(4, 5, 2, 32, 32))
-    sl = "y=20:40;x=30:50;c=0;z=2,4"
-    assert_allclose(img[sl].compute(), img.compute()[sl])
-    assert_allclose(img.affine(translation=[1, 50, 50]).compute(),
-                    img.compute().affine(translation=[1, 50, 50])
-                    )
+def test_functions_and_slicing(resource):
+    with ip.SetConst(RESOURCE=resource):
+        path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
+        img = ip.lazy_imread(path, chunks=(4, 5, 2, 32, 32))
+        sl = "y=20:40;x=30:50;c=0;z=2,4"
+        assert_allclose(img[sl].compute(), img.compute()[sl])
+        assert_allclose(img.affine(translation=[1, 50, 50]).compute(),
+                        img.compute().affine(translation=[1, 50, 50])
+                        )
 
 @pytest.mark.parametrize("fn", filters)
-def test_filters(fn):
-    path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
-    img = ip.lazy_imread(path, chunks=(4, 5, 2, 32, 32))
-    
-    assert_allclose(getattr(img, fn)().compute(),
-                    getattr(img.compute(), fn)()
-                    )
-    
+def test_filters(fn, resource):
+    with ip.SetConst(RESOURCE=resource):
+        path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
+        img = ip.lazy_imread(path, chunks=(4, 5, 2, 32, 32))
+        
+        assert_allclose(getattr(img, fn)().compute(),
+                        getattr(img.compute(), fn)()
+                        )
+        
 
 def test_numpy_function():
     img = ip.aslazy(ip.random.random_uint16((10, 100, 100)))
diff --git a/tests/test_methods.py b/tests/test_methods.py
index aecd424f..1714c351 100644
--- a/tests/test_methods.py
+++ b/tests/test_methods.py
@@ -38,11 +38,12 @@
 
 @pytest.mark.parametrize("f", filters)
 @pytest.mark.parametrize("dtype", dtypes)
-def test_filters(f, dtype):
-    path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
-    img = ip.imread(path, dtype=dtype)["c=1;z=2"]
-    assert img.axes == "tyx"
-    getattr(img, f)()
+def test_filters(f, dtype, resource):
+    with ip.SetConst(RESOURCE=resource):
+        path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"
+        img = ip.imread(path, dtype=dtype)["c=1;z=2"]
+        assert img.axes == "tyx"
+        getattr(img, f)()
 
 def test_sm():
     path = Path(__file__).parent / "_test_images" / "image_tzcyx.tif"

From c03976bf0df2ae6bb5708744105f5bad4e4deb15 Mon Sep 17 00:00:00 2001
From: hanjinliu <liuhanjin-sc@g.ecc.u-tokyo.ac.jp>
Date: Mon, 28 Feb 2022 17:02:58 +0900
Subject: [PATCH 12/13] cupy bug fixes

---
 impy/arrays/_utils/_corr.py | 2 +-
 impy/correlation.py         | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/impy/arrays/_utils/_corr.py b/impy/arrays/_utils/_corr.py
index 23ec8463..a5f54603 100644
--- a/impy/arrays/_utils/_corr.py
+++ b/impy/arrays/_utils/_corr.py
@@ -15,7 +15,7 @@ def subpixel_pcc(
     cross_correlation = xp.fft.ifftn(product)
     power = abs2(cross_correlation)
     if max_shifts is not None:
-        max_shifts = np.asarray(max_shifts)
+        max_shifts = xp.asarray(max_shifts)
         power = crop_by_max_shifts(power, max_shifts, max_shifts)
         
     maxima = xp.unravel_index(xp.argmax(power), power.shape)
diff --git a/impy/correlation.py b/impy/correlation.py
index 98f4ddbd..479473f1 100644
--- a/impy/correlation.py
+++ b/impy/correlation.py
@@ -453,8 +453,8 @@ def polar_pcc_maximum(
         max_degree = 180
     rmax = min(img0.shape)
     with Progress("polar_pcc_maximum_2d"):
-        imgp = ip_asarray(polar2d(img0, rmax, np.pi/180))
-        imgrotp = ip_asarray(polar2d(img1, rmax, np.pi/180))
+        imgp = ip_asarray(xp.asnumpy(polar2d(img0, rmax, np.pi/180)))
+        imgrotp = ip_asarray(xp.asnumpy(polar2d(img1, rmax, np.pi/180)))
         max_shifts = (max_degree, 1)
         shift = pcc_maximum(imgp, imgrotp, upsample_factor=upsample_factor,
                             max_shifts=max_shifts)

From 1fcf2cc22ffbc8786297f4ae21e1ac7573d80412 Mon Sep 17 00:00:00 2001
From: hanjinliu <liuhanjin-sc@g.ecc.u-tokyo.ac.jp>
Date: Mon, 28 Feb 2022 17:28:35 +0900
Subject: [PATCH 13/13] update to v1.26.0

---
 impy/__init__.py        |  5 ++---
 impy/_const.py          | 27 +++++++++++++++++++++++++++
 impy/arrays/imgarray.py |  2 +-
 tests/test_methods.py   |  1 -
 4 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/impy/__init__.py b/impy/__init__.py
index e6a13126..f4b5b3a8 100644
--- a/impy/__init__.py
+++ b/impy/__init__.py
@@ -1,11 +1,11 @@
-__version__ = "1.25.3.dev0"
+__version__ = "1.26.0"
 __author__ = "Hanjin Liu"
 __email__ = "liuhanjin-sc@g.ecc.u-tokyo.ac.jp"
 
 import logging
 from functools import wraps
 
-from ._const import Const, SetConst
+from ._const import Const, SetConst, silent, use
 
 from .collections import *
 from .core import *
@@ -14,7 +14,6 @@
 from .correlation import *
 from .arrays import ImgArray, LazyImgArray  # for typing
 from . import random
-import numpy as np
 
 r"""
 Inheritance
diff --git a/impy/_const.py b/impy/_const.py
index ef950290..286580bf 100644
--- a/impy/_const.py
+++ b/impy/_const.py
@@ -98,3 +98,30 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         for k, v in self._old_value:
             Const[k] = v
+
+def silent():
+    """
+    Do not show progress in this context.
+    
+    An alias of ``ip.SetConst(SHOW_PROGRESS=False)``
+    """
+    return SetConst(SHOW_PROGRESS=False)
+
+def use(resource, import_error: bool = False):
+    """
+    Use a resource (numpy or cupy) in this context.
+
+    Parameters
+    ----------
+    resource : str
+        Resource to use.
+    import_error : bool, default is False
+        If false, resource will not be switched to cupy if not available.
+        Raise ImportError if true.
+    """
+    if not import_error:
+        try:
+            import cupy
+        except ImportError:
+            resource = "numpy"
+    return SetConst(RESOURCE=resource)
diff --git a/impy/arrays/imgarray.py b/impy/arrays/imgarray.py
index 40fae6a0..19c5ce6d 100644
--- a/impy/arrays/imgarray.py
+++ b/impy/arrays/imgarray.py
@@ -1740,7 +1740,7 @@ def rof_filter(self, lmd: float = 0.05, tol: float = 1e-4, max_iter: int = 50, *
         -------
         ImgArray
             Filtered image
-        """        
+        """     
         return self.apply_dask(
             skres._denoise._denoise_tv_chambolle_nd, 
             c_axes=complement_axes(dims, self.axes),
diff --git a/tests/test_methods.py b/tests/test_methods.py
index 1714c351..436b61ec 100644
--- a/tests/test_methods.py
+++ b/tests/test_methods.py
@@ -31,7 +31,6 @@
            "doh_filter",
            "log_filter",
            "rolling_ball",
-           "rof_filter",
            ]
 
 dtypes = [np.uint8, np.uint16, np.float32]