cSAXS/cdtools
0
0 Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-11-12 17:23:54 +01:00
Code Issues Packages Projects Releases Wiki Activity

Get tests to past and removeopt_for_fft

Browse Source
This commit is contained in:
levita_a committed 2023-12-01 15:58:18 +01:00
1 parent 8ef3411604
commit a9e0157707
26 files changed
+96 -308
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/source/tutorial.rst
+1 -2
View File
@@ -358,8 +358,7 @@ To initialize the object from a dataset, we need to start by extracting the rele
probe_basis, probe_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
opt_for_fft=False)
distance)
# Next generate the object geometry
pix_translations = tools.interactions.translations_to_pixel(
examples/simple_ptycho_model.py
+1 -2
View File
@@ -39,8 +39,7 @@ class SimplePtycho(CDIModel):
probe_basis, probe_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
opt_for_fft=False)
distance)
# Next generate the object geometry
pix_translations = tools.interactions.translations_to_pixel(
src/cdtools/datasets/__init__.py
+1 -1
View File
@@ -29,7 +29,7 @@ attempting to do so
# I don't believe that __all__ really needed, but it's nice to define it
# to be explicit that import * is safe
__all__ = ['CDataset','Ptycho2DDataset','PolarizationSwepttycho2DDataset']
__all__ = ['CDataset','Ptycho2DDataset','PolarizationSweptPtycho2DDataset']
from cdtools.datasets.base import CDataset
from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py
+11 -8
View File
@@ -66,7 +66,7 @@ class PolarizationSweptPtycho2DDataset(Ptycho2DDataset):
*args, **kwargs)
self.polarization_indices = t.tensor(polarization_indices, dtype=t.int32)
self.polarization_states = t.tensor(polarization_states, dtype=t.float32)
self.polarization_states = t.tensor(polarization_states, dtype=t.complex64)
def _load(self, index):
@@ -154,15 +154,18 @@ class PolarizationSweptPtycho2DDataset(Ptycho2DDataset):
dataset.__class__ = cls
# Now, we save out the polarizer and analyzer states
polarization_indices = cdtdata.get_shot_to_shot_info(cxi_file, 'polarization_indices')
polarization_states = cdtdata.get_entry_info(cxi_file, 'polarization_states')
polarization_indices = cdtdata.get_shot_to_shot_info(
cxi_file, 'polarization_indices')
polarization_states = dataset.entry_info['polarization_states']
dataset.polarization_indices = t.tensor(polarization_indices, dtype=t.int32)
dataset.polarization_states = t.tensor(polarization_states, dtype=t.complex64)
dataset.polarization_indices = t.tensor(polarization_indices,
dtype=t.int32)
dataset.polarization_states = t.tensor(polarization_states,
dtype=t.complex64)
return dataset
def to_cxi(self, cxi_file, polarized=False):
def to_cxi(self, cxi_file):
"""Saves out a PolarizationSweptPtycho2DDataset as a .cxi file
This function saves all the compatible information in a
@@ -187,5 +190,5 @@ class PolarizationSweptPtycho2DDataset(Ptycho2DDataset):
# Now, we save out the polarization states
cdtdata.add_shot_to_shot_info(cxi_file, self.polarization_indices,
'polarization_indices')
cdtdata.add_entry_info(cxi_file, self.polarization_states,
'polarization_states')
cdtdata.add_entry_info(cxi_file, {'polarization_states':
self.polarization_states})
src/cdtools/models/base.py
+1 -1
View File
@@ -284,7 +284,7 @@ class CDIModel(t.nn.Module):
# If we don't want to run in a different thread, this is easy
if not thread:
for it in range(iterations):
yield run_iteration()
yield run_epoch()
# But if we do want to thread, it's annoying:
else:
src/cdtools/models/bragg_2d_ptycho.py
+1 -2
View File
@@ -209,7 +209,7 @@ class Bragg2DPtycho(CDIModel):
@classmethod
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True, correct_tilt=True, lens=False, opt_for_fft=False):
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True, correct_tilt=True, lens=False):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -236,7 +236,6 @@ class Bragg2DPtycho(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
# now we grab the sample surface normal
if hasattr(dataset, 'sample_info') and \
src/cdtools/models/fancy_ptycho.py
-2
View File
@@ -177,7 +177,6 @@ class FancyPtycho(CDIModel):
scattering_mode=None,
oversampling=1,
auto_center=False,
opt_for_fft=False,
fourier_probe=False,
loss='amplitude mse',
units='um',
@@ -216,7 +215,6 @@ class FancyPtycho(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
src/cdtools/models/fastccd_ptycho.py
-2
View File
@@ -160,7 +160,6 @@ class FastCCDPtycho(CDIModel):
scattering_mode=None,
oversampling=1,
auto_center=False,
opt_for_fft=False,
fourier_probe=False,
loss='amplitude mse',
units='um',
@@ -197,7 +196,6 @@ class FastCCDPtycho(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
src/cdtools/models/multimode_rpi.py
+2 -7
View File
@@ -4,16 +4,12 @@ from cdtools import tools
from cdtools.tools import plotting as p
from cdtools.tools.interactions import RPI_interaction
from cdtools.tools import initializers
from scipy.ndimage.morphology import binary_dilation
from scipy.ndimage import binary_dilation
import numpy as np
from copy import copy
__all__ = ['MultimodeRPI']
__all__ = ['RPI']
class MultimodeRPI(CDIModel):
@property
@@ -99,7 +95,7 @@ class MultimodeRPI(CDIModel):
@classmethod
def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0):
def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', weight_matrix=False, probe_threshold=0):
raise NotImplementedError()
wavelength = dataset.wavelength
@@ -128,7 +124,6 @@ class MultimodeRPI(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if not isinstance(probe,t.Tensor):
src/cdtools/models/multislice_2d_ptycho.py
-1
View File
@@ -181,7 +181,6 @@ class Multislice2DPtycho(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=False,
oversampling=oversampling)
src/cdtools/models/polarization_swept_ptycho.py
+34 -34
View File
@@ -117,7 +117,7 @@ class PolarizationSweptPtycho(CDIModel):
t.tensor(translation_scale, dtype=dtype))
if probe_support is None:
probe_support = t.ones_like(self.probe[0], dtype=t.bool)
probe_support = t.ones(self.probe.shape[-2:], dtype=t.bool)
self.register_buffer('probe_support',
t.tensor(probe_support, dtype=t.bool))
@@ -139,7 +139,7 @@ class PolarizationSweptPtycho(CDIModel):
self.simulate_finite_pixels = simulate_finite_pixels
self.polarization_states = t.nn.Parameter(
t.tensor(polarization_states, dtype=dtype))
t.tensor(polarization_states, dtype=t.complex64))
# by default, don't optimize this, but I think it might be
# interesting to try it because the polarization states
# are not very pure
@@ -172,7 +172,6 @@ class PolarizationSweptPtycho(CDIModel):
scattering_mode=None,
oversampling=1,
auto_center=False,
opt_for_fft=False,
fourier_probe=False,
loss='amplitude mse',
units='um',
@@ -212,7 +211,6 @@ class PolarizationSweptPtycho(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
@@ -298,7 +296,7 @@ class PolarizationSweptPtycho(CDIModel):
mask = None
if probe_support_radius is not None:
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
probe_support = t.zeros(probe.shape[-2:], dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
@@ -344,15 +342,14 @@ class PolarizationSweptPtycho(CDIModel):
if self.translation_offsets is not None:
pix_trans += (self.translation_scale *
self.translation_offsets[index])
# This restricts the basis probes to stay within the probe support
basis_prs = self.probe * self.probe_support[..., :, :]
print(basis_prs.shape)
# Now we expand the polarization states explicitly.
# pol_basis_prs has dimensions n_states x 2 x n_modes x m x l
pol_basis_prs = (
self.polarization_states[...,None,None,None]
* basis_prs[...,None,:,:,:])
print(pol_basis_prs.shape)
# For a Fourier-space probe, we take an IFT
if self.fourier_probe:
basis_prs = tools.propagators.inverse_far_field(basis_prs)
@@ -360,13 +357,13 @@ class PolarizationSweptPtycho(CDIModel):
# So, first we get the appropriate probe for each shot.
# prs is now n_frames x 2 x n_modes x m x l
prs = pol_basis_prs[polarization_indices]
print(prs.shape)
# Now we construct the probes for each shot from the basis probes
if self.weights is not None:
Ws = self.weights[index]
# And then we multiply by the weights, along the 0th dimension
prs = Ws[..., None, None, None, None] * basis_prs
prs = Ws[..., None, None, None, None] * prs#basis_prs
if self.simulate_probe_translation:
det_pix_trans = tools.interactions.translations_to_pixel(
self.det_basis,
@@ -394,23 +391,28 @@ class PolarizationSweptPtycho(CDIModel):
prs = tools.propagators.inverse_far_field(prs)
# Now we actually do the interaction, using the sinc subpixel
# translation model as per usual
print('hi')
print(prs[...,0,:,:,:].shape)
print(self.obj[:,0].shape)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,0,:,:,:], self.obj[:,0], pix_trans,
shift_probe=True, multiple_modes=True)
exit_waves = exit_waves + (
self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,1,:,:,:], self.obj[:,1], pix_trans,
shift_probe=True, multiple_modes=True)
)
print(exit_waves.shape)
# After the object, no point in treating the polarization modes
# any differently from the object/probe modes
print(exit_waves.flatten(start_dim=1,end_dim=-3).shape)
print('sup')
return exit_waves.flatten(start_dim=1,end_dim=-3)
# So, this stacks up all the output modes in lexicographic order
exit_waves = t.cat(
[self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,0,:,:,:], self.obj[0,0], pix_trans,
shift_probe=True, multiple_modes=True),
self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,0,:,:,:], self.obj[1,0], pix_trans,
shift_probe=True, multiple_modes=True)],
dim=1)
exit_waves = exit_waves + t.cat(
[self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,1,:,:,:], self.obj[0,1], pix_trans,
shift_probe=True, multiple_modes=True),
self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,1,:,:,:], self.obj[1,1], pix_trans,
shift_probe=True, multiple_modes=True)],
dim=1)
return exit_waves
def forward_propagator(self, wavefields):
@@ -519,14 +521,12 @@ class PolarizationSweptPtycho(CDIModel):
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# First we treat the purely incoherent case
# I don't love this pattern of using an if statement with a return
# to catch this case, but because it's so much simpler than the
# unified mode case I think it's appropriate
probe = self.probe.detach().cpu().numpy()
ortho_probes = analysis.orthogonalize_probes(probe)
ortho_probes = np.empty(probe.shape, dtype=probe.dtype)
for idx in range(probe.shape[0]):
ortho_probes[idx] = analysis.orthogonalize_probes(probe[idx])
self.probe.data = t.as_tensor(
ortho_probes,
device=self.probe.device,
@@ -572,13 +572,13 @@ class PolarizationSweptPtycho(CDIModel):
lambda self: len(self.weights.shape) >= 2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
self.obj[np.s_[:,:] + self.obj_view_slice],
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Object Phase',
lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
self.obj[np.s_[:,:] + self.obj_view_slice],
fig=fig,
basis=self.probe_basis,
units=self.units)),
src/cdtools/models/polarized_fancy_ptycho.py
+2 -4
View File
@@ -61,7 +61,7 @@ class PolarizedFancyPtycho(FancyPtycho):
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
@classmethod
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True):
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um', left_polarized=True):
# When using this method, remember to pass through the inputs
model = FancyPtycho.from_dataset(
@@ -78,7 +78,6 @@ class PolarizedFancyPtycho(FancyPtycho):
scattering_mode=scattering_mode,
oversampling=oversampling,
auto_center=auto_center,
opt_for_fft=opt_for_fft,
loss=loss,
units=units)
@@ -122,7 +121,7 @@ class PolarizedFancyPtycho(FancyPtycho):
polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)]
@classmethod
def from_dataset2(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um'):
def from_dataset2(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um'):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -159,7 +158,6 @@ class PolarizedFancyPtycho(FancyPtycho):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
probe_shape = t.stack((2, probe_shape), dim=-3)
src/cdtools/models/rpi.py
+2 -3
View File
@@ -4,7 +4,7 @@ from cdtools import tools
from cdtools.tools import plotting as p
from cdtools.tools.interactions import RPI_interaction
from cdtools.tools import initializers
from scipy.ndimage.morphology import binary_dilation
from scipy.ndimage import binary_dilation
import numpy as np
from copy import copy
import time
@@ -139,7 +139,7 @@ class RPI(CDIModel):
@classmethod
def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0, simulate_finite_pixels=False):
def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', weight_matrix=False, probe_threshold=0, simulate_finite_pixels=False):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -167,7 +167,6 @@ class RPI(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if not isinstance(probe,t.Tensor):
src/cdtools/models/s_matrix_ptycho.py
-1
View File
@@ -140,7 +140,6 @@ class SMatrixPtycho(CDIModel):
wavelength,
distance,
padding=padding,
opt_for_fft=False,
oversampling=oversampling)
if propagation_distances is None:
src/cdtools/models/time_resolved_ptycho_calibration.py
+1 -2
View File
@@ -125,7 +125,7 @@ class TimeResolvedPtychoCalibration(CDIModel):
@classmethod
def from_dataset(cls, dataset, zp_geometry, time_window, n_times, n_frames, randomize_ang=0, padding=0, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', simulate_probe_translation=False):
def from_dataset(cls, dataset, zp_geometry, time_window, n_times, n_frames, randomize_ang=0, padding=0, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um', simulate_probe_translation=False):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -156,7 +156,6 @@ class TimeResolvedPtychoCalibration(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
src/cdtools/models/time_resolved_rpi.py
+3 -7
View File
@@ -4,15 +4,12 @@ from cdtools import tools
from cdtools.tools import plotting as p
from cdtools.tools.interactions import RPI_interaction
from cdtools.tools import initializers
from scipy.ndimage.morphology import binary_dilation
from scipy.ndimage import binary_dilation
import numpy as np
from copy import copy
__all__ = ['MultimodeRPI']
__all__ = ['RPI']
__all__ = ['TimeResolvedRPI']
class TimeResolvedRPI(CDIModel):
@@ -89,7 +86,7 @@ class TimeResolvedRPI(CDIModel):
@classmethod
def from_dataset(cls, dataset, probe, framerate, obj_size=None, background=None, mask=None, padding=0, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, probe_threshold=0):
def from_dataset(cls, dataset, probe, framerate, obj_size=None, background=None, mask=None, padding=0, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', probe_threshold=0):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -117,7 +114,6 @@ class TimeResolvedRPI(CDIModel):
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if not isinstance(probe,t.Tensor):
src/cdtools/tools/analysis/analysis.py
-2
View File
@@ -764,11 +764,9 @@ def calc_fidelity(fields_1, fields_2, dims=2):
fields_1 = t.as_tensor(fields_1)
fields_2 = t.as_tensor(fields_2)
mult = fields_1.unsqueeze(-dims-2) * fields_2.unsqueeze(-dims-1).conj()
sumdims = tuple(d - dims for d in range(dims))
mat = t.sum(mult,dim=sumdims)
# Because I think this is the nuclear norm squared, I would like to swap
# Out the definition for this, but I need to test it before swapping.
# It also probably makes sense to implement sqrt_fidelity separately
src/cdtools/tools/data/data.py
+6
View File
@@ -76,6 +76,12 @@ def get_entry_info(cxi_file):
metadata[attr] = dateutil.parser.parse(str(e1[attr][()].decode()))
except ValueError:
metadata[attr] = str(e1[attr][()].decode())
array_attrs = ['polarization_states']
for attr in array_attrs:
if attr in e1:
metadata[attr] = np.array(e1[attr])
return metadata
src/cdtools/tools/initializers/initializers.py
+1 -13
View File
@@ -21,12 +21,7 @@ __all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
'RPI_spectral_init',
'generate_subdominant_modes']
# I need to remove a dependency on scipy.fftpack, so I did this.
# I should replace this with a real function
def next_fast_len(x):
return x
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0, oversampling=1):
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, padding=0, oversampling=1):
"""Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry
It takes in the parameters for a given detector - the basis defining
@@ -48,8 +43,6 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
The sample-detector distance, in m
center : torch.Tensor
If defined, the location of the zero frequency pixel
opt_for_fft : bool
Default is true, whether to increase detector size to improve fft performance
padding : int
Default is 0, the size of an extra border of nonphysical pixels around the detector
oversampling : int
@@ -85,11 +78,6 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
# In some edge cases this shape can be smaller than the detector shape
full_shape = t.max(full_shape, det_shape)
if opt_for_fft:
full_shape = t.as_tensor([next_fast_len(dim) for dim in full_shape],
dtype=t.int32)
# Then, generate a slice that pops the actual detector from the full
# detector shape
src/cdtools/tools/propagators/propagators.py
+7 -1
View File
@@ -540,6 +540,12 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
# adding to get the full 3D wavevectors. Note that we convert to complex
# before the square root to appropriately map negative numbers to
# complex frequencies
# The correct algorithm for sqrt(1-x**2) - 1 is
# - x**2 / (sqrt(1-x**2) + 1)
# This one doesn't have numerical stability issues near x=0.
# I need to figure out how to apply that knowledge here.
K = K_ip + perpendicular_dir[:,None,None] \
* t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared)))
@@ -558,7 +564,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
K_m_K_0 = K - K_0[:,None,None]
# We actually calculate the phase mask
phase_mask = t.tensordot(offset_vector,K_m_K_0, dims=1)
# And we take the conjugate if needed, which makes sure that the
# imaginary part is always positive (supressing evanescent waves)
if propagation_vector is not None \
tests/test_datasets.py
-180
View File
@@ -276,183 +276,3 @@ def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
assert t.allclose(pattern.to(device='cpu'),
t.tensor(expected['data'][3,:,:]))
#
# And we now test the derived class for polarization-dependent ptycho
#
def test_PolarizedPtycho2DDataset_init():
entry_info = {'start_time': datetime.datetime.now(),
'title' : 'A simple test'}
sample_info = {'name': 'A test sample',
'mass' : 3.4,
'unit_cell' : np.array([1,1,1,87,84.5,90])}
wavelength = 1e-9
detector_geometry = {'distance': 0.7,
'basis': np.array([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(),
'corner': np.array((2550e-6,3825e-6,0.3))}
mask = np.ones((256,256))
patterns = np.random.rand(20,256,256)
translations = np.random.rand(20,3)
analyzer = np.random.rand(20)
polarizer = np.random.rand(20)
dataset = PolarizedPtycho2DDataset(translations, polarizer,
analyzer, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
assert dataset.detector_geometry == detector_geometry
assert t.allclose(dataset.patterns, t.as_tensor(patterns))
assert t.allclose(dataset.translations, t.as_tensor(translations))
assert t.allclose(dataset.analyzer, t.as_tensor(analyzer, dtype=t.float32))
assert t.allclose(dataset.polarizer, t.as_tensor(polarizer, dtype=t.float32))
def test_PolarizedPtycho2DDataset_from_cxi(polarized_ptycho_cxi):
cxi, expected = polarized_ptycho_cxi
dataset = PolarizedPtycho2DDataset.from_cxi(cxi)
# The entry metadata loaded
for key in expected['entry metadata']:
assert dataset.entry_info[key] == expected['entry metadata'][key]
# Don't test for fidelity since this is tested in the data, just test
# that it is loaded
if expected['sample info'] is None:
assert dataset.sample_info is None
else:
assert dataset.sample_info is not None
assert np.isclose(dataset.wavelength,expected['wavelength'])
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
expected['detector']['distance'])
# Check that the other ones are loaded but not for fidelity
assert 'basis' in dataset.detector_geometry
if expected['detector']['corner'] is not None:
assert 'corner' in dataset.detector_geometry
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
if expected['dark'] is not None:
assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
dataset.background))
assert t.allclose(t.tensor(expected['data']),dataset.patterns)
assert t.allclose(t.tensor(expected['translations']),dataset.translations)
assert t.allclose(t.tensor(expected['analyzer_angle']),dataset.analyzer)
assert t.allclose(t.tensor(expected['polarizer_angle']),dataset.polarizer)
def test_PolarizedPtycho2DDataset_to_cxi(polarized_ptycho_cxi, tmp_path):
cxi, expected = polarized_ptycho_cxi
print('loading dataset')
dataset = PolarizedPtycho2DDataset.from_cxi(cxi)
print('dataset mask is type', dataset.mask.dtype)
with cdtdata.create_cxi(tmp_path / 'test_PolarizedPtycho2DDataset_to_cxi.cxi') as f:
dataset.to_cxi(f)
# Now we have to check that all the stuff was written
with h5py.File(tmp_path / 'test_PolarizedPtycho2DDataset_to_cxi.cxi', 'r') as f:
read_dataset = PolarizedPtycho2DDataset.from_cxi(f)
assert dataset.entry_info == read_dataset.entry_info
if dataset.sample_info is None:
assert read_dataset.sample_info is None
else:
assert read_dataset.sample_info is not None
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
read_dataset.detector_geometry['distance'])
# Check that the other ones are loaded but not for fidelity
assert 'basis' in read_dataset.detector_geometry
if dataset.detector_geometry['corner'] is not None:
assert 'corner' in read_dataset.detector_geometry
if dataset.mask is not None:
assert t.all(t.eq(dataset.mask,read_dataset.mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
assert t.allclose(dataset.patterns, read_dataset.patterns)
assert t.allclose(dataset.translations, read_dataset.translations)
assert t.allclose(dataset.analyzer, read_dataset.analyzer)
assert t.allclose(dataset.analyzer, read_dataset.analyzer)
def test_PolarizedPtycho2DDataset_to(polarized_ptycho_cxi):
dataset = PolarizedPtycho2DDataset.from_cxi(polarized_ptycho_cxi[0])
dataset.to(dtype=t.float64)
assert dataset.mask.dtype == t.bool
assert dataset.patterns.dtype == t.float64
assert dataset.translations.dtype == t.float64
assert dataset.analyzer.dtype == t.float64
assert dataset.polarizer.dtype == t.float64
# If cuda is available, check that moving the mask to CUDA works.
if t.cuda.is_available():
dataset.to(device='cuda:0')
assert dataset.mask.device == t.device('cuda:0')
assert dataset.background.device == t.device('cuda:0')
assert dataset.patterns.device == t.device('cuda:0')
assert dataset.translations.device == t.device('cuda:0')
assert dataset.analyzer.device == t.device('cuda:0')
assert dataset.polarizer.device == t.device('cuda:0')
def test_PolarizedPtycho2DDataset_ops(polarized_ptycho_cxi):
cxi, expected = polarized_ptycho_cxi
dataset = PolarizedPtycho2DDataset.from_cxi(cxi)
dataset.get_as('cpu')
assert len(dataset) == expected['data'].shape[0]
(idx, translation, polarizer, analyzer), pattern = dataset[3]
assert idx == 3
assert t.allclose(translation, t.tensor(expected['translations'][3,:]))
assert t.allclose(pattern, t.tensor(expected['data'][3,:,:]))
assert t.allclose(analyzer, t.tensor(expected['analyzer_angle'][3]))
assert t.allclose(polarizer, t.tensor(expected['polarizer_angle'][3]))
def test_PolarizedPtycho2DDataset_get_as(polarized_ptycho_cxi):
cxi, expected = polarized_ptycho_cxi
dataset = PolarizedPtycho2DDataset.from_cxi(cxi)
if t.cuda.is_available():
dataset.get_as('cuda:0')
assert len(dataset) == expected['data'].shape[0]
(idx, translation, polarizer, analyzer), pattern = dataset[3]
assert str(translation.device) == 'cuda:0'
assert str(pattern.device) == 'cuda:0'
assert idx == 3
assert t.allclose(translation.to(device='cpu'),
t.tensor(expected['translations'][3,:]))
assert t.allclose(pattern.to(device='cpu'),
t.tensor(expected['data'][3,:,:]))
assert t.allclose(polarizer.to(device='cpu'),
t.tensor(expected['polarizer_angle'][3]))
assert t.allclose(analyzer.to(device='cpu'),
t.tensor(expected['analyzer_angle'][3]))
tests/tools/test_analysis.py
+6 -4
View File
@@ -367,18 +367,20 @@ def test_calc_fidelity():
dm_1 = t.tensordot(dm_1.transpose(0,1), dm_1.conj(), dims=1).numpy()
dm_2 = t.reshape(fields_2, (3,-1))
dm_2 = t.tensordot(dm_2.transpose(0,1), dm_2.conj(), dims=1).numpy()
inner_mat = la.sqrtm(np.dot(np.dot(la.sqrtm(dm_1),dm_2),la.sqrtm(dm_1)))
sqrt_dm_1 = la.sqrtm(dm_1).astype(dm_1.dtype)
inner_mat = la.sqrtm(np.dot(np.dot(sqrt_dm_1,dm_2), sqrt_dm_1))
inner_mat = inner_mat.astype(dm_1.dtype) #la.sqrtm doubles the precision
fidelity = t.as_tensor(np.abs(np.trace(inner_mat))**2)
assert t.isclose(fidelity, analysis.calc_fidelity(fields_1, fields_2))
# Check that it reduces to the overlap for coherent fields
fields_1 = t.rand(1,30,17, dtype=t.complex128)
fields_2 = t.rand(1,30,17, dtype=t.complex128)
assert t.isclose(t.abs(t.sum(fields_1*fields_2.conj()))**2,
analysis.calc_fidelity(fields_1, fields_2))
# Checking that it works with extra dimensions
fields_1 = t.rand(3,3,30,17, dtype=t.complex128)
fields_2 = t.rand(3,1,30,17, dtype=t.complex128)
tests/tools/test_image_processing.py
+4 -4
View File
@@ -7,13 +7,13 @@ from scipy import ndimage
def test_centroid():
# Test single im
im = t.rand((30,40))
sp_centroid = ndimage.measurements.center_of_mass(im.numpy())
sp_centroid = ndimage.center_of_mass(im.numpy())
centroid = image_processing.centroid(im)
assert t.allclose(centroid, t.Tensor(sp_centroid))
# Test stack o' ims
ims = t.rand((5,30,40))
sp_centroids = [ndimage.measurements.center_of_mass(im.numpy())
sp_centroids = [ndimage.center_of_mass(im.numpy())
for im in ims]
centroids = image_processing.centroid(ims)
assert t.allclose(centroids, t.Tensor(sp_centroids))
@@ -22,14 +22,14 @@ def test_centroid():
def test_centroid_sq():
# Test single im
im = t.rand((30,40))
sp_centroid = ndimage.measurements.center_of_mass(im.numpy()**2)
sp_centroid = ndimage.center_of_mass(im.numpy()**2)
centroid = image_processing.centroid_sq(im)
assert t.allclose(centroid, t.Tensor(sp_centroid))
# Test complex with multiple ims
ims = t.rand((5,30,40)) + 1j * t.rand((5,30,40))
np_ims = ims.numpy()
sp_centroids = [ndimage.measurements.center_of_mass(np.abs(im)**2)
sp_centroids = [ndimage.center_of_mass(np.abs(im)**2)
for im in np_ims]
centroids = image_processing.centroid_sq(ims, comp=True)
assert t.allclose(centroids, t.Tensor(np.array(sp_centroids)))
tests/tools/test_initializers.py
+3 -16
View File
@@ -12,28 +12,16 @@ def test_exit_wave_geometry():
distance = 1.
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance, opt_for_fft=False)
distance)
assert full_shape == shape
assert t.ones(full_shape)[det_slice].shape == shape
assert t.allclose(rs_basis[0,1],t.Tensor([-8.928571428571428e-07]))
assert t.allclose(rs_basis[1,0],t.Tensor([-4.5662100456621004e-07]))
# Then test it's expanding functionality for a non-optimal array
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance, opt_for_fft=True)
exp_shape = t.Size([75,60])
assert full_shape == exp_shape
assert t.ones(full_shape)[det_slice].shape == shape
assert t.allclose(rs_basis[0,1],t.Tensor([-8.333333333333333e-07]))
assert t.allclose(rs_basis[1,0],t.Tensor([-4.444444444444444e-07]))
# Then test it's padding function
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance, opt_for_fft=False,
padding=2)
distance, padding=2)
exp_shape = t.Size([77,60])
assert full_shape == exp_shape
assert t.ones(full_shape)[det_slice].shape == shape
@@ -43,8 +31,7 @@ def test_exit_wave_geometry():
center = t.Tensor([20,42])
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance, center=center,
opt_for_fft=False)
distance, center=center)
exp_shape = t.Size([105,84])
assert full_shape == exp_shape
assert t.ones(full_shape)[det_slice].shape == shape
tests/tools/test_plotting.py
+4 -4
View File
@@ -2,18 +2,18 @@ from cdtools.tools import plotting
from cdtools.tools import initializers
import numpy as np
import torch as t
import scipy.misc
import scipy.datasets
import matplotlib.pyplot as plt
def test_plot_amplitude(show_plot):
# Test with tensor
im = t.as_tensor(scipy.misc.ascent(),dtype=t.complex128)
im = t.as_tensor(scipy.datasets.ascent(),dtype=t.complex128)
plotting.plot_amplitude(im, basis = np.array([[1,1], [1,1], [0,0]]), title = 'Test Amplitude')
if show_plot:
plt.show()
# Test with numpy array
im = scipy.misc.ascent().astype(np.complex128)
im = scipy.datasets.ascent().astype(np.complex128)
plotting.plot_amplitude(im, title = 'Test Amplitude')
if show_plot:
plt.show()
@@ -35,7 +35,7 @@ def test_plot_phase(show_plot):
def test_plot_colorize(show_plot):
# Test with tensor
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
im = gaussian * t.as_tensor(scipy.misc.ascent(), dtype=t.complex64)
im = gaussian * t.as_tensor(scipy.datasets.ascent(), dtype=t.complex64)
plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[1,1], [1,1], [0,0]]))
if show_plot:
plt.show()
tests/tools/test_propagators.py
+5 -5
View File
@@ -5,7 +5,7 @@ from cdtools.tools import image_processing
import numpy as np
import torch as t
import pytest
import scipy.misc
import scipy.datasets
from scipy import stats
from matplotlib import pyplot as plt
@@ -13,7 +13,7 @@ from matplotlib import pyplot as plt
@pytest.fixture(scope='module')
def exit_waves_1():
# Import scipy test image and add a random phase
obj = scipy.misc.ascent()[0:64,0:64].astype(np.complex128)
obj = scipy.datasets.ascent()[0:64,0:64].astype(np.complex128)
arr = np.random.random_sample((64,64))
obj *= (arr+(1-arr**2)**.5*1j)
obj = t.as_tensor(obj)
@@ -53,7 +53,7 @@ def test_generate_high_NA_k_intensity_map():
distance = 1#6e-3
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance, opt_for_fft=False)
distance)
k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
rs_basis, basis, shape, distance, wavelength,
@@ -84,7 +84,7 @@ def test_generate_high_NA_k_intensity_map():
distance = 0.01#6e-3
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance, opt_for_fft=False)
distance)
rs_basis_tilted = rs_basis.clone()
rs_basis_tilted[2,1] = rs_basis_tilted[0,1]
@@ -271,7 +271,7 @@ def test_near_field():
asp = propagators.generate_angular_spectrum_propagator(
E0.shape, spacing, wavelength, z)
asp[10,10].backward()
t.real(asp[10,10]).backward()
assert z.grad != 0
assert spacing.grad[0] != 0
assert wavelength.grad !=0