From 7ad0992ee7738cc5e68a3b852686df4279d9bc0c Mon Sep 17 00:00:00 2001 From: stevehenke <91344068+stevehenke@users.noreply.github.com> Date: Tue, 16 Jul 2024 15:12:21 -0500 Subject: [PATCH] Misc fixes (#92) * show intensity propagation * support backwards propagation; remove legacy fresnel propagator * fix nanomax scan reader --- ptychodus/api/probe.py | 9 +- ptychodus/api/propagator.py | 113 +++++++----------- ptychodus/controller/probe/propagator.py | 11 +- ptychodus/model/analysis/core.py | 2 +- ptychodus/model/analysis/propagator.py | 63 +++++----- ptychodus/model/analysis/settings.py | 1 + .../model/product/probe/averagePattern.py | 4 +- .../model/product/probe/builderFactory.py | 2 +- ptychodus/model/product/probe/fzp.py | 4 +- ptychodus/model/product/probe/settings.py | 2 +- ptychodus/plugins/nanoMaxScanFile.py | 6 +- 11 files changed, 102 insertions(+), 115 deletions(-) diff --git a/ptychodus/api/probe.py b/ptychodus/api/probe.py index aa8285a7..7b5eb7eb 100644 --- a/ptychodus/api/probe.py +++ b/ptychodus/api/probe.py @@ -7,7 +7,7 @@ import numpy from .geometry import ImageExtent, PixelGeometry -from .propagator import WavefieldArrayType +from .propagator import WavefieldArrayType, intensity from .typing import RealArrayType @@ -74,8 +74,8 @@ class Probe: @staticmethod def _calculateModeRelativePower(array: WavefieldArrayType) -> Sequence[float]: - power = numpy.sum((array * array.conj()).real, axis=(-2, -1)) - powersum = power.sum() + power = numpy.sum(intensity(array), axis=(-2, -1)) + powersum = numpy.sum(power) if powersum > 0.: power /= powersum @@ -179,8 +179,7 @@ def getCoherence(self) -> float: return numpy.sqrt(numpy.sum(numpy.square(self._modeRelativePower))) def getIntensity(self) -> RealArrayType: - intensity = numpy.real(self._array * numpy.conjugate(self._array)) - return intensity.sum(axis=-3) + return numpy.sum(intensity(self._array), axis=-3) class ProbeFileReader(ABC): diff --git a/ptychodus/api/propagator.py b/ptychodus/api/propagator.py index e3046c9c..3446e5ed 100644 --- a/ptychodus/api/propagator.py +++ b/ptychodus/api/propagator.py @@ -10,6 +10,10 @@ WavefieldArrayType: TypeAlias = ComplexArrayType +def intensity(wavefield: WavefieldArrayType) -> RealArrayType: + return numpy.real(numpy.multiply(wavefield, numpy.conjugate(wavefield))) + + @dataclass(frozen=True) class PropagatorParameters: wavelength_m: float @@ -43,18 +47,13 @@ def z(self) -> float: @property def fresnel_number(self) -> float: '''fresnel number''' - return numpy.square(self.dx) / self.z - - @property - def diffraction_plane_pixel_width(self) -> float: - '''diffraction plane pixel width in wavelengths''' - return self.z / (self.width_px * self.dx) + return numpy.square(self.dx) / numpy.absolute(self.z) def get_spatial_coordinates(self) -> tuple[RealArrayType, RealArrayType]: - jj, ii = numpy.mgrid[:self.height_px, :self.width_px] - xx = ii - (self.width_px - 1) / 2 - yy = jj - (self.height_px - 1) / 2 - return yy, xx + JJ, II = numpy.mgrid[:self.height_px, :self.width_px] + XX = II - self.width_px // 2 + YY = JJ - self.height_px // 2 + return YY, XX def get_frequency_coordinates(self) -> tuple[RealArrayType, RealArrayType]: fx = fftshift(fftfreq(self.width_px)) @@ -79,7 +78,7 @@ def __init__(self, parameters: PropagatorParameters) -> None: FY, FX = parameters.get_frequency_coordinates() F2 = numpy.square(FX) + numpy.square(ar * FY) ratio = F2 / numpy.square(parameters.dx) - tf = numpy.exp(i2piz * numpy.sqrt(1 - ratio)), + tf = numpy.exp(i2piz * numpy.sqrt(1 - ratio)) self._transfer_function = numpy.where(ratio < 1, tf, 0) @@ -90,14 +89,14 @@ def propagate(self, wavefield: WavefieldArrayType) -> WavefieldArrayType: class FresnelTransferFunctionPropagator(Propagator): def __init__(self, parameters: PropagatorParameters) -> None: - Fr = parameters.fresnel_number ar = parameters.pixel_aspect_ratio - i2pi = 2j * numpy.pi + i2piz = 2j * numpy.pi * parameters.z FY, FX = parameters.get_frequency_coordinates() F2 = numpy.square(FX) + numpy.square(ar * FY) + ratio = F2 / numpy.square(parameters.dx) - self._transfer_function = numpy.exp(i2pi * (parameters.z - 0.5 * F2 / Fr)) + self._transfer_function = numpy.exp(i2piz * (1 - ratio / 2)) def propagate(self, wavefield: WavefieldArrayType) -> WavefieldArrayType: return fftshift(ifft2(self._transfer_function * fft2(ifftshift(wavefield)))) @@ -106,79 +105,51 @@ def propagate(self, wavefield: WavefieldArrayType) -> WavefieldArrayType: class FresnelTransformPropagator(Propagator): def __init__(self, parameters: PropagatorParameters) -> None: - Fr = parameters.fresnel_number - ar = parameters.pixel_aspect_ratio - - i2piz = 2j * numpy.pi * parameters.z ipi = 1j * numpy.pi - iar = 1j * ar + Fr = parameters.fresnel_number + ar = parameters.pixel_aspect_ratio + N = parameters.width_px + M = parameters.height_px YY, XX = parameters.get_spatial_coordinates() - FY, FX = parameters.get_frequency_coordinates() - F2 = numpy.square(FX) + numpy.square(ar * FY) - - self._A = numpy.exp(F2 * ipi / Fr) * numpy.exp(i2piz) * Fr / iar - self._B = numpy.exp(ipi * Fr * (numpy.square(XX) + numpy.square(YY / ar))) - - def propagate(self, wavefield: WavefieldArrayType) -> WavefieldArrayType: - return self._A * fftshift(fft2(ifftshift(wavefield * self._B))) - -class FresnelTransformLegacyPropagator(Propagator): + C0 = Fr / (1j * ar) + C1 = numpy.exp(2j * numpy.pi * parameters.z) + C2 = numpy.exp((numpy.square(XX / N) + numpy.square(ar * YY / M)) * ipi / Fr) + is_forward = (parameters.propagation_distance_m >= 0.) - def __init__(self, parameters: PropagatorParameters) -> None: - self._parameters = parameters + self._is_forward = is_forward + self._A = C2 * C1 * C0 if is_forward else C2 * C1 / C0 + self._B = numpy.exp(ipi * Fr * (numpy.square(XX) + numpy.square(YY / ar))) def propagate(self, wavefield: WavefieldArrayType) -> WavefieldArrayType: - dxy = self._parameters.pixel_width_m - z = self._parameters.propagation_distance_m - wavelength = self._parameters.wavelength_m - - (M, N) = wavefield.shape - k = 2 * numpy.pi / wavelength - - # the coordinate grid - M_grid = numpy.arange(-1 * numpy.floor(M / 2), numpy.ceil(M / 2)) - N_grid = numpy.arange(-1 * numpy.floor(N / 2), numpy.ceil(N / 2)) - lx = M_grid * dxy - ly = N_grid * dxy - - XX, YY = numpy.meshgrid(lx, ly) - - # the coordinate grid on the output plane - fc = 1 / dxy - fu = wavelength * z * fc - lu = M_grid * fu / M - lv = N_grid * fu / N - Fx, Fy = numpy.meshgrid(lu, lv) - - if z > 0: - pf = numpy.exp(1j * k * z) * numpy.exp(1j * k * (Fx**2 + Fy**2) / 2 / z) - kern = wavefield * numpy.exp(1j * k * (XX**2 + YY**2) / 2 / z) - cgh = fft2(fftshift(kern)) - OUT = fftshift(cgh * fftshift(pf)) + if self._is_forward: + return self._A * fftshift(fft2(ifftshift(wavefield * self._B))) else: - pf = numpy.exp(1j * k * z) * numpy.exp(1j * k * (XX**2 + YY**2) / 2 / z) - cgh = ifft2(fftshift(wavefield * numpy.exp(1j * k * (Fx**2 + Fy**2) / 2 / z))) - OUT = fftshift(cgh) * pf - - return OUT + return self._B * fftshift(ifft2(ifftshift(wavefield * self._A))) class FraunhoferPropagator(Propagator): def __init__(self, parameters: PropagatorParameters) -> None: + ipi = 1j * numpy.pi + Fr = parameters.fresnel_number ar = parameters.pixel_aspect_ratio + N = parameters.width_px + M = parameters.height_px + YY, XX = parameters.get_spatial_coordinates() - i2piz = 2j * numpy.pi * parameters.z - ipi = 1j * numpy.pi - iar = 1j * ar - - FY, FX = parameters.get_frequency_coordinates() - F2 = numpy.square(FX) + numpy.square(ar * FY) + C0 = Fr / (1j * ar) + C1 = numpy.exp(2j * numpy.pi * parameters.z) + C2 = numpy.exp((numpy.square(XX / N) + numpy.square(ar * YY / M)) * ipi / Fr) + is_forward = (parameters.propagation_distance_m >= 0.) - self._A = numpy.exp(F2 * ipi / Fr) * numpy.exp(i2piz) * Fr / iar + self._is_forward = is_forward + self._A = C2 * C1 * C0 if is_forward else C2 * C1 / C0 def propagate(self, wavefield: WavefieldArrayType) -> WavefieldArrayType: - return self._A * fftshift(fft2(ifftshift(wavefield))) + if self._is_forward: + return self._A * fftshift(fft2(ifftshift(wavefield))) + else: + return fftshift(ifft2(ifftshift(wavefield * self._A))) diff --git a/ptychodus/controller/probe/propagator.py b/ptychodus/controller/probe/propagator.py index 32f263f3..6f64b6f1 100644 --- a/ptychodus/controller/probe/propagator.py +++ b/ptychodus/controller/probe/propagator.py @@ -55,16 +55,19 @@ def _updateCurrentCoordinate(self, step: int) -> None: logger.error('Bad slider range!') try: - self._xyVisualizationWidgetController.setArray(self._propagator.getXYProjection(step), - self._propagator.getPixelGeometry()) - except IndexError: + xyProjection = self._propagator.getXYProjection(step) + except (IndexError, ValueError): self._xyVisualizationWidgetController.clearArray() except Exception as err: logger.exception(err) ExceptionDialog.showException('Update Current Coordinate', err) + else: + self._xyVisualizationWidgetController.setArray(xyProjection, + self._propagator.getPixelGeometry()) # TODO auto-units - self._dialog.coordinateLabel.setText(f'{lerpValue:.4g} m') + lerpValue *= Decimal('1e6') + self._dialog.coordinateLabel.setText(f'{lerpValue:.1f} \u00B5m') def _propagate(self) -> None: view = self._dialog.parametersView diff --git a/ptychodus/model/analysis/core.py b/ptychodus/model/analysis/core.py index a1ea2904..32f8d8ee 100644 --- a/ptychodus/model/analysis/core.py +++ b/ptychodus/model/analysis/core.py @@ -29,7 +29,7 @@ def __init__(self, settingsRegistry: SettingsRegistry, self._probePropagationSettings = ProbePropagationSettings(settingsRegistry) self.probePropagator = ProbePropagator(self._probePropagationSettings, productRepository) - self.probePropagatorVisualizationEngine = VisualizationEngine(isComplex=True) + self.probePropagatorVisualizationEngine = VisualizationEngine(isComplex=False) self.exposureAnalyzer = ExposureAnalyzer(productRepository) self.exposureVisualizationEngine = VisualizationEngine(isComplex=False) self.fourierRingCorrelator = FourierRingCorrelator(objectRepository) diff --git a/ptychodus/model/analysis/propagator.py b/ptychodus/model/analysis/propagator.py index 0f299c10..c70a7165 100644 --- a/ptychodus/model/analysis/propagator.py +++ b/ptychodus/model/analysis/propagator.py @@ -8,8 +8,10 @@ from ptychodus.api.geometry import PixelGeometry from ptychodus.api.observer import Observable -from ptychodus.api.probe import Probe, WavefieldArrayType -from ptychodus.api.propagator import AngularSpectrumPropagator, PropagatorParameters +from ptychodus.api.probe import Probe +from ptychodus.api.propagator import (AngularSpectrumPropagator, PropagatorParameters, + WavefieldArrayType, intensity) +from ptychodus.api.typing import RealArrayType from ..product import ProductRepository from .settings import ProbePropagationSettings @@ -26,11 +28,13 @@ def __init__(self, settings: ProbePropagationSettings, repository: ProductReposi self._productIndex = -1 self._propagatedWavefield: WavefieldArrayType | None = None + self._propagatedIntensity: RealArrayType | None = None def setProduct(self, productIndex: int) -> None: if self._productIndex != productIndex: self._productIndex = productIndex self._propagatedWavefield = None + self._propagatedIntensity = None self.notifyObservers() def getProductName(self) -> str: @@ -43,9 +47,10 @@ def propagate(self, *, beginCoordinateInMeters: Decimal, endCoordinateInMeters: probe = item.getProbe().getProbe() wavelengthInMeters = item.getGeometry().probeWavelengthInMeters propagatedWavefield = numpy.zeros( - (numberOfSteps, probe.array.shape[-2], probe.array.shape[-1]), + (numberOfSteps, *probe.array.shape), dtype=probe.array.dtype, ) + propagatedIntensity = numpy.zeros((numberOfSteps, *probe.array.shape[-2:])) distanceInMeters = numpy.linspace(float(beginCoordinateInMeters), float(endCoordinateInMeters), numberOfSteps) pixelGeometry = probe.getPixelGeometry() @@ -60,13 +65,16 @@ def propagate(self, *, beginCoordinateInMeters: Decimal, endCoordinateInMeters: propagation_distance_m=zInMeters, ) propagator = AngularSpectrumPropagator(propagatorParameters) - # FIXME Each mode is propagated independently and the total intensity is the sum of the absolute square of their individual intensities. The focus was located by identifying the maximum intensity position which agreed with the focal spot position found by the edge scans. - wf = propagator.propagate(probe.array[0]) - propagatedWavefield[idx, :, :] = wf + + for mode in range(probe.array.shape[-3]): + wf = propagator.propagate(probe.array[mode]) + propagatedWavefield[idx, mode, :, :] = wf + propagatedIntensity[idx, :, :] += intensity(wf) self._settings.beginCoordinateInMeters.value = beginCoordinateInMeters self._settings.endCoordinateInMeters.value = endCoordinateInMeters self._propagatedWavefield = propagatedWavefield + self._propagatedIntensity = propagatedIntensity self.notifyObservers() def getBeginCoordinateInMeters(self) -> Decimal: @@ -84,35 +92,34 @@ def getPixelGeometry(self) -> PixelGeometry: return probe.getPixelGeometry() def getNumberOfSteps(self) -> int: - if self._propagatedWavefield is None: - return 1 + if self._propagatedIntensity is None: + return self._settings.numberOfSteps.value - return self._propagatedWavefield.shape[-3] + return self._propagatedIntensity.shape[0] - def getXYProjection(self, step: int) -> WavefieldArrayType: - if self._propagatedWavefield is None: - probe = self._getProbe() - return probe.array[0] + def getXYProjection(self, step: int) -> RealArrayType: + if self._propagatedIntensity is None: + raise ValueError('No propagated wavefield!') - return self._propagatedWavefield[step] + return self._propagatedIntensity[step] - def getZXProjection(self) -> WavefieldArrayType: - if self._propagatedWavefield is None: + def getZXProjection(self) -> RealArrayType: + if self._propagatedIntensity is None: raise ValueError('No propagated wavefield!') - sz = self._propagatedWavefield.shape[-2] - arrayL = self._propagatedWavefield[:, (sz - 1) // 2, :] - arrayR = self._propagatedWavefield[:, sz // 2, :] - return numpy.transpose(arrayL + arrayR) / 2 # type: ignore + sz = self._propagatedIntensity.shape[-2] + cutPlaneL = self._propagatedIntensity[:, (sz - 1) // 2, :] + cutPlaneR = self._propagatedIntensity[:, sz // 2, :] + return numpy.transpose(numpy.add(cutPlaneL, cutPlaneR) / 2) - def getZYProjection(self) -> WavefieldArrayType: - if self._propagatedWavefield is None: + def getZYProjection(self) -> RealArrayType: + if self._propagatedIntensity is None: raise ValueError('No propagated wavefield!') - sz = self._propagatedWavefield.shape[-1] - arrayL = self._propagatedWavefield[:, :, (sz - 1) // 2] - arrayR = self._propagatedWavefield[:, :, sz // 2] - return numpy.transpose(arrayL + arrayR) / 2 # type: ignore + sz = self._propagatedIntensity.shape[-1] + cutPlaneL = self._propagatedIntensity[:, :, (sz - 1) // 2] + cutPlaneR = self._propagatedIntensity[:, :, sz // 2] + return numpy.transpose(numpy.add(cutPlaneL, cutPlaneR) / 2) def getSaveFileFilterList(self) -> Sequence[str]: return [self.getSaveFileFilter()] @@ -121,7 +128,7 @@ def getSaveFileFilter(self) -> str: return 'NumPy Zipped Archive (*.npz)' def savePropagatedProbe(self, filePath: Path) -> None: - if self._propagatedWavefield is None: + if self._propagatedWavefield is None or self._propagatedIntensity is None: raise ValueError('No propagated wavefield!') pixelGeometry = self.getPixelGeometry() @@ -137,4 +144,6 @@ def savePropagatedProbe(self, filePath: Path) -> None: pixelGeometry.widthInMeters, 'wavefield', self._propagatedWavefield, + 'intensity', + self._propagatedIntensity, ) diff --git a/ptychodus/model/analysis/settings.py b/ptychodus/model/analysis/settings.py index 0e34ab51..f7638a0f 100644 --- a/ptychodus/model/analysis/settings.py +++ b/ptychodus/model/analysis/settings.py @@ -15,6 +15,7 @@ def __init__(self, registry: SettingsRegistry) -> None: 'BeginCoordinateInMeters', '-1e-3') self.endCoordinateInMeters = self._settingsGroup.createRealEntry( 'EndCoordinateInMeters', '+1e-3') + self.numberOfSteps = self._settingsGroup.createIntegerEntry('NumberOfSteps', 100) def update(self, observable: Observable) -> None: if observable is self._settingsGroup: diff --git a/ptychodus/model/product/probe/averagePattern.py b/ptychodus/model/product/probe/averagePattern.py index 104bd529..b2b172a7 100644 --- a/ptychodus/model/product/probe/averagePattern.py +++ b/ptychodus/model/product/probe/averagePattern.py @@ -3,7 +3,7 @@ import numpy from ptychodus.api.probe import Probe, ProbeGeometryProvider -from ptychodus.api.propagator import FresnelTransformLegacyPropagator, PropagatorParameters +from ptychodus.api.propagator import FresnelTransformPropagator, PropagatorParameters from ...patterns import ActiveDiffractionDataset, Detector from .builder import ProbeBuilder @@ -32,7 +32,7 @@ def build(self, geometryProvider: ProbeGeometryProvider) -> Probe: pixel_height_m=pixelGeometry.heightInMeters, propagation_distance_m=-geometryProvider.detectorDistanceInMeters, ) - propagator = FresnelTransformLegacyPropagator(propagatorParameters) + propagator = FresnelTransformPropagator(propagatorParameters) array = propagator.propagate(numpy.sqrt(detectorIntensity).astype(complex)) return Probe( diff --git a/ptychodus/model/product/probe/builderFactory.py b/ptychodus/model/product/probe/builderFactory.py index 432aac60..6f2766ee 100644 --- a/ptychodus/model/product/probe/builderFactory.py +++ b/ptychodus/model/product/probe/builderFactory.py @@ -33,8 +33,8 @@ def __init__(self, settings: ProbeSettings, detector: Detector, self._fileReaderChooser = fileReaderChooser self._fileWriterChooser = fileWriterChooser self._builders: Mapping[str, Callable[[], ProbeBuilder]] = { - 'average_pattern': self._createAveragePatternBuilder, 'disk': lambda: DiskProbeBuilder(settings), + 'average_pattern': self._createAveragePatternBuilder, 'fresnel_zone_plate': self._createFresnelZonePlateBuilder, 'rectangular': lambda: RectangularProbeBuilder(settings), 'super_gaussian': lambda: SuperGaussianProbeBuilder(settings), diff --git a/ptychodus/model/product/probe/fzp.py b/ptychodus/model/product/probe/fzp.py index 70642211..8442780f 100644 --- a/ptychodus/model/product/probe/fzp.py +++ b/ptychodus/model/product/probe/fzp.py @@ -7,7 +7,7 @@ from ptychodus.api.geometry import PixelGeometry from ptychodus.api.plugins import PluginChooser from ptychodus.api.probe import FresnelZonePlate, Probe, ProbeGeometryProvider -from ptychodus.api.propagator import FresnelTransformLegacyPropagator, PropagatorParameters +from ptychodus.api.propagator import FresnelTransformPropagator, PropagatorParameters from .builder import ProbeBuilder from .settings import ProbeSettings @@ -100,7 +100,7 @@ def build(self, geometryProvider: ProbeGeometryProvider) -> Probe: pixel_height_m=fzpPixelGeometry.heightInMeters, propagation_distance_m=distanceInMeters, ) - propagator = FresnelTransformLegacyPropagator(propagatorParameters) + propagator = FresnelTransformPropagator(propagatorParameters) array = propagator.propagate(fzpTransmissionFunction) return Probe( diff --git a/ptychodus/model/product/probe/settings.py b/ptychodus/model/product/probe/settings.py index 487c6f08..ee458e20 100644 --- a/ptychodus/model/product/probe/settings.py +++ b/ptychodus/model/product/probe/settings.py @@ -42,7 +42,7 @@ def __init__(self, registry: SettingsRegistry) -> None: self.centralBeamstopDiameterInMeters = self._settingsGroup.createRealEntry( 'CentralBeamstopDiameterInMeters', '60e-6') self.defocusDistanceInMeters = self._settingsGroup.createRealEntry( - 'DefocusDistanceInMeters', '800e-6') + 'DefocusDistanceInMeters', '0') def update(self, observable: Observable) -> None: if observable is self._settingsGroup: diff --git a/ptychodus/plugins/nanoMaxScanFile.py b/ptychodus/plugins/nanoMaxScanFile.py index 052ca949..a5d7e1fc 100644 --- a/ptychodus/plugins/nanoMaxScanFile.py +++ b/ptychodus/plugins/nanoMaxScanFile.py @@ -29,7 +29,11 @@ def read(self, filePath: Path) -> Scan: raise ScanPointParseError('Coordinate array shape mismatch!') for idx, (x, y) in enumerate(zip(positionX, positionY)): - point = ScanPoint(idx, x, y) + point = ScanPoint( + idx, + x * self.MICRONS_TO_METERS, + y * self.MICRONS_TO_METERS, + ) pointList.append(point) return Scan(pointList)