Skip to content

Commit

Permalink
Misc fixes (#92)
Browse files Browse the repository at this point in the history 
* show intensity propagation
* support backwards propagation; remove legacy fresnel propagator
* fix nanomax scan reader
  • Loading branch information
@stevehenke
stevehenke authored Jul 16, 2024
1 parent 40d0e28 commit 7ad0992
Show file tree
Hide file tree
Showing 11 changed files with 102 additions and 115 deletions.
9 changes: 4 additions & 5 deletions ptychodus/api/probe.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import numpy

from .geometry import ImageExtent, PixelGeometry
from .propagator import WavefieldArrayType
from .propagator import WavefieldArrayType, intensity
from .typing import RealArrayType


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down
113 changes: 42 additions & 71 deletions ptychodus/api/propagator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand All @@ -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)

Expand All @@ -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))))
Expand All @@ -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)))
11 changes: 7 additions & 4 deletions ptychodus/controller/probe/propagator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion ptychodus/model/analysis/core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
63 changes: 36 additions & 27 deletions ptychodus/model/analysis/propagator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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()
Expand All @@ -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:
Expand All @@ -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()]
Expand All @@ -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()
Expand All @@ -137,4 +144,6 @@ def savePropagatedProbe(self, filePath: Path) -> None:
pixelGeometry.widthInMeters,
'wavefield',
self._propagatedWavefield,
'intensity',
self._propagatedIntensity,
)
1 change: 1 addition & 0 deletions ptychodus/model/analysis/settings.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
4 changes: 2 additions & 2 deletions ptychodus/model/product/probe/averagePattern.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion ptychodus/model/product/probe/builderFactory.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
Loading

0 comments on commit 7ad0992

Please sign in to comment.