From 4802e50939edfca8ba9585bbf97f71560e4eff09 Mon Sep 17 00:00:00 2001 From: Ashish Tripathi Date: Fri, 5 Apr 2024 18:51:10 -0500 Subject: [PATCH 1/5] NEW: Add Fresnel Propagation Operator Co-authored-by: Ashish Tripathi --- src/tike/operators/cupy/__init__.py | 1 + src/tike/operators/cupy/fresnelspectprop.py | 143 ++++++++++++++++++++ tests/operators/test_propagation.py | 25 +++- 3 files changed, 167 insertions(+), 2 deletions(-) create mode 100644 src/tike/operators/cupy/fresnelspectprop.py diff --git a/src/tike/operators/cupy/__init__.py b/src/tike/operators/cupy/__init__.py index cb2a9d6f..787d2751 100644 --- a/src/tike/operators/cupy/__init__.py +++ b/src/tike/operators/cupy/__init__.py @@ -16,6 +16,7 @@ from .pad import * from .patch import * from .propagation import * +from .fresnelspectprop import * from .ptycho import * from .rotate import * from .shift import * diff --git a/src/tike/operators/cupy/fresnelspectprop.py b/src/tike/operators/cupy/fresnelspectprop.py new file mode 100644 index 00000000..d973ea0f --- /dev/null +++ b/src/tike/operators/cupy/fresnelspectprop.py @@ -0,0 +1,143 @@ +"""Defines a free-space propagation operator based on the CuPy FFT module.""" + +__author__ = "Ashish Tripathi" +__copyright__ = "Copyright (c) 2024, UChicago Argonne, LLC." + +import typing + +import numpy.typing as npt +import numpy as np + +from .cache import CachedFFT +from .operator import Operator + + +class FresnelSpectProp(CachedFFT, Operator): + """Fresnel spectrum propagation (short range) using CuPy. + + Take an (..., N, N) array and apply the Fourier transform to the last two + dimensions. + + Attributes + ---------- + pixel_size : float + The realspace size of a pixel in meters + delta_z : float + The realspace propagation distance in meters + wavelength : float + The wavelength of the light in meters + + Parameters + ---------- + nearplane: (..., detector_shape, detector_shape) complex64 + The wavefronts before propagation. + farplane: (..., detector_shape, detector_shape) complex64 + The wavefronts after propagation. + """ + + def __init__( + self, + norm: str = "ortho", + pixel_size: float = 1.0, + delta_z: float = 1.0, + wavelength: float = 1.0, + **kwargs, + ): + self.norm = norm + self.pixel_size = pixel_size + self.delta_z = delta_z + self.wavelength = wavelength + + def fwd( + self, + nearplane: npt.NDArray[np.csingle], + overwrite: bool = False, + **kwargs, + ) -> npt.NDArray[np.csingle]: + """forward (parallel to beam direction) Fresnel spectrum propagtion operator""" + propagator = self._create_fresnel_spectrum_propagator( + (nearplane.shape[-2], nearplane.shape[-1]), + self.pixel_size, + self.delta_z, + self.wavelength, + ) + + nearplane_fft2 = self._fft2( + nearplane, + norm=self.norm, + axes=(-2, -1), + overwrite_x=overwrite, + ) + + farplane = self._ifft2( + nearplane_fft2 * propagator, + norm=self.norm, + axes=(-2, -1), + overwrite_x=overwrite, + ) + + return farplane + + def adj( + self, + farplane: npt.NDArray[np.csingle], + overwrite: bool = False, + **kwargs, + ) -> npt.NDArray[np.csingle]: + """backward (anti-parallel to beam direction) Fresnel spectrum propagtion operator""" + propagator = self._create_fresnel_spectrum_propagator( + (farplane.shape[-2], farplane.shape[-1]), + self.pixel_size, + self.delta_z, + self.wavelength, + ) + + farplane_fft2 = self._fft2( + farplane, + norm=self.norm, + axes=(-2, -1), + overwrite_x=overwrite, + ) + + nearplane = self._ifft2( + farplane_fft2 + * self.xp.conj( + propagator, + ), # IS IT OK TO ALWAYS TAKE CONJ? OR SHOULD WE DO THIS ONCE AND REUSE? + norm=self.norm, + axes=(-2, -1), + overwrite_x=overwrite, + ) + + return nearplane + + def _create_fresnel_spectrum_propagator( + self, + N: typing.Tuple[int, int], + pixel_size: float = 1.0, + delta_z: float = 1.0, + wavelength: float = 1.0, + ) -> np.ndarray: + """ + Parameters + ---------- + pixel_size : real width of pixel in meters + delta_z: propagation distance in meters + wavelength: wavelength of light in meters + """ + # FIXME: Check that dimension ordering is consistent + rr2 = self.xp.linspace(-0.5 * N[1], 0.5 * N[1] - 1, num=N[1]) ** 2 + cc2 = self.xp.linspace(-0.5 * N[0], 0.5 * N[0] - 1, num=N[0]) ** 2 + + prb_FOV = self.xp.asarray([pixel_size, pixel_size], dtype=self.xp.float32) + + x = -1j * self.xp.pi * wavelength * delta_z + rr2 = self.xp.exp(x * rr2[..., None] / (prb_FOV[0] ** 2)) + cc2 = self.xp.exp(x * cc2[..., None] / (prb_FOV[1] ** 2)) + + fresnel_spectrum_propagator = self.xp.ndarray.astype( + self.xp.fft.fftshift(self.xp.outer(self.xp.transpose(rr2), cc2)), + dtype=self.xp.csingle, + ) + + return fresnel_spectrum_propagator diff --git a/tests/operators/test_propagation.py b/tests/operators/test_propagation.py index 4722bfb5..69358e47 100644 --- a/tests/operators/test_propagation.py +++ b/tests/operators/test_propagation.py @@ -4,8 +4,7 @@ import unittest import numpy as np -from tike.operators import Propagation -import tike.precision +from tike.operators import Propagation, FresnelSpectProp from .util import random_complex, OperatorTests @@ -37,5 +36,27 @@ def setUp(self, nwaves=13, probe_shape=127): print(self.operator) +class TestFresnelSpectrumPropagation(unittest.TestCase, OperatorTests): + """Test the FresnelSpectProp operator.""" + + def setUp(self, nwaves=13, probe_shape=127): + """Load a dataset for reconstruction.""" + self.operator = FresnelSpectProp( + nwaves=nwaves, + detector_shape=probe_shape, + probe_shape=probe_shape, + ) + self.operator.__enter__() + self.xp = self.operator.xp + np.random.seed(0) + self.m = self.xp.asarray( + random_complex(nwaves, probe_shape, probe_shape)) + self.m_name = 'nearplane' + self.d = self.xp.asarray( + random_complex(nwaves, probe_shape, probe_shape)) + self.d_name = 'farplane' + self.kwargs = {} + print(self.operator) + if __name__ == '__main__': unittest.main() From 480c07f5dcbfabbebad437b590667017113fec59 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Wed, 29 May 2024 18:53:54 -0500 Subject: [PATCH 2/5] REF: Rename parameters to clearer names --- src/tike/operators/cupy/fresnelspectprop.py | 29 +++++++++------------ 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/src/tike/operators/cupy/fresnelspectprop.py b/src/tike/operators/cupy/fresnelspectprop.py index d973ea0f..fa5c8d42 100644 --- a/src/tike/operators/cupy/fresnelspectprop.py +++ b/src/tike/operators/cupy/fresnelspectprop.py @@ -15,23 +15,22 @@ class FresnelSpectProp(CachedFFT, Operator): """Fresnel spectrum propagation (short range) using CuPy. - Take an (..., N, N) array and apply the Fourier transform to the last two - dimensions. + Take an (..., W, H) compelx array representing a wavefront and propagate. Attributes ---------- pixel_size : float The realspace size of a pixel in meters - delta_z : float + distance : float The realspace propagation distance in meters wavelength : float The wavelength of the light in meters Parameters ---------- - nearplane: (..., detector_shape, detector_shape) complex64 + nearplane: (..., W, H) complex64 The wavefronts before propagation. - farplane: (..., detector_shape, detector_shape) complex64 + farplane: (..., W, H) complex64 The wavefronts after propagation. """ @@ -39,13 +38,13 @@ def __init__( self, norm: str = "ortho", pixel_size: float = 1.0, - delta_z: float = 1.0, + distance: float = 1.0, wavelength: float = 1.0, **kwargs, ): self.norm = norm self.pixel_size = pixel_size - self.delta_z = delta_z + self.distance = distance self.wavelength = wavelength def fwd( @@ -58,7 +57,7 @@ def fwd( propagator = self._create_fresnel_spectrum_propagator( (nearplane.shape[-2], nearplane.shape[-1]), self.pixel_size, - self.delta_z, + self.distance, self.wavelength, ) @@ -88,7 +87,7 @@ def adj( propagator = self._create_fresnel_spectrum_propagator( (farplane.shape[-2], farplane.shape[-1]), self.pixel_size, - self.delta_z, + self.distance, self.wavelength, ) @@ -100,10 +99,8 @@ def adj( ) nearplane = self._ifft2( - farplane_fft2 - * self.xp.conj( - propagator, - ), # IS IT OK TO ALWAYS TAKE CONJ? OR SHOULD WE DO THIS ONCE AND REUSE? + # FIXME: IS IT OK TO ALWAYS TAKE CONJ? OR SHOULD WE DO THIS ONCE AND REUSE? + farplane_fft2 * self.xp.conj(propagator), norm=self.norm, axes=(-2, -1), overwrite_x=overwrite, @@ -115,14 +112,14 @@ def _create_fresnel_spectrum_propagator( self, N: typing.Tuple[int, int], pixel_size: float = 1.0, - delta_z: float = 1.0, + distance: float = 1.0, wavelength: float = 1.0, ) -> np.ndarray: """ Parameters ---------- pixel_size : real width of pixel in meters - delta_z: propagation distance in meters + distance: propagation distance in meters wavelength: wavelength of light in meters """ # FIXME: Check that dimension ordering is consistent @@ -131,7 +128,7 @@ def _create_fresnel_spectrum_propagator( prb_FOV = self.xp.asarray([pixel_size, pixel_size], dtype=self.xp.float32) - x = -1j * self.xp.pi * wavelength * delta_z + x = -1j * self.xp.pi * wavelength * distance rr2 = self.xp.exp(x * rr2[..., None] / (prb_FOV[0] ** 2)) cc2 = self.xp.exp(x * cc2[..., None] / (prb_FOV[1] ** 2)) From b61d6d075a62504bc46622f14f966fbeb27fe965 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Wed, 29 May 2024 18:55:31 -0500 Subject: [PATCH 3/5] REF: Reduce number of lines in code --- src/tike/operators/cupy/fresnelspectprop.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/src/tike/operators/cupy/fresnelspectprop.py b/src/tike/operators/cupy/fresnelspectprop.py index fa5c8d42..20c1dafc 100644 --- a/src/tike/operators/cupy/fresnelspectprop.py +++ b/src/tike/operators/cupy/fresnelspectprop.py @@ -115,22 +115,13 @@ def _create_fresnel_spectrum_propagator( distance: float = 1.0, wavelength: float = 1.0, ) -> np.ndarray: - """ - Parameters - ---------- - pixel_size : real width of pixel in meters - distance: propagation distance in meters - wavelength: wavelength of light in meters - """ # FIXME: Check that dimension ordering is consistent rr2 = self.xp.linspace(-0.5 * N[1], 0.5 * N[1] - 1, num=N[1]) ** 2 cc2 = self.xp.linspace(-0.5 * N[0], 0.5 * N[0] - 1, num=N[0]) ** 2 - prb_FOV = self.xp.asarray([pixel_size, pixel_size], dtype=self.xp.float32) - x = -1j * self.xp.pi * wavelength * distance - rr2 = self.xp.exp(x * rr2[..., None] / (prb_FOV[0] ** 2)) - cc2 = self.xp.exp(x * cc2[..., None] / (prb_FOV[1] ** 2)) + rr2 = self.xp.exp(x * rr2[..., None] / (pixel_size**2)) + cc2 = self.xp.exp(x * cc2[..., None] / (pixel_size**2)) fresnel_spectrum_propagator = self.xp.ndarray.astype( self.xp.fft.fftshift(self.xp.outer(self.xp.transpose(rr2), cc2)), From e4fa6c0976ce77ffd421718418059da7da00098a Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Mon, 3 Jun 2024 13:36:46 -0500 Subject: [PATCH 4/5] REF: Set realistic defaults for wave parameters in fresnel prop --- src/tike/operators/cupy/fresnelspectprop.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/tike/operators/cupy/fresnelspectprop.py b/src/tike/operators/cupy/fresnelspectprop.py index 20c1dafc..b01e8aac 100644 --- a/src/tike/operators/cupy/fresnelspectprop.py +++ b/src/tike/operators/cupy/fresnelspectprop.py @@ -37,9 +37,9 @@ class FresnelSpectProp(CachedFFT, Operator): def __init__( self, norm: str = "ortho", - pixel_size: float = 1.0, - distance: float = 1.0, - wavelength: float = 1.0, + pixel_size: float = 1e-5, + distance: float = 1e-6, + wavelength: float = 1e-9, **kwargs, ): self.norm = norm From 26b4c9ed72a63b5c8f76588c997ae721ced8b877 Mon Sep 17 00:00:00 2001 From: a4894z <78556099+a4894z@users.noreply.github.com> Date: Wed, 5 Jun 2024 12:17:41 -0500 Subject: [PATCH 5/5] Update fresnelspectprop.py changed pixel_size default --- src/tike/operators/cupy/fresnelspectprop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tike/operators/cupy/fresnelspectprop.py b/src/tike/operators/cupy/fresnelspectprop.py index b01e8aac..abeb9f85 100644 --- a/src/tike/operators/cupy/fresnelspectprop.py +++ b/src/tike/operators/cupy/fresnelspectprop.py @@ -37,7 +37,7 @@ class FresnelSpectProp(CachedFFT, Operator): def __init__( self, norm: str = "ortho", - pixel_size: float = 1e-5, + pixel_size: float = 1e-7, distance: float = 1e-6, wavelength: float = 1e-9, **kwargs,