fix cupy (#38)

* fix cupy

* style(pre-commit.ci): auto fixes [...]

* fix cupy random

* fix jax and cupy

* style(pre-commit.ci): auto fixes [...]

* fix lint

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: tlambert03

src/microsim/psf.py

@@ -157,8 +157,9 @@ def rz_to_xyz(
     xp = NumpyAPI.create(xp)
 
     # Create XY grid of radius values.
-    rmap = radius_map(xyshape, off) * sf
+    rmap = radius_map(xyshape, off, xp=xp) * sf
     nz = rz.shape[0]
+
     out = xp.asarray(
         [
             xp.map_coordinates(
@@ -201,7 +202,7 @@ def vectorial_psf(
     ).astype(xp.float_dtype)
 
     offsets = xp.asarray(pos[:2]) / (dxy * 1e-6)
-    _psf = rz_to_xyz(rz, (ny, nx), sf, off=offsets)  # type: ignore [arg-type]
+    _psf = rz_to_xyz(rz, (ny, nx), sf, off=offsets, xp=xp)  # type: ignore [arg-type]
     if normalize:
         _psf /= xp.max(_psf)
     return _psf

src/microsim/schema/backend.py

@@ -14,6 +14,8 @@
 
     import jax
 
+    from microsim._data_array import ArrayProtocol
+
     _Shape: TypeAlias = tuple[int, ...]
 
     # Anything that can be coerced to a shape tuple
@@ -47,18 +49,19 @@ def create(cls, backend: BackendName | NumpyAPI | None) -> NumpyAPI:
 
         return NumpyAPI()
 
+    _random_seed: int | None = None
+    _float_dtype: np.dtype | None = None
+
     def __init__(self) -> None:
         from scipy import signal, special, stats
         from scipy.ndimage import map_coordinates
 
-        self._random_seed: int | None = None
         self.xp = np
         self.signal = signal
         self.stats = stats
         self.j0 = special.j0
         self.j1 = special.j1
         self.map_coordinates = map_coordinates
-        self._float_dtype: np.dtype | None = None
 
     @property
     def float_dtype(self) -> np.dtype | None:
@@ -103,6 +106,11 @@ def poisson_rvs(
     ) -> npt.NDArray:
         return self.stats.poisson.rvs(lam, size=shape)  # type: ignore
 
+    def norm_rvs(
+        self, loc: ArrayProtocol, scale: npt.ArrayLike | None = None
+    ) -> ArrayProtocol:
+        return self.stats.norm.rvs(loc, scale)  # type: ignore
+
     def fftconvolve(
         self, a: ArrT, b: ArrT, mode: Literal["full", "valid", "same"] = "full"
     ) -> ArrT:
@@ -140,7 +148,6 @@ def __init__(self) -> None:
 
         from ._jax_bessel import j0, j1
 
-        self._random_seed: int | None = None
         self.xp = jax.numpy
         self.signal = signal
         self.stats = stats
@@ -173,6 +180,15 @@ def poisson_rvs(  # type: ignore
 
         return poisson(self._key, lam, shape=shape)
 
+    def norm_rvs(
+        self, loc: ArrayProtocol, scale: npt.ArrayLike | None = None
+    ) -> ArrayProtocol:
+        from jax.random import normal
+
+        std_samples = normal(self._key, shape=loc.shape)
+        # scale and shift
+        return std_samples * scale + loc  # type: ignore
+
     def fftconvolve(
         self, a: ArrT, b: ArrT, mode: Literal["full", "valid", "same"] = "full"
     ) -> ArrT:
@@ -211,6 +227,16 @@ def __init__(self) -> None:
         self.j1 = special.j1
         self.map_coordinates = map_coordinates
 
+    def poisson_rvs(
+        self, lam: npt.ArrayLike, shape: Sequence[int] | None = None
+    ) -> npt.NDArray:
+        return self.xp.random.poisson(lam, shape)  # type: ignore
+
+    def norm_rvs(
+        self, loc: ArrayProtocol, scale: npt.ArrayLike | None = None
+    ) -> ArrayProtocol:
+        return self.xp.random.normal(loc, scale)  # type: ignore
+
     def fftconvolve(
         self, a: ArrT, b: ArrT, mode: Literal["full", "valid", "same"] = "full"
     ) -> ArrT:

src/microsim/schema/detectors/_camera.py

@@ -81,9 +81,11 @@ def adc_gain(self) -> float:
     def max_intensity(self) -> int:
         return int(2**self.bit_depth - 1)
 
-    def quantize_electrons(self, total_electrons: npt.NDArray) -> npt.NDArray:
-        voltage = stats.norm.rvs(total_electrons, self.read_noise) * self.gain
-        return np.round((voltage / self.adc_gain) + self.offset)  # type: ignore
+    def quantize_electrons(
+        self, total_electrons: npt.NDArray, xp: NumpyAPI
+    ) -> npt.NDArray:
+        voltage = xp.norm_rvs(total_electrons, self.read_noise) * self.gain
+        return xp.round((voltage / self.adc_gain) + self.offset)  # type: ignore
 
 
 class CameraCCD(Camera): ...

src/microsim/schema/detectors/_simulate.py

@@ -74,7 +74,7 @@ def simulate_camera(
         total_electrons = camera.apply_em_gain(total_electrons)
 
     # model read noise
-    gray_values = camera.quantize_electrons(total_electrons)
+    gray_values = camera.quantize_electrons(total_electrons, xp)
 
     # sCMOS binning
     if binning > 1 and isinstance(camera, CameraCMOS):

src/microsim/util.py

@@ -220,6 +220,8 @@ def ortho_plot(
     import matplotlib.pyplot as plt
     from matplotlib.colors import PowerNorm
 
+    if hasattr(img, "get"):
+        img = img.get()
     img = np.asarray(img)
     """Plot XY and XZ slices of a 3D array."""
     _, ax = plt.subplots(ncols=2, figsize=(10, 5))

tests/test_simulation.py

@@ -67,7 +67,7 @@ def test_schema(
     assert type(out1.data).__module__.split(".")[0].startswith(np_backend)
 
     out2 = sim1.run()
-    if seed is None:
+    if seed is None and np_backend != "jax":
         assert not np.allclose(out1, out2)
     else:
         np.testing.assert_allclose(out1, out2)