mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
106 lines
4.3 KiB
Python
106 lines
4.3 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
|
|
from CDTools.tools import cmath
|
|
import torch as t
|
|
import numpy as np
|
|
|
|
#
|
|
# This file will host tools to turn a propagated wavefield into a measured
|
|
# intensity pattern on a detector
|
|
#
|
|
|
|
__all__ = ['intensity', 'incoherent sum', 'quadratic_background']
|
|
|
|
|
|
def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None):
|
|
"""Returns the intensity of a wavefield
|
|
|
|
The intensity is defined as the magnitude squared of the
|
|
wavefront. If a detector slice is given, the returned array
|
|
will only include that slice from the simulated wavefront.
|
|
|
|
Args:
|
|
wavefield (torch.Tensor) : A JxMxNx2 stack of complex wavefields
|
|
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
|
|
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
|
|
|
|
Returns:
|
|
torch.Tensor : A real MxN array storing the wavefield's intensities
|
|
"""
|
|
if detector_slice is None:
|
|
output = cmath.cabssq(wavefield) + epsilon
|
|
else:
|
|
if wavefield.dim() == 3:
|
|
output = cmath.cabssq(wavefield[detector_slice]) + epsilon
|
|
else:
|
|
output = cmath.cabssq(wavefield[(np.s_[:],) + detector_slice]) + epsilon
|
|
if saturation is None:
|
|
return output
|
|
else:
|
|
return t.clamp(output,0,saturation)
|
|
|
|
|
|
def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None):
|
|
"""Returns the incoherent sum of the intensities of the wavefields
|
|
|
|
The intensity is defined as the sum of the magnitudes squared of
|
|
the wavefields. If a detector slice is given, the returned array
|
|
will only include that slice from the simulated wavefronts.
|
|
|
|
The first index is the set of incoherently adding patterns, and
|
|
the second index is the index of the diffraction pattern to measure.
|
|
The next two indices index the wavefield. The final index is the complex
|
|
index.
|
|
|
|
Args:
|
|
wavefields (torch.Tensor) : An LxJxMxNx2 stack of complex wavefields
|
|
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
|
|
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
|
|
|
|
Returns:
|
|
torch.Tensor : A real JXMxN array storing the incoherently summed intensities
|
|
"""
|
|
# This syntax just adds an axis to the slice to preserve the J direction
|
|
if detector_slice is None:
|
|
output = t.sum(cmath.cabssq(wavefields),dim=0) + epsilon
|
|
else:
|
|
if wavefields.dim() == 4:
|
|
output = t.sum(cmath.cabssq(wavefields[(np.s_[:],)+detector_slice]),dim=0) + epsilon
|
|
else:
|
|
output = t.sum(cmath.cabssq(wavefields[(np.s_[:],np.s_[:])+detector_slice]),dim=0) + epsilon
|
|
|
|
if saturation is None:
|
|
return output
|
|
else:
|
|
return t.clamp(output,0,saturation)
|
|
|
|
|
|
def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None):
|
|
"""Returns the intensity of a wavefield plus a background
|
|
|
|
The intensity is calculated via the given measurment function
|
|
Then, the square of the given background is added. This kind
|
|
of background model is commonly used to enforce positivity of the
|
|
background model.
|
|
|
|
Args:
|
|
wavefield (torch.Tensor) : A JxMxNx2 stack of complex wavefields
|
|
background (torch.Tensor) : An tensor storing the square root of the detector background
|
|
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
|
|
measurement (function) : Optional, the measurement function to use. The default is measurements.intensity
|
|
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
|
|
|
|
Returns:
|
|
torch.Tensor : A real MxN array storing the wavefield's intensities
|
|
"""
|
|
if detector_slice is None:
|
|
output = measurement(wavefield, epsilon=epsilon) + background**2
|
|
else:
|
|
output = measurement(wavefield, detector_slice, epsilon=epsilon) \
|
|
+ background**2
|
|
|
|
if saturation is None:
|
|
return output
|
|
else:
|
|
return t.clamp(output,0,saturation)
|