cSAXS/cdtools
0
0 Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-11-14 02:03:53 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into FRC

Browse Source
This commit is contained in:
Damian Guenzing authored and GitHub committed 2025-07-14 11:58:07 -07:00
2 parents 027831d3a8 25ca3a9317
commit 9264228c89
14 files changed
+1009 -1046
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.flake8
+2
View File
@@ -0,0 +1,2 @@
[flake8]
ignore = E501, W503
tests/conftest.py
+71 -70
View File
@@ -1,9 +1,10 @@
import numpy as np
import torch as t
import h5py
import pytest
import datetime
import h5py
import numpy as np
import pytest
import torch as t
#
#
@@ -47,6 +48,7 @@ def pytest_collection_modifyitems(config, items):
if "slow" in item.keywords:
item.add_marker(skip_slow)
@pytest.fixture
def reconstruction_device(request):
return request.config.getoption("--reconstruction_device")
@@ -57,7 +59,6 @@ def show_plot(request):
return request.config.getoption("--plot")
@pytest.fixture(scope='module')
def ptycho_cxi_1():
"""Creates an example file for CXI ptychography. This file is defined
@@ -67,11 +68,11 @@ def ptycho_cxi_1():
"""
expected = {}
f = h5py.File('ptycho_cxi_1','w',driver='core',backing_store=False)
f = h5py.File('ptycho_cxi_1', 'w', driver='core', backing_store=False)
# Start by defining the basic structure
f.create_dataset('cxi_version', data=150)
f.create_dataset('number_of_entries',data=1)
f.create_dataset('number_of_entries', data=1)
# Then define a bunch of metadata for entry_1
e1f = f.create_group('entry_1')
@@ -104,19 +105,19 @@ def ptycho_cxi_1():
s1f['concentration'] = s1e['concentration']
s1e['mass'] = np.float32(np.random.rand())
s1f['mass'] = s1e['mass']
s1e['temperature'] = np.float32(np.random.rand()*100)
s1e['temperature'] = np.float32(np.random.rand() * 100)
s1f['temperature'] = s1e['temperature']
s1e['thickness'] = np.float32(np.random.rand()*1e-7)
s1e['thickness'] = np.float32(np.random.rand() * 1e-7)
s1f['thickness'] = s1e['thickness']
s1e['unit_cell_volume'] = np.float32(np.random.rand() * 1e-27)
s1f['unit_cell_volume'] = s1e['unit_cell_volume']
s1e['unit_cell'] = np.array([1,1,1,90,90,90]).astype(np.float32)
s1f.create_dataset('unit_cell',data = s1e['unit_cell'])
s1e['unit_cell'] = np.array([1, 1, 1, 90, 90, 90]).astype(np.float32)
s1f.create_dataset('unit_cell', data=s1e['unit_cell'])
i1f = e1f.create_group('instrument_1')
source1f = i1f.create_group('source_1')
energy = np.float32(1.3618e-16) #Joules, = 850 eV
energy = np.float32(1.3618e-16) # Joules, = 850 eV
source1f['energy'] = energy
expected['wavelength'] = np.float32(1.9864459e-25) / energy
source1f['wavelength'] = expected['wavelength']
@@ -126,48 +127,48 @@ def ptycho_cxi_1():
d1e = expected['detector']
d1e['distance'] = np.float32(0.3)
d1f['distance'] = d1e['distance']
d1e['basis'] = np.array([[0,-30e-6,0],
[-20e-6,0,0]]).astype(np.float32).transpose()
d1f.create_dataset('basis_vectors',data=d1e['basis'])
d1e['basis'] = np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).astype(np.float32).transpose()
d1f.create_dataset('basis_vectors', data=d1e['basis'])
d1f['x_pixel_size'] = np.float32(20e-6)
d1f['y_pixel_size'] = np.float32(30e-6)
d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
d1e['corner'] = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
d1f.create_dataset('corner_position', data=d1e['corner'])
# Remember the format for the CXI file differs from the format used
# internally
mask = np.zeros((256,256)).astype(np.int32)
mask[5,8] = 1
expected['mask'] = np.ones((256,256)).astype(bool)
expected['mask'][5,8] = 0
d1f.create_dataset('mask',data=mask)
mask = np.zeros((256, 256)).astype(np.int32)
mask[5, 8] = 1
expected['mask'] = np.ones((256, 256)).astype(bool)
expected['mask'][5, 8] = 0
d1f.create_dataset('mask', data=mask)
# There is no specification for this in the CXI file format :(
qe_mask = np.ones((256,256)).astype(np.float32)
qe_mask = np.ones((256, 256)).astype(np.float32)
expected['qe_mask'] = qe_mask
d1f.create_dataset('qe_mask',data=qe_mask)
d1f.create_dataset('qe_mask', data=qe_mask)
# Create an initial background
dark = np.ones((256,256)) * 0.01
dark = np.ones((256, 256)) * 0.01
expected['dark'] = dark
d1f.create_dataset('data_dark', data=dark)
data1f = e1f.create_group('data_1')
data = np.random.rand(100,256,256).astype(np.float32)
data = np.random.rand(100, 256, 256).astype(np.float32)
expected['data'] = data
d1f.create_dataset('data',data=data)
d1f.create_dataset('data', data=data)
data1f['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data')
d1f['data'].attrs['axes'] = np.bytes_('translation:y:x')
expected['axes'] = ['translation','y','x']
expected['axes'] = ['translation', 'y', 'x']
g1f = s1f.create_group('geometry_1')
orientation = np.array([1.,0,0,0,1,0])
orientation = np.array([1., 0, 0, 0, 1, 0])
g1f.create_dataset('orientation', data=orientation)
s1e['orientation'] = np.array([[1.,0,0],[0,1,0],[0,0,1]])
translations = np.arange(300).reshape((100,3)).astype(np.float32)
g1f.create_dataset('translation',data=translations)
s1e['orientation'] = np.array([[1., 0, 0], [0, 1, 0], [0, 0, 1]])
translations = np.arange(300).reshape((100, 3)).astype(np.float32)
g1f.create_dataset('translation', data=translations)
data1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
d1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
expected['translations'] = -translations
@@ -193,11 +194,11 @@ def ptycho_cxi_2():
"""
expected = {}
f = h5py.File('ptycho_cxi_2','w',driver='core',backing_store=False)
f = h5py.File('ptycho_cxi_2', 'w', driver='core', backing_store=False)
# Start by defining the basic structure
f.create_dataset('cxi_version', data=150)
f.create_dataset('number_of_entries',data=1)
f.create_dataset('number_of_entries', data=1)
# Then define a bunch of metadata for entry_1
e1f = f.create_group('entry_1')
@@ -210,13 +211,13 @@ def ptycho_cxi_2():
s1f = e1f.create_group('sample_1')
expected['sample info'] = {}
s1e = expected['sample info']
s1e['temperature'] = np.float32(np.random.rand()*100)
s1e['temperature'] = np.float32(np.random.rand() * 100)
s1f['temperature'] = s1e['temperature']
i1f = e1f.create_group('instrument_1')
source1f = i1f.create_group('source_1')
energy = np.float32(1.3618e-16) #Joules, = 850 eV
energy = np.float32(1.3618e-16) # Joules, = 850 eV
expected['wavelength'] = np.float32(1.9864459e-25) / energy
source1f['wavelength'] = expected['wavelength']
@@ -224,11 +225,11 @@ def ptycho_cxi_2():
expected['detector'] = {}
d1e = expected['detector']
d1e['distance'] = np.float32(0.3)
d1e['basis'] = np.array([[0,-30e-6,0],
[-20e-6,0,0]]).astype(np.float32).transpose()
d1e['basis'] = np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).astype(np.float32).transpose()
d1f['x_pixel_size'] = np.float32(20e-6)
d1f['y_pixel_size'] = np.float32(30e-6)
d1e['corner'] = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
d1e['corner'] = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
d1f.create_dataset('corner_position', data=d1e['corner'])
# Remember the format for the CXI file differs from the format used
@@ -236,24 +237,23 @@ def ptycho_cxi_2():
expected['mask'] = None
expected['qe_mask'] = None
# Test with a set of dark images
dark = np.ones((10,256,256)) * 0.01
expected['dark'] = np.nanmean(dark,axis=0)
dark = np.ones((10, 256, 256)) * 0.01
expected['dark'] = np.nanmean(dark, axis=0)
d1f.create_dataset('data_dark', data=dark)
e1f.create_group('data_1')
data1f = e1f.create_group('data_1')
data = np.random.rand(100,256,256).astype(np.float32)
data = np.random.rand(100, 256, 256).astype(np.float32)
expected['data'] = data
d1f.create_dataset('data',data=data)
d1f.create_dataset('data', data=data)
expected['axes'] = None
g1f = s1f.create_group('geometry_1')
translations = np.arange(300).reshape((100,3)).astype(np.float32)
g1f.create_dataset('translation',data=translations)
translations = np.arange(300).reshape((100, 3)).astype(np.float32)
g1f.create_dataset('translation', data=translations)
expected['translations'] = -translations
yield f, expected
@@ -276,11 +276,11 @@ def ptycho_cxi_3():
"""
expected = {}
f = h5py.File('ptycho_cxi_3','w',driver='core',backing_store=False)
f = h5py.File('ptycho_cxi_3', 'w', driver='core', backing_store=False)
# Start by defining the basic structure
f.create_dataset('cxi_version', data=150)
f.create_dataset('number_of_entries',data=1)
f.create_dataset('number_of_entries', data=1)
# Then define a bunch of metadata for entry_1
e1f = f.create_group('entry_1')
@@ -297,7 +297,7 @@ def ptycho_cxi_3():
i1f = e1f.create_group('instrument_1')
source1f = i1f.create_group('source_1')
energy = np.float32(1.3618e-16) #Joules, = 850 eV
energy = np.float32(1.3618e-16) # Joules, = 850 eV
source1f['energy'] = energy
expected['wavelength'] = np.float32(1.9864459e-25) / energy
@@ -306,41 +306,41 @@ def ptycho_cxi_3():
d1e = expected['detector']
d1e['distance'] = np.float32(0.3)
d1f['distance'] = d1e['distance']
d1e['basis'] = np.array([[0,-30e-6,0],
[-20e-6,0,0]]).astype(np.float32).transpose()
d1f.create_dataset('basis_vectors',data=d1e['basis'])
d1e['basis'] = np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).astype(np.float32).transpose()
d1f.create_dataset('basis_vectors', data=d1e['basis'])
d1e['corner'] = None
# Remember the format for the CXI file differs from the format used
# internally
mask = np.ones((256,256)).astype(np.uint32) * 0x00001000
mask[15,47] = 38
expected['mask'] = np.ones((256,256)).astype(bool)
expected['mask'][15,47] = 0
d1f.create_dataset('mask',data=mask)
mask = np.ones((256, 256)).astype(np.uint32) * 0x00001000
mask[15, 47] = 38
expected['mask'] = np.ones((256, 256)).astype(bool)
expected['mask'][15, 47] = 0
d1f.create_dataset('mask', data=mask)
expected['qe_mask'] = None
expected['dark'] = None
data1f = e1f.create_group('data_1')
data = np.random.rand(100,256,256).astype(np.float32)
data = np.random.rand(100, 256, 256).astype(np.float32)
expected['data'] = data
data1f.create_dataset('data',data=data)
data1f.create_dataset('data', data=data)
data1f['data'].attrs['axes'] = np.bytes_('translation:y:x')
expected['axes'] = ['translation','y','x']
expected['axes'] = ['translation', 'y', 'x']
translations = np.arange(300).reshape((100,3)).astype(np.float32)
data1f.create_dataset('translation',data=translations)
translations = np.arange(300).reshape((100, 3)).astype(np.float32)
data1f.create_dataset('translation', data=translations)
expected['translations'] = -translations
yield f, expected
f.close()
@pytest.fixture(scope='module')
def polarized_ptycho_cxi(ptycho_cxi_1):
f, expected = ptycho_cxi_1
@@ -352,7 +352,7 @@ def polarized_ptycho_cxi(ptycho_cxi_1):
data1f.create_dataset('polarizer_angle', data=expected['polarizer_angle'])
yield f, expected
# As specific issues start to crop up with loading CXI files from different
# beamlines, put a fixture here that replicates the issue so that we can
@@ -374,6 +374,7 @@ def gold_ball_cxi(pytestconfig):
return str(pytestconfig.rootpath) + \
'/examples/example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
@pytest.fixture(scope='module')
def lab_ptycho_cxi(pytestconfig):
return str(pytestconfig.rootpath) + \
@@ -382,8 +383,8 @@ def lab_ptycho_cxi(pytestconfig):
@pytest.fixture(scope='module')
def example_nested_dicts(pytestconfig):
example_tensor = t.as_tensor(np.array([1,4.5,7]))
example_array = np.ones([10,20,30])
example_tensor = t.as_tensor(np.array([1, 4.5, 7]))
example_array = np.ones([10, 20, 30])
example_scalar = 4.5
example_single_element_array = np.array([0.3])
example_string = 'testing'
tests/models/test_fancy_ptycho.py
+16 -18
View File
@@ -1,13 +1,12 @@
import pytest
import cdtools
import torch as t
import cdtools
from matplotlib import pyplot as plt
# Force all reconstructions to use the same RNG seed
t.manual_seed(0)
def test_center_probe(lab_ptycho_cxi):
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
model = cdtools.models.FancyPtycho.from_dataset(
@@ -28,8 +27,8 @@ def test_center_probe(lab_ptycho_cxi):
fourier_model.probe.data = cdtools.tools.propagators.far_field(
base_probe
)
fourier_base_probe = fourier_model.probe.detach().clone()
fourier_model.probe.detach().clone()
fourier_model.center_probes()
fourier_centered_probe = fourier_model.probe.detach().clone()
ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
@@ -45,26 +44,27 @@ def test_center_probe(lab_ptycho_cxi):
atol=1e-4,
rtol=1e-3
)
@pytest.mark.slow
def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
print('\nTesting performance on the standard transmission ptycho dataset')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
n_modes=3,
oversampling=2,
exponentiate_obj=True,
dm_rank=2,
probe_support_radius=120,
propagation_distance=5e-3,
units='mm',
propagation_distance=5e-3,
units='mm',
obj_view_crop=-50,
use_qe_mask=True, # test this in the case where no qe mask is defined
use_qe_mask=True, # test this in the case where no qe mask is defined
)
print('Running reconstruction on provided reconstruction_device,',
reconstruction_device)
model.to(device=reconstruction_device)
@@ -75,11 +75,11 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50):
for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
model.tidy_probes()
if show_plot:
@@ -94,7 +94,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
print('\nTesting performance on the standard gold balls dataset')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
pad = 10
@@ -119,7 +119,7 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
reconstruction_device)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
print(model.report())
if show_plot and model.epoch % 10 == 0:
@@ -135,7 +135,7 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
model.tidy_probes()
if show_plot:
@@ -146,5 +146,3 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
# and choosing a rough value. If it triggers this assertion error,
# something changed to make the final quality worse!
assert model.loss_history[-1] < 0.0001
tests/models/test_simple_ptycho.py
+6 -5
View File
@@ -1,17 +1,18 @@
import pytest
import cdtools
from matplotlib import pyplot as plt
import torch as t
import cdtools
# Force all reconstructions to use the same RNG seed
t.manual_seed(0)
@pytest.mark.slow
def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
model = cdtools.models.SimplePtycho.from_dataset(dataset)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
@@ -19,10 +20,10 @@ def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# If this fails, the reconstruction got worse
assert model.loss_history[-1] < 0.013
tests/test_datasets.py
+89 -108
View File
@@ -18,20 +18,20 @@ from cdtools.tools import data as cdtdata
def test_CDataset_init():
entry_info = {'start_time': datetime.datetime.now(),
'title' : 'A simple test'}
'title': 'A simple test'}
sample_info = {'name': 'A test sample',
'mass' : 3.4,
'unit_cell' : np.array([1,1,1,87,84.5,90])}
'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))
'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))
dataset = CDataset(entry_info, sample_info,
wavelength, detector_geometry, mask)
assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(bool))))
assert t.all(t.eq(dataset.mask, t.tensor(mask.astype(bool))))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
@@ -53,7 +53,7 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
else:
assert dataset.sample_info is not None
assert np.isclose(dataset.wavelength,expected['wavelength'])
assert np.isclose(dataset.wavelength, expected['wavelength'])
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
@@ -64,18 +64,16 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
assert 'corner' in dataset.detector_geometry
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
assert t.all(t.eq(t.tensor(expected['mask']), dataset.mask))
if expected['qe_mask'] is not None:
assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
assert t.all(t.eq(t.tensor(expected['qe_mask']), dataset.qe_mask))
if expected['dark'] is not None:
assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
dataset.background))
def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
for cxi, expected in test_ptycho_cxis:
dataset = CDataset.from_cxi(cxi)
@@ -95,27 +93,24 @@ def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
# Just check one of the loaded attributes
# 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))
assert t.all(t.eq(dataset.mask, read_dataset.mask))
if dataset.qe_mask is not None:
assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
assert t.all(t.eq(dataset.qe_mask, read_dataset.qe_mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
def test_CDataset_to(ptycho_cxi_1):
dataset = CDataset.from_cxi(ptycho_cxi_1[0])
@@ -136,28 +131,28 @@ def test_CDataset_to(ptycho_cxi_1):
def test_Ptycho2DDataset_init():
entry_info = {'start_time': datetime.datetime.now(),
'title' : 'A simple test'}
'title': 'A simple test'}
sample_info = {'name': 'A test sample',
'mass' : 3.4,
'unit_cell' : np.array([1,1,1,87,84.5,90])}
'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))
qe_mask = 1.2*np.ones((256,256), dtype=np.float32)
patterns = np.random.rand(20,256,256)
translations = np.random.rand(20,3)
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
'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))
qe_mask = 1.2 * np.ones((256, 256), dtype=np.float32)
patterns = np.random.rand(20, 256, 256)
translations = np.random.rand(20, 3)
assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
dataset = Ptycho2DDataset(translations, 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
@@ -167,15 +162,15 @@ def test_Ptycho2DDataset_init():
# Also test one with a qe_mask
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask,
qe_mask=qe_mask)
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask,
qe_mask=qe_mask)
assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
assert t.all(t.eq(dataset.qe_mask,t.as_tensor(qe_mask)))
assert t.all(t.eq(dataset.mask, t.BoolTensor(mask)))
assert t.all(t.eq(dataset.qe_mask, t.as_tensor(qe_mask)))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
@@ -199,7 +194,7 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
else:
assert dataset.sample_info is not None
assert np.isclose(dataset.wavelength,expected['wavelength'])
assert np.isclose(dataset.wavelength, expected['wavelength'])
# Just check one of the loaded attributes
assert np.isclose(dataset.detector_geometry['distance'],
@@ -210,18 +205,17 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
assert 'corner' in dataset.detector_geometry
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
assert t.all(t.eq(t.tensor(expected['mask']), dataset.mask))
if expected['qe_mask'] is not None:
assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
assert t.all(t.eq(t.tensor(expected['qe_mask']), dataset.qe_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['data']), dataset.patterns)
assert t.allclose(t.tensor(expected['translations']), dataset.translations)
def test_Ptycho2DDataset_from_cxi_64bit(test_ptycho_cxis):
@@ -279,20 +273,19 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
assert np.isclose(dataset.wavelength, read_dataset.wavelength)
# Just check one of the loaded attributes
# 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))
assert t.all(t.eq(dataset.mask, read_dataset.mask))
if dataset.qe_mask is not None:
assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
assert t.all(t.eq(dataset.qe_mask, read_dataset.qe_mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
@@ -303,7 +296,6 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
def test_Ptycho2DDataset_to(ptycho_cxi_1):
dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
dataset.to(dtype=t.float64)
assert dataset.mask.dtype == t.bool
assert dataset.qe_mask.dtype == t.float64
@@ -327,8 +319,8 @@ def test_Ptycho2DDataset_ops(ptycho_cxi_1):
assert len(dataset) == expected['data'].shape[0]
(idx, translation), 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(translation, t.tensor(expected['translations'][3, :]))
assert t.allclose(pattern, t.tensor(expected['data'][3, :, :]))
def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
@@ -341,12 +333,12 @@ def test_Ptycho2DDataset_get_as(ptycho_cxi_1):
(idx, translation), 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,:]))
t.tensor(expected['translations'][3, :]))
assert t.allclose(pattern.to(device='cpu'),
t.tensor(expected['data'][3,:,:]))
t.tensor(expected['data'][3, :, :]))
def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
@@ -365,18 +357,15 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
masked_patterns = dataset.mask * dataset.patterns
assert t.allclose(
copied_dataset.patterns,
masked_patterns[:,::2,::2] +
masked_patterns[:,1::2,::2] +
masked_patterns[:,::2,1::2] +
masked_patterns[:,1::2,1::2]
masked_patterns[:, ::2, ::2] + masked_patterns[:, 1::2, ::2] + masked_patterns[:, ::2, 1::2] + masked_patterns[:, 1::2, 1::2]
)
if dataset.qe_mask is None:
manually_downsampled_mask = t.logical_and(
t.logical_and(dataset.mask[::2,::2],
dataset.mask[1::2,::2]),
t.logical_and(dataset.mask[::2,1::2],
dataset.mask[1::2,1::2])
t.logical_and(dataset.mask[::2, ::2],
dataset.mask[1::2, ::2]),
t.logical_and(dataset.mask[::2, 1::2],
dataset.mask[1::2, 1::2])
)
assert t.allclose(
copied_dataset.mask,
@@ -384,10 +373,10 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
)
else:
manually_downsampled_mask = t.logical_or(
t.logical_or(dataset.mask[::2,::2],
dataset.mask[1::2,::2]),
t.logical_or(dataset.mask[::2,1::2],
dataset.mask[1::2,1::2])
t.logical_or(dataset.mask[::2, ::2],
dataset.mask[1::2, ::2]),
t.logical_or(dataset.mask[::2, 1::2],
dataset.mask[1::2, 1::2])
)
assert t.allclose(
copied_dataset.mask,
@@ -396,8 +385,7 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
masked_qe_mask = dataset.mask * dataset.qe_mask
manually_downsampled_qe_mask = (
masked_qe_mask[::2,::2] + masked_qe_mask[1::2,::2]
+ masked_qe_mask[::2,1::2] + masked_qe_mask[1::2,1::2]
masked_qe_mask[::2, ::2] + masked_qe_mask[1::2, ::2] + masked_qe_mask[::2, 1::2] + masked_qe_mask[1::2, 1::2]
) / 4
assert t.allclose(
@@ -405,15 +393,11 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
manually_downsampled_qe_mask
)
if dataset.background is not None:
assert t.allclose(
copied_dataset.background,
dataset.background[::2,::2] +
dataset.background[1::2,::2] +
dataset.background[::2,1::2] +
dataset.background[1::2,1::2]
)
dataset.background[::2, ::2] + dataset.background[1::2, ::2] + dataset.background[::2, 1::2] + dataset.background[1::2, 1::2]
)
# And then we just test the shape for a few factors, and check that
# it doesn't fail on edge cases (e.g. factor=1)
@@ -428,10 +412,10 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
assert np.allclose(expected_pattern_shape,
np.array(copied_dataset.patterns.shape))
assert np.allclose(np.array(dataset.mask.shape) // factor,
np.array(copied_dataset.mask.shape))
if dataset.background is not None:
assert np.allclose(np.array(dataset.background.shape) // factor,
np.array(copied_dataset.background.shape))
@@ -468,13 +452,13 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
copied_dataset = deepcopy(dataset)
# Test 1: Complain when the the bounds of an ROI are correctly defined,
# but it does not contain any sample positions inside of it. The
# translations in ptycho_cxi_1 ranges from 0m to -300m in both x and y
# (it looks like a line scan). We select an ROI at (0, -300) which should
# but it does not contain any sample positions inside of it. The
# translations in ptycho_cxi_1 ranges from 0m to -300m in both x and y
# (it looks like a line scan). We select an ROI at (0, -300) which should
# not contain any translation positions.
with pytest.raises(ValueError) as excinfo:
copied_dataset.crop_translations(roi=(0,-5,-295,-300))
assert('(i.e., patterns and translations will be empty)') in str(excinfo.value)
copied_dataset.crop_translations(roi=(0, -5, -295, -300))
assert '(i.e., patterns and translations will be empty)' in str(excinfo.value)
# Test 2: Draw an ROI that's centered in the middle of the x/y translation range
# and make sure that the first and last x/y elements in dataset.translate
@@ -482,40 +466,37 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
# Make tuples that will store the x and y positions. This will be used for
# making permutations of the x/y positions in the ROI later.
x_left, y_top = dataset.translations[10,:2]
x_right, y_bottom = dataset.translations[-11,:2]
x_left, y_top = dataset.translations[10, :2]
x_right, y_bottom = dataset.translations[-11, :2]
x_permutations = ((x_left, x_right), (x_right, x_left))
y_permutations = ((y_top, y_bottom), (y_bottom, y_top))
roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in
itertools.product(x_permutations, y_permutations))
roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in itertools.product(x_permutations, y_permutations))
# Get the dataset
copied_dataset.crop_translations(roi=roi_permutations[0])
# Execute the actual test
assert (copied_dataset.translations[0,0] in x_permutations[0]) and \
(copied_dataset.translations[-1,0] in x_permutations[0])
assert (copied_dataset.translations[0,1] in y_permutations[0]) and \
(copied_dataset.translations[-1,1] in y_permutations[0])
assert (copied_dataset.translations[0, 0] in x_permutations[0]) and \
(copied_dataset.translations[-1, 0] in x_permutations[0])
assert (copied_dataset.translations[0, 1] in y_permutations[0]) and \
(copied_dataset.translations[-1, 1] in y_permutations[0])
# Test 3: Check if the shape of dataset.patterns and dataset.translate is correct
# (designed to be 20 fewer rows here)
# In the future, this should include a check for dataset.intensities once
# an appropriate cxi file is set up for conftest.
expected_patterns_shape = np.concatenate([[dataset.patterns.shape[0] - 20],
dataset.patterns.shape[-2:]])
expected_translations_shape = np.concatenate([[dataset.translations.shape[0] - 20],
dataset.translations.shape[-1:]])
expected_patterns_shape = np.concatenate([[dataset.patterns.shape[0] - 20], dataset.patterns.shape[-2:]])
expected_translations_shape = np.concatenate([[dataset.translations.shape[0] - 20], dataset.translations.shape[-1:]])
assert np.allclose(np.array(copied_dataset.patterns.shape), expected_patterns_shape)
assert np.allclose(np.array(copied_dataset.translations.shape), expected_translations_shape)
# Test 4: Make sure that we always get the same result no matter what order we
# define the left/right and bottom/top values in roi, provided that roi[:2]
# define the left/right and bottom/top values in roi, provided that roi[:2]
# and roi[2:] correspond with the x and y coordinates, respectively.
# Check each permutation
@@ -526,6 +507,6 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
copied_dataset.crop_translations(roi=roi)
# Check if the contents of dataset.patterns and dataset.translate is correct
assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10,:])
assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10, :])
assert t.allclose(copied_dataset.translations, dataset.translations[10:-10,:])
assert t.allclose(copied_dataset.translations, dataset.translations[10:-10, :])
tests/tools/test_analysis.py
+258 -213
View File
@@ -16,10 +16,10 @@ def test_product_svd():
A = np.random.rand(*shape_A) + 1j * np.random.rand(*shape_A)
B = np.random.rand(*shape_B) + 1j * np.random.rand(*shape_B)
AB = np.matmul(A,B)
AB = np.matmul(A, B)
U_1, S_1, Vh_1 = t.linalg.svd(t.as_tensor(AB), full_matrices=False)
U_2, S_2, Vh_2 = analysis.product_svd(t.as_tensor(A),t.as_tensor(B))
U_2, S_2, Vh_2 = analysis.product_svd(t.as_tensor(A), t.as_tensor(B))
check_AB = U_2 @ t.diag_embed(S_2).to(dtype=Vh_2.dtype) @ Vh_2
# So, at a minimum, U S Vh = AB
@@ -29,23 +29,23 @@ def test_product_svd():
# singular vector, so all we can ask for in the comparison is that the
# magnitudes here are
assert np.allclose(S_1[:rank].numpy(), S_2.numpy())
prod_U = U_1[:,:rank].transpose(0,1).conj() @ U_2
prod_Vh = Vh_1[:rank,:] @ Vh_2.transpose(0,1).conj()
prod_U = U_1[:, :rank].transpose(0, 1).conj() @ U_2
prod_Vh = Vh_1[:rank, :] @ Vh_2.transpose(0, 1).conj()
assert np.allclose(t.abs(prod_U).numpy(), np.eye(rank))
assert np.allclose(t.abs(prod_Vh).numpy(), np.eye(rank))
# Confirms that the phases are consistent between the two, I think
# it's redundant with the first check but I'm not sure
assert np.allclose(prod_Vh.numpy(), prod_U.numpy())
# test with numpy
U_3, S_3, Vh_3 = analysis.product_svd(A,B)
U_3, S_3, Vh_3 = analysis.product_svd(A, B)
assert isinstance(U_3, np.ndarray)
assert isinstance(S_3, np.ndarray)
assert isinstance(Vh_3, np.ndarray)
assert np.allclose(S_1[:rank].numpy(), S_3)
prod_U = U_1[:,:rank].transpose(0,1).numpy().conj() @ U_3
prod_Vh = Vh_1[:rank,:].numpy() @ Vh_3.transpose().conj()
prod_U = U_1[:, :rank].transpose(0, 1).numpy().conj() @ U_3
prod_Vh = Vh_1[:rank, :].numpy() @ Vh_3.transpose().conj()
assert np.allclose(np.abs(prod_U), np.eye(rank))
assert np.allclose(np.abs(prod_Vh), np.eye(rank))
# Confirms that the phases are consistent between the two, I think
@@ -56,42 +56,49 @@ def test_product_svd():
def test_orthogonalize_probes():
op = analysis.orthogonalize_probes
probe_xs = np.arange(64) - 32
probe_ys = np.arange(76) - 38
probe_Ys, probe_Xs = np.meshgrid(probe_ys, probe_xs)
probe_Rs = np.sqrt(probe_Xs**2 + probe_Ys**2)
probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2 + 1j)),
3*np.exp(-probe_Rs**2 / (2 * 12**2 - 3j)),
1*np.exp(-probe_Rs**2 / (2 * 15**2))])
probes = np.array(
[
10 * np.exp(-(probe_Rs**2) / (2 * 10**2 + 1j)),
3 * np.exp(-(probe_Rs**2) / (2 * 12**2 - 3j)),
1 * np.exp(-(probe_Rs**2) / (2 * 15**2)),
]
)
weight_matrix_none = None
weight_matrix_single = np.random.randn(1,3) + 1j * np.random.randn(1,3)
weight_matrix_small = np.random.randn(2,3) + 1j * np.random.randn(2,3)
weight_matrix_medium = np.random.randn(3,3) + 1j * np.random.randn(3,3)
weight_matrix_large = np.random.randn(7,3) + 1j * np.random.randn(7,3)
weight_matrix_single = np.random.randn(1, 3) + 1j * np.random.randn(1, 3)
weight_matrix_small = np.random.randn(2, 3) + 1j * np.random.randn(2, 3)
weight_matrix_medium = np.random.randn(3, 3) + 1j * np.random.randn(3, 3)
weight_matrix_large = np.random.randn(7, 3) + 1j * np.random.randn(7, 3)
weight_matrices = [
weight_matrix_none,
weight_matrix_single,
weight_matrix_small,
weight_matrix_medium,
weight_matrix_large
weight_matrix_large,
]
for weight_matrix in weight_matrices:
ortho_probes_np, rwm_np = op(probes, weight_matrix=weight_matrix,
return_reexpressed_weights=True)
ortho_probes_np, rwm_np = op(
probes, weight_matrix=weight_matrix, return_reexpressed_weights=True
)
assert isinstance(ortho_probes_np, np.ndarray)
assert isinstance(rwm_np, np.ndarray)
probes_t = t.as_tensor(probes)
wm_t = (t.as_tensor(weight_matrix) if weight_matrix is not None
else weight_matrix)
ortho_probes_t, rwm_t = op(probes_t, weight_matrix=wm_t,
return_reexpressed_weights=True)
probes_t = t.as_tensor(probes)
wm_t = (
t.as_tensor(weight_matrix) if weight_matrix is not None else weight_matrix
)
ortho_probes_t, rwm_t = op(
probes_t, weight_matrix=wm_t, return_reexpressed_weights=True
)
assert t.is_tensor(ortho_probes_t)
assert t.is_tensor(rwm_t)
@@ -106,15 +113,17 @@ def test_orthogonalize_probes():
realized_probes = probes
calculated_probes = np.tensordot(rwm_np, ortho_probes_np, axes=1)
assert np.allclose(realized_probes, calculated_probes)
# Now we test if the orthogonalized probes are orthogonalized
reshaped_probes = ortho_probes_np.reshape(
(ortho_probes_np.shape[0],
ortho_probes_np.shape[1] * ortho_probes_np.shape[2]))
products = np.matmul(reshaped_probes,
reshaped_probes.conj().transpose())
(
ortho_probes_np.shape[0],
ortho_probes_np.shape[1] * ortho_probes_np.shape[2],
)
)
products = np.matmul(reshaped_probes, reshaped_probes.conj().transpose())
if weight_matrix is not None:
output_nmodes = min(weight_matrix.shape[0], probes.shape[0])
@@ -122,23 +131,22 @@ def test_orthogonalize_probes():
output_nmodes = probes.shape[0]
for i in range(output_nmodes):
for j in range(i+1, output_nmodes):
assert np.isclose(products[i,j], 0)
for j in range(i + 1, output_nmodes):
assert np.isclose(products[i, j], 0)
# And now we test if they multiply to the same density matrix as
# the original probes + weight matrix
reshaped_realized_probes = realized_probes.reshape(
(realized_probes.shape[0],
realized_probes.shape[1] * realized_probes.shape[2]))
(
realized_probes.shape[0],
realized_probes.shape[1] * realized_probes.shape[2],
)
)
dm_original = np.matmul(
reshaped_realized_probes.conj().transpose(),
reshaped_realized_probes
)
dm_output = np.matmul(
reshaped_probes.conj().transpose(),
reshaped_probes
reshaped_realized_probes.conj().transpose(), reshaped_realized_probes
)
dm_output = np.matmul(reshaped_probes.conj().transpose(), reshaped_probes)
assert np.allclose(dm_original, dm_output)
# And finally, we confirm that what we have are the eigenvectors/values
@@ -151,53 +159,58 @@ def test_orthogonalize_probes():
# is undefined
assert np.allclose(np.abs(cross_products), np.abs(products))
def test_standardize():
# Start by making a probe and object that should meet the standardization
# conditions
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
probe = initializers.gaussian((230, 240), (20, 20), curvature=(0.01, 0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe) ** 2))
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.isclose(1, np.sum(np.abs(probe) ** 2) / len(probe.ravel()))
assert np.angle(np.sum(probe)) < 2e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
(obj.shape[1]//8)*3:(obj.shape[1]//8)*5]
obj = 30 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5))
obj_slice = np.s_[
(obj.shape[0] // 8) * 3:(obj.shape[0] // 8) * 5,
(obj.shape[1] // 8) * 3:(obj.shape[1] // 8) * 5,
]
obj = obj * np.exp(-1j * np.angle(np.sum(obj[obj_slice])))
assert np.isclose(0,np.angle(np.sum(obj[obj_slice])))
assert np.isclose(0, np.angle(np.sum(obj[obj_slice])))
# Then make a nonstandard version of them and standardize it
# First, don't add a phase ramp and test
test_probe = probe * 37.6 * np.exp(1j*0.35)
test_obj = obj / 37.6 * np.exp(1j*1.43)
test_probe = probe * 37.6 * np.exp(1j * 0.35)
test_obj = obj / 37.6 * np.exp(1j * 1.43)
s_probe, s_obj = analysis.standardize(test_probe, test_obj)
assert np.allclose(probe, s_probe)
assert np.allclose(obj, s_obj)
# Test that it works on torch tensors
s_probe, s_obj = analysis.standardize(t.as_tensor(test_probe,dtype=t.complex64), t.as_tensor(test_obj,dtype=t.complex64))
s_probe, s_obj = analysis.standardize(
t.as_tensor(test_probe, dtype=t.complex64),
t.as_tensor(test_obj, dtype=t.complex64),
)
s_probe = s_probe.numpy()
s_obj = s_obj.numpy()
assert np.allclose(probe, s_probe)
assert np.allclose(obj, s_obj)
# Then do one with a phase ramp
phase_ramp_dir = (np.random.rand(2) - 0.5)
phase_ramp_dir = np.random.rand(2) - 0.5
probe_Xs, probe_Ys = np.mgrid[:probe.shape[0],:probe.shape[1]]
phase_ramp = np.exp(1j*probe_Ys * phase_ramp_dir[1]+
1j*probe_Xs * phase_ramp_dir[0])
probe_Xs, probe_Ys = np.mgrid[: probe.shape[0], : probe.shape[1]]
phase_ramp = np.exp(
1j * probe_Ys * phase_ramp_dir[1] + 1j * probe_Xs * phase_ramp_dir[0]
)
test_probe = test_probe * phase_ramp
obj_Xs, obj_Ys = np.mgrid[:obj.shape[0],:obj.shape[1]]
obj_phase_ramp = np.exp(-1j*obj_Ys * phase_ramp_dir[1]+
-1j*obj_Xs * phase_ramp_dir[0])
obj_Xs, obj_Ys = np.mgrid[: obj.shape[0], : obj.shape[1]]
obj_phase_ramp = np.exp(
-1j * obj_Ys * phase_ramp_dir[1] + -1j * obj_Xs * phase_ramp_dir[0]
)
test_obj = test_obj * obj_phase_ramp
s_probe, s_obj = analysis.standardize(test_probe, test_obj, correct_ramp=True)
@@ -206,12 +219,12 @@ def test_standardize():
assert np.max(s_obj - obj) / np.max(np.abs(obj)) < 1e-4
# Finally a test with the phase ramp and multiple probes
subdominant_probe = 0.1*np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
subdominant_probe = (0.1 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5)))
subdominant_probe = subdominant_probe * np.exp(-1j * np.angle(np.sum(subdominant_probe)))
test_subdominant_probe = subdominant_probe * 37.6
test_subdominant_probe = test_subdominant_probe * phase_ramp
incoh_probe = np.array([test_probe,test_subdominant_probe])
incoh_probe = np.array([test_probe, test_subdominant_probe])
s_probe, s_obj = analysis.standardize(incoh_probe, test_obj, correct_ramp=True)
@@ -220,7 +233,6 @@ def test_standardize():
assert np.max(s_probe[1] - subdominant_probe) / np.max(np.abs(subdominant_probe)) < 1e-4
def test_synthesize_reconstructions():
# I can only really test for a lack of failures, so I think my plan
# will be to create a dataset that just needs to be added and see that
@@ -228,140 +240,153 @@ def test_synthesize_reconstructions():
# Start by making a probe and object that should meet the standardization
# conditions
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
probe = initializers.gaussian((230, 240), (20, 20), curvature=(0.01, 0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe) ** 2))
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.isclose(1, np.sum(np.abs(probe) ** 2) / len(probe.ravel()))
assert np.abs(np.angle(np.sum(probe))) < 2e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
(obj.shape[1]//8)*3:(obj.shape[1]//8)*5]
obj = 30 * np.random.rand(230, 240) * np.exp(1j * (np.random.rand(230, 240) - 0.5))
obj_slice = np.s_[
(obj.shape[0] // 8) * 3:(obj.shape[0] // 8) * 5,
(obj.shape[1] // 8) * 3:(obj.shape[1] // 8) * 5,
]
obj = obj * np.exp(-1j * np.angle(np.sum(obj[obj_slice])))
assert np.isclose(0,np.angle(np.sum(obj[obj_slice])))
assert np.isclose(0, np.angle(np.sum(obj[obj_slice])))
# Now I make stacks of identical probes and objects
probes = [probe,probe,probe,probe]
probes = [probe, probe, probe, probe]
probe = np.copy(probe)
objects = [obj,obj,obj,obj]
objects = [obj, obj, obj, obj]
obj = np.copy(obj)
s_probe, s_obj, obj_stack = analysis.synthesize_reconstructions(probes,objects)
s_probe, s_obj, obj_stack = analysis.synthesize_reconstructions(probes, objects)
assert np.max(s_probe - probe) < 2e-5
assert np.max(s_obj - obj) < 2e-5
for t_obj in obj_stack:
assert np.max(t_obj - obj) < 5e-5
def test_calc_consistency_prtf():
# Create an object with a specific structure
obj = 30 * np.random.rand(1030,1040) * np.exp(1j * (np.random.rand(1030,1040) - 0.5))
obj = (30 * np.random.rand(1030, 1040) * np.exp(1j * (np.random.rand(1030, 1040) - 0.5)))
#
synth_obj = np.sqrt(0.7) * obj
obj_stack = [obj,obj,obj,obj]
obj_stack = [obj, obj, obj, obj]
basis = np.array([[0, 2, 0], [3, 0, 0]])
basis = np.array([[0,2,0],
[3,0,0]])
freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis)
assert np.allclose(prtf, 0.7)
freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis, nbins=30)
assert np.allclose(prtf, 0.7)
# Check that it also works with torch input
t_synth_obj = t.as_tensor(synth_obj)
t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, basis, nbins=30)
freqs, prtf = analysis.calc_consistency_prtf(
t_synth_obj, t_obj_stack, basis, nbins=30
)
assert np.allclose(prtf.numpy(), 0.7)
# And also when the basis is in torch
t_synth_obj = t.as_tensor(synth_obj)
t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30)
freqs, prtf = analysis.calc_consistency_prtf(
t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30
)
assert np.allclose(prtf.numpy(), 0.7)
# Check that is uses the right number of bins
assert len(prtf) == 30
assert len(freqs) == 30
# Check that the maximum frequency is correct for the basis
assert np.isclose(freqs[-1]-freqs[-2] + freqs[-1], np.sqrt(1/4**2 + 1/6**2))
assert np.isclose(freqs[-1] - freqs[-2] + freqs[-1], np.sqrt(1 / 4**2 + 1 / 6**2))
def test_calc_deconvolved_cross_correlation():
obj1 = np.random.rand(200,300) + 1j * np.random.rand(200,300)
obj2 = np.random.rand(200,300) + 1j * np.random.rand(200,300)
obj1 = np.random.rand(200, 300) + 1j * np.random.rand(200, 300)
obj2 = np.random.rand(200, 300) + 1j * np.random.rand(200, 300)
cor_fft = np.fft.fft2(obj1) * np.conj(np.fft.fft2(obj2))
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
np_cor = np.fft.ifft2(cor_fft / np.abs(cor_fft))
# test with numpy inputs
test_cor = analysis.calc_deconvolved_cross_correlation(obj1,obj2, im_slice=np.s_[:,:])
test_cor = analysis.calc_deconvolved_cross_correlation(
obj1, obj2, im_slice=np.s_[:, :]
)
assert np.allclose(test_cor, np_cor)
# test with pytorch inputs
obj1_t = t.as_tensor(obj1)
obj2_t = t.as_tensor(obj2)
test_cor_t = analysis.calc_deconvolved_cross_correlation(obj1_t,obj2_t, im_slice=np.s_[:,:])
test_cor_t = analysis.calc_deconvolved_cross_correlation(
obj1_t, obj2_t, im_slice=np.s_[:, :]
)
assert np.allclose(test_cor_t.numpy(), np_cor)
def test_calc_frc():
obj1 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
obj2 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
obj1 = np.random.rand(270, 230) + 1j * np.random.rand(270, 230)
obj2 = np.random.rand(270, 230) + 1j * np.random.rand(270, 230)
basis = np.array([[0,2,0],
[3,0,0]])
basis = np.array([[0, 2, 0], [3, 0, 0]])
nbins = 100
snr = 2
cor_fft = np.fft.fftshift(np.fft.fft2(obj1[10:-10,20:-20])) * \
np.fft.fftshift(np.conj(np.fft.fft2(obj2[10:-10,20:-20])))
F1 = np.abs(np.fft.fftshift(np.fft.fft2(obj1[10:-10,20:-20])))**2
F2 = np.abs(np.fft.fftshift(np.fft.fft2(obj2[10:-10,20:-20])))**2
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[0],d=di))
j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
cor_fft = np.fft.fftshift(np.fft.fft2(obj1[10:-10, 20:-20])) * np.fft.fftshift(
np.conj(np.fft.fft2(obj2[10:-10, 20:-20]))
)
numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft)
denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1)
denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2)
n_pix, bins = np.histogram(Rs,bins=nbins)
F1 = np.abs(np.fft.fftshift(np.fft.fft2(obj1[10:-10, 20:-20]))) ** 2
F2 = np.abs(np.fft.fftshift(np.fft.fft2(obj2[10:-10, 20:-20]))) ** 2
di = np.linalg.norm(basis[:, 0])
dj = np.linalg.norm(basis[:, 1])
i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[0], d=di))
j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[1], d=dj))
Js, Is = np.meshgrid(j_freqs, i_freqs)
Rs = np.sqrt(Is**2 + Js**2)
numerator, bins = np.histogram(Rs, bins=nbins, weights=cor_fft)
denominator_F1, bins = np.histogram(Rs, bins=nbins, weights=F1)
denominator_F2, bins = np.histogram(Rs, bins=nbins, weights=F2)
n_pix, bins = np.histogram(Rs, bins=nbins)
bins = bins[:-1]
frc = numerator / np.sqrt(denominator_F1*denominator_F2)
# This moves from combined-image SNR to single-image SNR
frc = numerator / np.sqrt(denominator_F1 * denominator_F2)
# This moves from combined-image SNR to single-image SNR
snr /= 2
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
(1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / (
1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix)
)
test_bins, test_frc, test_threshold = analysis.calc_frc(
obj1, obj2, basis, im_slice=np.s_[10:-10,20:-20],
nbins=100, snr=2, limit='corner')
obj1,
obj2,
basis,
im_slice=np.s_[10:-10, 20:-20],
nbins=100,
snr=2,
limit="corner",
)
assert np.allclose(bins, test_bins)
assert np.allclose(frc, test_frc)
assert np.allclose(threshold, test_threshold)
@@ -375,33 +400,40 @@ def test_calc_frc():
obj1_torch,
obj2_torch,
basis_torch,
im_slice=np.s_[10:-10,20:-20], nbins=100, snr=2, limit='corner')
im_slice=np.s_[10:-10, 20:-20],
nbins=100,
snr=2,
limit="corner",
)
assert np.allclose(bins, test_bins_t.numpy())
assert np.allclose(frc, test_frc_t.numpy())
assert np.allclose(threshold, test_threshold_t.numpy())
def test_calc_rms_error():
field_1 = t.rand(14,19, dtype=t.complex64)
field_2 = t.rand(14,19, dtype=t.complex64)
field_1 = t.rand(14, 19, dtype=t.complex64)
field_2 = t.rand(14, 19, dtype=t.complex64)
# Check that the calculation is insensitive to phase
assert t.allclose(analysis.calc_rms_error(field_1, field_2),
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2))
assert t.allclose(
analysis.calc_rms_error(field_1, field_2),
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2),
)
# And that it is sensitive to phase if we turn off the
# And that it is sensitive to phase if we turn off the
assert not t.allclose(
analysis.calc_rms_error(field_1, field_2, align_phases=False),
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2,
align_phases=False))
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2, align_phases=False),
)
# Check that the result is positive
assert analysis.calc_rms_error(field_1, field_2) > 0
# And that it is a smaller number with align_phases on
assert (analysis.calc_rms_error(field_1, field_2) <=
analysis.calc_rms_error(field_1, field_2, align_phases=False))
assert analysis.calc_rms_error(field_1, field_2) <= analysis.calc_rms_error(
field_1, field_2, align_phases=False
)
# Now we check against an explicit implementation:
gamma = field_1 * t.conj(field_2)
@@ -409,115 +441,127 @@ def test_calc_rms_error():
# This is an alternate way of doing the calculation. Actually, would this
# be a better implementation anyway? Probably no difference tbh.
rms_error_nophase = t.sqrt((t.mean(t.abs(field_1)**2) +
t.mean(t.abs(field_2)**2) -
2 * t.abs(t.mean(field_1 * t.conj(field_2)))))
assert t.allclose(rms_error_nophase,
analysis.calc_rms_error(field_1, field_2))
rms_error_nophase = t.sqrt(
(
t.mean(t.abs(field_1) ** 2)
+ t.mean(t.abs(field_2) ** 2)
- 2 * t.abs(t.mean(field_1 * t.conj(field_2)))
)
)
assert t.allclose(rms_error_nophase, analysis.calc_rms_error(field_1, field_2))
rms_error_phase = t.sqrt((t.mean(t.abs(field_1)**2) +
t.mean(t.abs(field_2)**2) -
2 * t.real(t.mean(field_1 * t.conj(field_2)))))
rms_error_phase = t.sqrt(
(
t.mean(t.abs(field_1) ** 2)
+ t.mean(t.abs(field_2) ** 2)
- 2 * t.real(t.mean(field_1 * t.conj(field_2)))
)
)
assert t.allclose(rms_error_phase,
analysis.calc_rms_error(field_1, field_2,
align_phases=False))
assert t.allclose(
rms_error_phase, analysis.calc_rms_error(field_1, field_2, align_phases=False)
)
# Now let's test that it works along a dimension:
field_1 = t.rand(3,14,19, dtype=t.complex64)
field_2 = t.rand(3,14,19, dtype=t.complex64)
field_1 = t.rand(3, 14, 19, dtype=t.complex64)
field_2 = t.rand(3, 14, 19, dtype=t.complex64)
result = analysis.calc_rms_error(field_1, field_2, normalize=True)
assert (result.shape == t.Size([3]))
assert result.shape == t.Size([3])
for i in range(3):
assert t.allclose(analysis.calc_rms_error(field_1[i],
field_2[i],
normalize=True),
result[i])
assert t.allclose(
analysis.calc_rms_error(field_1[i], field_2[i], normalize=True), result[i]
)
def test_calc_fidelity():
fields_1 = t.rand(2,30,17, dtype=t.complex128)
fields_2 = t.rand(3,30,17, dtype=t.complex128)
fields_1 = t.rand(2, 30, 17, dtype=t.complex128)
fields_2 = t.rand(3, 30, 17, dtype=t.complex128)
dm_1 = t.reshape(fields_1, (2,-1))
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()
dm_1 = t.reshape(fields_1, (2, -1))
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()
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)
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)
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))
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)
field_3 = t.rand(1,30,17, dtype=t.complex128)
fields_1 = t.rand(3, 3, 30, 17, dtype=t.complex128)
fields_2 = t.rand(3, 1, 30, 17, dtype=t.complex128)
field_3 = t.rand(1, 30, 17, dtype=t.complex128)
fidelities = analysis.calc_fidelity(fields_1, fields_2)
fidelities_2 = analysis.calc_fidelity(fields_1, field_3)
for i in range(3):
assert t.isclose(analysis.calc_fidelity(fields_1[i], fields_2[i]),
fidelities[i])
assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3),
fidelities_2[i])
assert t.isclose(
analysis.calc_fidelity(fields_1[i], fields_2[i]), fidelities[i]
)
assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3), fidelities_2[i])
# Check that the diensionality argument works
fields_1 = t.rand(3,2,12, dtype=t.complex128)
fields_2 = t.rand(3,2,12, dtype=t.complex128)
fields_1 = t.rand(3, 2, 12, dtype=t.complex128)
fields_2 = t.rand(3, 2, 12, dtype=t.complex128)
assert (analysis.calc_fidelity(fields_1, fields_2, dims=1).shape
== t.Size([3]))
fields_1 = t.rand(3,2,12,4,5, dtype=t.complex128)
fields_2 = t.rand(3,2,12,4,5, dtype=t.complex128)
assert analysis.calc_fidelity(fields_1, fields_2, dims=1).shape == t.Size([3])
fields_1 = t.rand(3, 2, 12, 4, 5, dtype=t.complex128)
fields_2 = t.rand(3, 2, 12, 4, 5, dtype=t.complex128)
assert analysis.calc_fidelity(fields_1, fields_2, dims=3).shape == t.Size([3])
assert (analysis.calc_fidelity(fields_1, fields_2, dims=3).shape
== t.Size([3]))
def test_calc_generalized_rms_error():
# Test that it matches the rms error 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(analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_rms_error(fields_1[0], fields_2[0],
align_phases=True))
fields_1 = t.rand(1, 30, 17, dtype=t.complex128)
fields_2 = t.rand(1, 30, 17, dtype=t.complex128)
assert t.isclose(
analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_rms_error(fields_1[0], fields_2[0], align_phases=True),
)
# Test that it is independent of field order
fields_1 = t.rand(5,30,17, dtype=t.complex128)
fields_2 = t.rand(3,30,17, dtype=t.complex128)
fields_1 = t.rand(5, 30, 17, dtype=t.complex128)
fields_2 = t.rand(3, 30, 17, dtype=t.complex128)
fields_3 = fields_2.flip(0)
assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_generalized_rms_error(fields_1, fields_3))
assert t.isclose(
analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_generalized_rms_error(fields_1, fields_3),
)
# Test with leading dimensions
fields_1 = t.rand(3,4,2,10,17, dtype=t.complex128)
fields_2 = t.rand(3,4,3,10,17, dtype=t.complex128)
assert (analysis.calc_generalized_rms_error(fields_1, fields_2).shape
== t.Size([3,4]))
fields_1 = t.rand(3, 4, 2, 10, 17, dtype=t.complex128)
fields_2 = t.rand(3, 4, 3, 10, 17, dtype=t.complex128)
assert analysis.calc_generalized_rms_error(fields_1, fields_2).shape == t.Size(
[3, 4]
)
# And test with different number of dimensions dims
# Test that it is independent of field order
fields_1 = t.rand(3,6,17, dtype=t.complex128)
fields_2 = t.rand(3,1,17, dtype=t.complex128)
fields_1 = t.rand(3, 6, 17, dtype=t.complex128)
fields_2 = t.rand(3, 1, 17, dtype=t.complex128)
fields_3 = fields_2.flip(0)
assert (analysis.calc_generalized_rms_error(fields_1, fields_2, dims=1).shape == t.Size([3]))
@@ -620,3 +664,4 @@ def test_line_based_frc():
line_based_frc(np.random.rand(30, 17, 3), np.random.rand(30, 17, 3))
line_based_frc(np.random.rand(30, 17, 3), np.random.rand(30, 17))
line_based_frc("string", np.random.rand(30, 17))
tests/tools/test_data.py
+76 -87
View File
@@ -1,60 +1,55 @@
from cdtools.tools import data
import numpy as np
import torch as t
import h5py
import pytest
import os
import datetime
import numbers
import pathlib
import h5py
import numpy as np
import torch as t
from cdtools.tools import data
#
# We start with a bunch of tests of the data loading capabilities
#
def test_get_entry_info(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
entry_info = data.get_entry_info(cxi)
for key in expected['entry metadata']:
assert entry_info[key] == expected['entry metadata'][key]
def test_get_sample_info(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
sample_info = data.get_sample_info(cxi)
if sample_info is None and \
('sample info' not in expected or
expected['sample info'] is None):
('sample info' not in expected or expected['sample info'] is None):
# Valid if no sample info is defined at all
continue
for key in expected['sample info']:
if isinstance(expected['sample info'][key],np.ndarray):
if isinstance(expected['sample info'][key], np.ndarray):
assert np.allclose(sample_info[key],
expected['sample info'][key])
else:
assert sample_info[key] == expected['sample info'][key]
def test_get_wavelength(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
assert np.isclose(expected['wavelength'],data.get_wavelength(cxi))
assert np.isclose(expected['wavelength'], data.get_wavelength(cxi))
def test_get_detector_geometry(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
distance, basis, corner = data.get_detector_geometry(cxi)
assert np.isclose(distance,expected['detector']['distance'])
assert np.allclose(basis,expected['detector']['basis'])
assert np.isclose(distance, expected['detector']['distance'])
assert np.allclose(basis, expected['detector']['basis'])
if isinstance(expected['detector']['corner'], np.ndarray):
assert np.allclose(corner, expected['detector']['corner'])
else:
assert corner == expected['detector']['corner']
def test_get_mask(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
mask = data.get_mask(cxi)
@@ -62,7 +57,7 @@ def test_get_mask(test_ptycho_cxis):
continue
assert np.all(mask == expected['mask'])
def test_get_qe_mask(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
qe_mask = data.get_qe_mask(cxi)
@@ -78,8 +73,8 @@ def test_get_dark(test_ptycho_cxis):
assert expected['dark'] is None
else:
assert np.allclose(dark, expected['dark'])
def test_get_data(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
patterns, axes = data.get_data(cxi)
@@ -99,8 +94,6 @@ def test_get_ptycho_translations(test_ptycho_cxis):
assert np.allclose(data.get_ptycho_translations(cxi),
expected['translations'])
#
# Then, write a test for the data saving. It should create a .cxi file
# using the data seving tools, and then check that when read with the
@@ -110,13 +103,13 @@ def test_get_ptycho_translations(test_ptycho_cxis):
def test_create_cxi(tmp_path):
data.create_cxi(tmp_path / 'test_create.cxi')
with h5py.File(tmp_path / 'test_create.cxi','r') as f:
with h5py.File(tmp_path / 'test_create.cxi', 'r') as f:
assert f['cxi_version'][()] == 160
assert 'entry_1' in f
def test_add_entry_info(tmp_path):
entry_info = {'experiment_identifier':'test of cxi file writing tools',
entry_info = {'experiment_identifier': 'test of cxi file writing tools',
'title': 'my cool experiment',
'start_time': datetime.datetime.now(),
'end_time': datetime.datetime.now()}
@@ -124,12 +117,11 @@ def test_add_entry_info(tmp_path):
with data.create_cxi(tmp_path / 'test_add_entry_info.cxi') as f:
data.add_entry_info(f, entry_info)
with h5py.File(tmp_path / 'test_add_entry_info.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_entry_info.cxi', 'r') as f:
read_entry_info = data.get_entry_info(f)
print(read_entry_info)
for key in entry_info:
if isinstance(entry_info[key], np.ndarray):
assert np.allclose(entry_info[key], read_entry_info[key])
@@ -138,18 +130,18 @@ def test_add_entry_info(tmp_path):
def test_add_sample_info(tmp_path):
sample_info = {'name':'A nice fake sample',
sample_info = {'name': 'A nice fake sample',
'concentration': 10,
'mass': 5.3,
'temperature': 76,
'description': 'A very nice sample',
'unit_cell': np.array([1,1,1,90.,90.,90.])}
'unit_cell': np.array([1, 1, 1, 90., 90., 90.])}
with data.create_cxi(tmp_path / 'test_add_sample_info.cxi') as f:
data.add_sample_info(f, sample_info)
with h5py.File(tmp_path / 'test_add_sample_info.cxi','r') as f:
read_sample_info = data.get_sample_info(f)
with h5py.File(tmp_path / 'test_add_sample_info.cxi', 'r') as f:
read_sample_info = data.get_sample_info(f)
for key in sample_info:
if isinstance(sample_info[key], np.ndarray):
@@ -158,7 +150,7 @@ def test_add_sample_info(tmp_path):
assert np.isclose(sample_info[key], read_sample_info[key])
else:
assert sample_info[key] == read_sample_info[key]
def test_add_source(tmp_path):
wavelength = 1e-9
@@ -167,26 +159,26 @@ def test_add_source(tmp_path):
with data.create_cxi(tmp_path / 'test_add_source.cxi') as f:
data.add_source(f, wavelength)
with h5py.File(tmp_path / 'test_add_source.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_source.cxi', 'r') as f:
# Check this directly since we want to make sure it saved
# the wavelength and energy
read_wavelength = f['entry_1/instrument_1/source_1/wavelength'][()]
read_energy = f['entry_1/instrument_1/source_1/energy'][()]
assert np.isclose( wavelength, read_wavelength)
assert np.isclose( energy, read_energy)
assert np.isclose(wavelength, read_wavelength)
assert np.isclose(energy, read_energy)
def test_add_detector(tmp_path):
distance = 0.34
basis = np.array([[0,-30e-6,0],
[-20e-6,0,0]]).astype(np.float32).transpose()
corner = np.array((2550e-6,3825e-6,0.3)).astype(np.float32)
basis = np.array([[0, -30e-6, 0],
[-20e-6, 0, 0]]).astype(np.float32).transpose()
corner = np.array((2550e-6, 3825e-6, 0.3)).astype(np.float32)
with data.create_cxi(tmp_path / 'test_add_detector.cxi') as f:
data.add_detector(f, distance, basis, corner=corner)
with h5py.File(tmp_path / 'test_add_detector.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_detector.cxi', 'r') as f:
# Check this directly since we want to make sure it saved
# the pixel sizes
d1 = f['entry_1/instrument_1/detector_1']
@@ -198,112 +190,110 @@ def test_add_detector(tmp_path):
assert np.isclose(distance, read_distance)
assert np.allclose(basis, read_basis)
assert np.isclose(np.linalg.norm(basis[:,1]), read_x_pix)
assert np.isclose(np.linalg.norm(basis[:,0]), read_y_pix)
assert np.allclose(corner,read_corner)
assert np.isclose(np.linalg.norm(basis[:, 1]), read_x_pix)
assert np.isclose(np.linalg.norm(basis[:, 0]), read_y_pix)
assert np.allclose(corner, read_corner)
def test_add_mask(tmp_path):
mask = (np.random.rand(350,600) > 0.1).astype(np.uint8)
mask = (np.random.rand(350, 600) > 0.1).astype(np.uint8)
with data.create_cxi(tmp_path / 'test_add_mask.cxi') as f:
data.add_mask(f, mask)
with h5py.File(tmp_path / 'test_add_mask.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_mask.cxi', 'r') as f:
read_mask = data.get_mask(f)
assert np.all(mask == read_mask)
def test_add_qe_mask(tmp_path):
qe_mask = np.random.rand(350,199).astype(np.float32)
qe_mask = np.random.rand(350, 199).astype(np.float32)
with data.create_cxi(tmp_path / 'test_add_qe_mask.cxi') as f:
data.add_qe_mask(f, qe_mask)
with h5py.File(tmp_path / 'test_add_qe_mask.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_qe_mask.cxi', 'r') as f:
read_qe_mask = data.get_qe_mask(f)
assert np.allclose(qe_mask, read_qe_mask)
def test_add_dark(tmp_path):
dark = np.random.rand(350,620)
dark = np.random.rand(350, 620)
with data.create_cxi(tmp_path / 'test_add_dark.cxi') as f:
data.add_dark(f, dark)
with h5py.File(tmp_path / 'test_add_dark.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_dark.cxi', 'r') as f:
read_dark = data.get_dark(f)
print(dark.shape)
assert np.allclose(dark, read_dark)
def test_add_data(tmp_path):
# First test from numpy, with axes
fake_data = np.random.rand(100,256,256)
axes = ['translation','y','x']
fake_data = np.random.rand(100, 256, 256)
axes = ['translation', 'y', 'x']
with data.create_cxi(tmp_path / 'test_add_data.cxi') as f:
data.add_data(f, fake_data, axes)
with h5py.File(tmp_path / 'test_add_data.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_data.cxi', 'r') as f:
# Check this directly since we want to make sure it saved
# it in all the places it should have
read_data_1 = f['entry_1/data_1/data'][()]
read_data_2 = f['entry_1/instrument_1/detector_1/data'][()]
read_axes = str(f['entry_1/instrument_1/detector_1/data'].attrs['axes'].decode())
assert np.allclose(fake_data, read_data_1)
assert np.allclose(fake_data, read_data_2)
assert 'translation:y:x' == read_axes
# Then test from torch, without axes
fake_data = t.from_numpy(fake_data)
with data.create_cxi(tmp_path / 'test_add_data_torch.cxi') as f:
data.add_data(f, fake_data)
with h5py.File(tmp_path / 'test_add_data_torch.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_data_torch.cxi', 'r') as f:
read_data, axes = data.get_data(f)
assert np.allclose(fake_data.numpy(),read_data)
assert np.allclose(fake_data.numpy(), read_data)
def test_add_shot_to_shot_info(tmp_path):
analyzer = np.random.rand(100)
with data.create_cxi(tmp_path / 'test_add_shot_to_shot_info.cxi') as f:
data.add_shot_to_shot_info(f, analyzer, 'analyzer_angle')
with h5py.File(tmp_path / 'test_add_shot_to_shot_info.cxi') as f:
# Check this directly since we want to make sure it saved
# it in all the places it should have
read_analyzer_1 = f['entry_1/data_1/analyzer_angle'][()]
read_analyzer_2 = \
f['entry_1/instrument_1/detector_1/analyzer_angle'][()]
read_analyzer_2 = f['entry_1/instrument_1/detector_1/analyzer_angle'][()]
read_analyzer_3 = f['entry_1/sample_1/geometry_1/analyzer_angle'][()]
assert np.allclose(analyzer, read_analyzer_1)
assert np.allclose(analyzer, read_analyzer_2)
assert np.allclose(analyzer, read_analyzer_3)
def test_add_ptycho_translations(tmp_path):
translations = np.random.rand(3,100)
translations = np.random.rand(3, 100)
with data.create_cxi(tmp_path / 'test_add_ptycho_translations.cxi') as f:
data.add_ptycho_translations(f, translations)
with h5py.File(tmp_path / 'test_add_ptycho_translations.cxi','r') as f:
with h5py.File(tmp_path / 'test_add_ptycho_translations.cxi', 'r') as f:
# Check this directly since we want to make sure it saved
# it in all the places it should have
read_translations_1 = f['entry_1/data_1/translation'][()]
read_translations_2 = \
f['entry_1/instrument_1/detector_1/translation'][()]
read_translations_2 = f['entry_1/instrument_1/detector_1/translation'][()]
read_translations_3 = f['entry_1/sample_1/geometry_1/translation'][()]
assert np.allclose(-translations, read_translations_1)
@@ -312,8 +302,7 @@ def test_add_ptycho_translations(tmp_path):
def test_nested_dict_to_h5(tmp_path, example_nested_dicts):
### Tests both nested_dict_to_h5 and h5_to_nested_dict
# Tests both nested_dict_to_h5 and h5_to_nested_dict
def check_dict_equality(truth, to_test):
for key in truth.keys():
if isinstance(truth[key], dict):
@@ -328,23 +317,24 @@ def test_nested_dict_to_h5(tmp_path, example_nested_dicts):
assert truth[key] == to_test[key]
else:
assert 0
for test_dict in example_nested_dicts:
filename = tmp_path / 'example_dataset.h5'
data.nested_dict_to_h5(filename, test_dict)
roundtrip = data.h5_to_nested_dict(filename)
check_dict_equality(test_dict, roundtrip)
def test_h5_to_nested_dict(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
# Just test that it runs without errors for these ones.
# A round-trip test is in test_nested_dict_to_h5
d = data.h5_to_nested_dict(cxi)
data.h5_to_nested_dict(cxi)
def test_nested_dict_to_numpy(example_nested_dicts):
def check_dict_numpyness(truth, to_test):
def check_dict_numpyness(truth, to_test):
for key in truth.keys():
if isinstance(truth[key], dict):
check_dict_numpyness(truth[key], to_test[key])
@@ -361,12 +351,11 @@ def test_nested_dict_to_numpy(example_nested_dicts):
for test_dict in example_nested_dicts:
numpy_dict = data.nested_dict_to_numpy(test_dict)
check_dict_numpyness(test_dict, numpy_dict)
def test_nested_dict_to_torch(example_nested_dicts):
check_dict_numpyness(test_dict, numpy_dict)
def check_dict_torchiness(truth, to_test):
def test_nested_dict_to_torch(example_nested_dicts):
def check_dict_torchiness(truth, to_test):
for key in truth.keys():
if isinstance(truth[key], dict):
check_dict_torchiness(truth[key], to_test[key])
tests/tools/test_image_processing.py
+57 -58
View File
@@ -1,18 +1,19 @@
import numpy as np
import torch as t
from scipy import ndimage
from cdtools.tools import image_processing, interactions
from scipy import ndimage
def test_centroid():
# Test single im
im = t.rand((30,40))
im = t.rand((30, 40))
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))
ims = t.rand((5, 30, 40))
sp_centroids = [ndimage.center_of_mass(im.numpy())
for im in ims]
centroids = image_processing.centroid(ims)
@@ -21,33 +22,33 @@ def test_centroid():
def test_centroid_sq():
# Test single im
im = t.rand((30,40))
im = t.rand((30, 40))
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))
ims = t.rand((5, 30, 40)) + 1j * t.rand((5, 30, 40))
np_ims = ims.numpy()
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)))
def test_sinc_subpixel_shift():
im = np.zeros((512,512), dtype=np.complex128)
im[256,256] = 1
im = np.zeros((512, 512), dtype=np.complex128)
im[256, 256] = 1
# test it by creating a single pixel object and seeing that it is
# shifted correctly
xs = np.arange(512) - 256
Ys,Xs = np.meshgrid(xs,xs)
sinc_im = np.sinc(Xs-0.3) * np.sinc(Ys-0.6)
Ys, Xs = np.meshgrid(xs, xs)
sinc_im = np.sinc(Xs - 0.3) * np.sinc(Ys - 0.6)
torch_im = t.as_tensor(im)
test_im = image_processing.sinc_subpixel_shift(torch_im,(0.3,0.6))
test_im = image_processing.sinc_subpixel_shift(torch_im, (0.3, 0.6))
# The fidelity isn't great due to the FFT-based approach, so we need
# a pretty relaxed condition
@@ -57,84 +58,82 @@ def test_sinc_subpixel_shift():
def test_find_pixel_shift():
# Test two real ims
big_im = t.rand((30,70))
im1 = big_im[3:,:-20]
im2 = big_im[:-3,20:]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([-3,20]))
big_im = t.rand((30, 70))
im1 = big_im[3:, :-20]
im2 = big_im[:-3, 20:]
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([-3, 20]))
# Test a real and complex im
big_im = t.rand((30,70))
im1 = big_im[:-5,10:].to(dtype=t.complex64)
im2 = big_im[5:,:-10]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,-10]))
assert t.all(image_processing.find_pixel_shift(im2,im1) == t.LongTensor([-5,10]))
big_im = t.rand((30, 70))
im1 = big_im[:-5, 10:].to(dtype=t.complex64)
im2 = big_im[5:, :-10]
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([5, -10]))
assert t.all(image_processing.find_pixel_shift(im2, im1) == t.LongTensor([-5, 10]))
# Test two complex ims
big_im = t.rand((45,45)) + 1j * t.rand((45,45))
im1 = big_im[:-5,:-4]
im2 = big_im[5:,4:]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,4]))
big_im = t.rand((45, 45)) + 1j * t.rand((45, 45))
im1 = big_im[:-5, :-4]
im2 = big_im[5:, 4:]
assert t.all(image_processing.find_pixel_shift(im1, im2) == t.LongTensor([5, 4]))
def test_find_subpixel_shift():
# We can do this by creating a test probe and a test object
test_probe = t.rand((70,70)) + 1j * t.rand((70,70))
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
test_probe = t.rand((70, 70)) + 1j * t.rand((70, 70))
test_obj = t.ones((300, 300)) + 1j * t.rand((300, 300))
shift = t.tensor((0.8, 0.75))
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift, multiple_modes=False)
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0,0), resolution=50)
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0, 0), resolution=50)
# tolerance of 0.03 on this measurement
assert t.all(t.abs(shift - retrieved_shift) < 0.03)
def test_find_shift():
# We can do this by creating a test probe and a test object
test_probe = t.rand((200,200)) + 1j * t.rand((200,200))
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
test_probe = t.rand((200, 200)) + 1j * t.rand((200, 200))
test_obj = t.ones((300, 300)) + 1j * t.rand((300, 300))
shift = t.tensor((0.8, 0.75))
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift,
multiple_modes=False)[:-40,:-6]
multiple_modes=False)[:-40, :-6]
retrieved_shift = image_processing.find_shift(im, test_probe[40:,6:], resolution=50)
retrieved_shift = image_processing.find_shift(im, test_probe[40:, 6:], resolution=50)
# tolerance of 0.03 on this measurement
assert t.all(t.abs(shift + t.Tensor((40,6)) - retrieved_shift) < 0.03)
assert t.all(t.abs(shift + t.Tensor((40, 6)) - retrieved_shift) < 0.03)
def test_convolve_1d():
test_image = np.random.rand(400,300)
#test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
xs = np.linspace(-100,100,300)
kernel = 1/(1+xs**2)
test_image = np.random.rand(400, 300)
# test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
xs = np.linspace(-100, 100, 300)
kernel = 1 / (1 + xs**2)
# First, we test with everything real, dim=1
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
t.as_tensor(kernel),dim=1)
t.as_tensor(kernel), dim=1)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
assert np.allclose(convolved.numpy(),np_result)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image, axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
assert np.allclose(convolved.numpy(), np_result)
xs = np.linspace(-100,100,400)
kernel = 1/(1+xs**2)
xs = np.linspace(-100, 100, 400)
kernel = 1 / (1 + xs**2)
# Then with dim=0, and a non-fftshifted kernel
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
t.as_tensor(np.fft.ifftshift(kernel)),
fftshift_kernel=False)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0))
assert np.allclose(convolved.numpy(),np_result)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image, axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:, None], axis=0))
assert np.allclose(convolved.numpy(), np_result)
# And finally with complex input
convolved = image_processing.convolve_1d(t.as_tensor(test_image,dtype=t.complex64),
t.as_tensor(kernel,dtype=t.complex64)).numpy()
np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)
assert np.allclose(convolved,np_result)
convolved = image_processing.convolve_1d(t.as_tensor(test_image, dtype=t.complex64),
t.as_tensor(kernel, dtype=t.complex64)).numpy()
np_result = np.fft.ifft(np.fft.fft(test_image, axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:, None], axis=0)
assert np.allclose(convolved, np_result)
tests/tools/test_initializers.py
+58 -66
View File
@@ -1,26 +1,27 @@
from cdtools.tools import initializers
from cdtools.datasets import Ptycho2DDataset
import numpy as np
import torch as t
from cdtools.tools import initializers
from cdtools.datasets import Ptycho2DDataset
def test_exit_wave_geometry():
# First test a simple case where nothing need change
basis = t.Tensor([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(0,1)
shape = t.Size([73,56])
basis = t.Tensor([[0, -30e-6, 0],
[-20e-6, 0, 0]]).transpose(0, 1)
shape = t.Size([73, 56])
wavelength = 1e-9
distance = 1.
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
assert t.allclose(rs_basis[0,1],t.Tensor([-8.928571428571428e-07]))
assert t.allclose(rs_basis[1,0],t.Tensor([-4.5662100456621004e-07]))
assert t.allclose(rs_basis[0, 1], t.Tensor([-8.928571428571428e-07]))
assert t.allclose(rs_basis[1, 0], t.Tensor([-4.5662100456621004e-07]))
def test_calc_object_setup():
# First just try a simple case
probe_shape = t.Size([120,57])
translations = t.rand((30,2)) * 300
probe_shape = t.Size([120, 57])
translations = t.rand((30, 2)) * 300
t_max = t.max(translations, dim=0)[0]
t_min = t.min(translations, dim=0)[0]
obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations)
@@ -28,65 +29,58 @@ def test_calc_object_setup():
assert t.allclose(min_translation, t_min)
assert obj_shape == t.Size(exp_shape)
# Then add some padding
padding = 5
obj_shape, min_translation = initializers.calc_object_setup(probe_shape, translations, padding=padding)
assert t.allclose(min_translation, t_min - padding)
assert obj_shape == t.Size(exp_shape + 2 * padding)
def test_gaussian():
# Generate gaussian as a numpy array (square array)
shape = [10, 10]
sigma = [2.5, 2.5]
center = ((shape[0]-1)/2, (shape[1]-1)/2)
center = ((shape[0] - 1) / 2, (shape[1] - 1) / 2)
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
np_result = 10 * np.exp(-0.5 * ((x - center[1]) / sigma[1])**2 - 0.5 * ((y - center[0]) / sigma[0])**2)
init_result = initializers.gaussian(shape, sigma, amplitude=10).numpy()
assert np.allclose(init_result, np_result)
# Generate gaussian as a numpy array (rectangular array)
shape = [10, 5]
sigma = [2.5, 3]
center = ((shape[0]-1)/2, (shape[1]-1)/2)
center = ((shape[0] - 1) / 2, (shape[1] - 1) / 2)
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
np_result = np.exp(-0.5 * ((x - center[1]) / sigma[1])**2
- 0.5 * ((y - center[0]) / sigma[0])**2)
init_result = initializers.gaussian(shape, sigma).numpy()
assert np.allclose(init_result, np_result)
# Generate gaussian with curvature
shape = [20, 30]
sigma = [2.5, 5]
curvature = [1,0.6]
center = ((shape[0]-1)/2 + 3, (shape[1]-1)/2 - 1.4)
curvature = [1, 0.6]
center = ((shape[0] - 1) / 2 + 3, (shape[1] - 1) / 2 - 1.4)
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = (10+0j)*np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
np_result *= np.exp(0.5j*curvature[1]*(x-center[1])**2
+0.5j*curvature[0]*(y-center[0])**2)
init_result = initializers.gaussian(shape, sigma, center=center,
curvature=curvature, amplitude=10).numpy()
np_result = (10 + 0j) * np.exp(-0.5 * ((x - center[1]) / sigma[1])**2 - 0.5 * ((y - center[0]) / sigma[0])**2)
np_result *= np.exp(0.5j * curvature[1] * (x - center[1])**2 + 0.5j * curvature[0] * (y - center[0])**2)
init_result = initializers.gaussian(shape, sigma, center=center, curvature=curvature, amplitude=10).numpy()
assert np.allclose(init_result, np_result)
def test_gaussian_probe(ptycho_cxi_1):
dataset = Ptycho2DDataset.from_cxi(ptycho_cxi_1[0])
det_basis = t.Tensor(dataset.detector_geometry['basis'])
det_shape = t.Size(dataset.patterns.shape[-2:])
wavelength = dataset.wavelength
distance = dataset.detector_geometry['distance']
basis = initializers.exit_wave_geometry(det_basis,
det_shape,
wavelength,
distance)
det_shape,
wavelength,
distance)
# Basis is around 60nm in the i(y) direction, 85nm in the j(x) direction
# Full window is therefore about 15 um in i(y) and 20 um in the j(x) dir
@@ -94,15 +88,13 @@ def test_gaussian_probe(ptycho_cxi_1):
sigma = 5e-7
# Build a stage explicitly with numpy to compare against
x = (np.arange(256) - 127.5) * (-basis[0,1]).numpy()
y = (np.arange(256) - 127.5) * (-basis[1,0]).numpy()
Xs,Ys = np.meshgrid(x,y)
Rs = np.sqrt(Xs**2+Ys**2)
x = (np.arange(256) - 127.5) * (-basis[0, 1]).numpy()
y = (np.arange(256) - 127.5) * (-basis[1, 0]).numpy()
Xs, Ys = np.meshgrid(x, y)
Rs = np.sqrt(Xs**2 + Ys**2)
# Now we first test the non-propagated probe
np_probe = np.exp(-1/(2*sigma**2) * Rs**2)
np_probe = np.exp(- 1 / (2 * sigma**2) * Rs**2)
normalization = 0
for params, im in dataset:
@@ -110,27 +102,26 @@ def test_gaussian_probe(ptycho_cxi_1):
normalization /= len(dataset)
normalization_1 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
probe = initializers.gaussian_probe(
dataset, basis, det_shape, sigma).numpy()
assert np.allclose(probe, normalization_1*np_probe)
assert np.allclose(probe, normalization_1 * np_probe)
# And then a propagated probe
z = 1e-4 #nm
z = 1e-4 # nm
k = 2 * np.pi / wavelength
w0 = np.sqrt(2)*sigma
w0 = np.sqrt(2) * sigma
zr = np.pi * w0**2 / wavelength
wz = w0 * np.sqrt(1 + (z / zr)**2)
Rz = z * (1 + (zr / z)**2)
np_probe = np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
Rz = z * (1 + (zr / z)**2)
np_probe = np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
normalization_2 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
normalization_2 = np.sqrt(normalization / np.sum(np.abs(np_probe)**2))
probe = initializers.gaussian_probe(dataset, basis, det_shape, sigma,
propagation_distance=z).numpy()
assert np.allclose(probe, normalization_2*np_probe)
assert np.allclose(probe, normalization_2 * np_probe)
def test_SHARP_style_probe(ptycho_cxi_1):
@@ -148,11 +139,13 @@ def test_SHARP_style_probe(ptycho_cxi_1):
wavelength,
distance)
assert basis.shape == t.Size([3, 2])
probe = initializers.SHARP_style_probe(dataset)
assert probe.shape == t.Size([256,256])
assert probe.shape == t.Size([256, 256])
probe = initializers.SHARP_style_probe(dataset, propagation_distance=20e-6)
assert probe.shape == t.Size([256,256])
assert probe.shape == t.Size([256, 256])
def test_RPI_spectral_init():
@@ -160,28 +153,27 @@ def test_RPI_spectral_init():
# since the original implementation is in numpy and there aren't any clear
# cases that can be calculated analytically.
pattern = np.random.rand(230,253).astype(np.float32)
probe = np.random.rand(230,253).astype(np.complex64)
obj_shape = [37,53]
pattern = np.random.rand(230, 253).astype(np.float32)
probe = np.random.rand(230, 253).astype(np.complex64)
obj_shape = [37, 53]
mask = t.Tensor(np.random.rand(*pattern.shape) > 0.04)
background = t.as_tensor(np.random.rand(*pattern.shape),dtype=t.float32) * 0.05
background = t.as_tensor(np.random.rand(*pattern.shape), dtype=t.float32) * 0.05
probe = t.as_tensor(probe)
pattern = t.as_tensor(pattern)
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape)
assert list(obj.shape) == [1]+obj_shape
assert list(obj.shape) == [1] + obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, mask=mask)
assert list(obj.shape) == [2]+obj_shape
assert list(obj.shape) == [2] + obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, background=background)
assert list(obj.shape) == [2]+obj_shape
assert list(obj.shape) == [2] + obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, mask=mask,
background=background)
assert list(obj.shape) == [2]+obj_shape
assert list(obj.shape) == [2] + obj_shape
tests/tools/test_interactions.py
+98 -114
View File
@@ -1,9 +1,10 @@
from cdtools.tools import interactions
import numpy as np
import torch as t
from numpy import fft
from numpy.fft import fftshift, ifftshift
from numpy import fft
import pytest
import torch as t
from cdtools.tools import interactions
# Have a random probe and a random object and test the two
@@ -14,59 +15,60 @@ import pytest
@pytest.fixture(scope='module')
def random_probe():
return np.random.rand(256,256) * np.exp(2j * np.pi * np.random.rand(256,256))
return np.random.rand(256, 256) * np.exp(2j * np.pi * np.random.rand(256, 256))
@pytest.fixture(scope='module')
def random_obj():
return np.random.rand(900,900) * np.exp(2j * np.pi * np.random.rand(900,900))
return np.random.rand(900, 900) * np.exp(2j * np.pi * np.random.rand(900, 900))
@pytest.fixture(scope='module')
def single_pixel_probe(scope='module'):
probe = np.zeros((256,256), dtype=np.complex128)
probe[128,128] = 1
probe = np.zeros((256, 256), dtype=np.complex128)
probe[128, 128] = 1
return probe
def test_translations_to_pixel():
# First, try the case where everything is ones and simple
basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
translations = t.rand((10,3))
basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t()
translations = t.rand((10, 3))
output = interactions.translations_to_pixel(basis, translations)
assert t.allclose(output, -translations[:,:2].flip(1))
assert t.allclose(output, -translations[:, :2].flip(1))
# Next, try a case with a single translation
translation = t.rand((3))
output = interactions.translations_to_pixel(basis, translation)
assert t.allclose(output, -translation[:2].flip(0))
# Then, try a case with no surface normal but with a real conversion
basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t()
translations = t.rand((10,3))
basis = t.Tensor([[0, -2, 0], [-1, 0, 0.1]]).t()
translations = t.rand((10, 3))
output = interactions.translations_to_pixel(basis, translations)
basis_vectors_inv = t.pinverse(basis)
translations[:,2] = 0 # manually project off z component
assert t.allclose(output, t.mm(translations,basis_vectors_inv.t()))
translations[:, 2] = 0 # manually project off z component
assert t.allclose(output, t.mm(translations, basis_vectors_inv.t()))
# Finally, try a case with a known surface normal (reflection)
basis = t.Tensor([[0,-1,0],[0,0,1]]).t()
surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)])
translations = t.rand((10,3))
basis = t.Tensor([[0, -1, 0], [0, 0, 1]]).t()
surface_normal = t.Tensor([np.sqrt(2), 0, -np.sqrt(2)])
translations = t.rand((10, 3))
output = interactions.translations_to_pixel(basis, translations,
surface_normal=surface_normal)
exp_translations = t.stack((-translations[:,1],translations[:,0]),dim=1)
exp_translations = t.stack((-translations[:, 1], translations[:, 0]), dim=1)
assert t.allclose(output, exp_translations)
def test_pixel_to_translations():
# First, try the case where everything is ones and simple
basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
translations = t.rand((10,3))
translations[:,2] = 0
basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t()
translations = t.rand((10, 3))
translations[:, 2] = 0
output = interactions.translations_to_pixel(basis, translations)
roundtrip = interactions.pixel_to_translations(basis, output)
assert t.allclose(translations, roundtrip)
# Next, try a case with a single translation
translation = t.rand((3))
translation[2] = 0
@@ -74,65 +76,59 @@ def test_pixel_to_translations():
roundtrip = interactions.pixel_to_translations(basis, output)
assert t.allclose(translation, roundtrip)
# Then, try a case with no surface normal but with a real conversion
basis = t.Tensor([[0,-2,0],[-1,0,0.1]]).t()
translations = t.rand((10,3))
translations[:,2] = 0 # manually project off z component
basis = t.Tensor([[0, -2, 0], [-1, 0, 0.1]]).t()
translations = t.rand((10, 3))
translations[:, 2] = 0 # manually project off z component
output = interactions.translations_to_pixel(basis, translations)
roundtrip = interactions.pixel_to_translations(basis, output)
assert t.allclose(translations, roundtrip)
# Finally, try a case with a known surface normal (reflection)
basis = t.Tensor([[0,-1,0],[0,0,1]]).t()
surface_normal = t.Tensor([np.sqrt(2),0,-np.sqrt(2)])
translations = t.rand((10,3))
translations[:,2] = 0 # manually project off z component
basis = t.Tensor([[0, -1, 0], [0, 0, 1]]).t()
surface_normal = t.Tensor([np.sqrt(2), 0, -np.sqrt(2)])
translations = t.rand((10, 3))
translations[:, 2] = 0 # manually project off z component
output = interactions.translations_to_pixel(basis, translations,
surface_normal=surface_normal)
roundtrip = interactions.pixel_to_translations(basis, output,
surface_normal=surface_normal)
surface_normal=surface_normal)
assert t.allclose(translations, roundtrip)
def test_project_translations_to_sample():
# First, try the case where everything is ones and simple
basis = t.Tensor([[0,-1,0],[-1,0,0]]).t()
translations = t.rand((10,3))
basis = t.Tensor([[0, -1, 0], [-1, 0, 0]]).t()
translations = t.rand((10, 3))
pixels, props = interactions.project_translations_to_sample(basis, translations)
assert np.allclose(pixels[:,0].numpy(),-translations[:,1])
assert np.allclose(pixels[:,1].numpy(),-translations[:,0])
assert np.allclose(props.numpy(),-translations[:,2:].numpy())
assert np.allclose(pixels[:, 0].numpy(), -translations[:, 1])
assert np.allclose(pixels[:, 1].numpy(), -translations[:, 0])
assert np.allclose(props.numpy(), -translations[:, 2:].numpy())
# Next, a simple tilt along one axis. This is a 45 degree rotation
# around the positive y-axis
# Thus, y-axis translations are unaffected, but x-axis translations
# induce a motion of 1/sqrt(2) in the j- pixel space, as well as
# creating a propagation (negative propagation for positive x)
basis = t.Tensor([[0,-1e-3,0],[-np.sqrt(2)*1e-3,0,np.sqrt(2)*1e-3]]).t()
translations = t.rand((10,3))
basis = t.Tensor([[0, -1e-3, 0], [-np.sqrt(2) * 1e-3, 0, np.sqrt(2) * 1e-3]]).t()
translations = t.rand((10, 3))
pixels, props = interactions.project_translations_to_sample(basis, translations)
print(props.numpy())
print(-translations[:,2:].numpy() - translations[:,:1].numpy())
assert np.allclose(pixels[:,0].numpy(),-translations[:,1]*1e3)
assert np.allclose(pixels[:,1].numpy(),-translations[:,0]*1e3/np.sqrt(2))
assert np.allclose(props.numpy(),-translations[:,2:].numpy() - translations[:,:1].numpy())
print(-translations[:, 2:].numpy() - translations[:, :1].numpy())
assert np.allclose(pixels[:, 0].numpy(), -translations[:, 1] * 1e3)
assert np.allclose(pixels[:, 1].numpy(), -translations[:, 0] * 1e3 / np.sqrt(2))
assert np.allclose(props.numpy(), -translations[:, 2:].numpy() - translations[:, :1].numpy())
# Finally, we check a non-orthogonal case
def test_ptycho_2D_round(random_probe, random_obj):
# Test a stack of images
translations = np.random.rand(10,2) * 500
exit_waves_np = [random_probe * \
random_obj[tr[0]:tr[0]+random_probe.shape[0],
tr[1]:tr[1]+random_probe.shape[1]] for
translations = np.random.rand(10, 2) * 500
exit_waves_np = [random_probe * random_obj[tr[0]:tr[0] + random_probe.shape[0],
tr[1]:tr[1] + random_probe.shape[1]] for
tr in np.round(translations).astype(int)]
exit_waves_t = interactions.ptycho_2D_round(t.as_tensor(random_probe),
t.as_tensor(random_obj),
@@ -146,14 +142,13 @@ def test_ptycho_2D_round(random_probe, random_obj):
assert np.allclose(exit_wave_t.numpy(), exit_waves_np[0])
def test_ptycho_2D_linear(single_pixel_probe, random_obj):
# For this one, I just want to check one translation, but
# I need to check both formats
translations = np.array([[46.7,53.2]])
translation = np.array([46.7,53.2])
translations = np.array([[46.7, 53.2]])
translation = np.array([46.7, 53.2])
exit_waves_probe = interactions.ptycho_2D_linear(
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
@@ -167,14 +162,9 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
shift_probe=True)
# Check that the outputs match
assert t.allclose(exit_waves_probe[0],exit_wave_probe)
assert t.allclose(exit_waves_probe[0], exit_wave_probe)
exit_waves_obj = interactions.ptycho_2D_linear(
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.tensor(translations),
shift_probe=False)
exit_waves_obj = interactions.ptycho_2D_linear(t.as_tensor(single_pixel_probe), t.as_tensor(random_obj), t.tensor(translations), shift_probe=False)
exit_wave_obj = interactions.ptycho_2D_linear(
t.as_tensor(single_pixel_probe),
@@ -183,38 +173,37 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
shift_probe=False)
# Check that the outputs match
assert t.allclose(exit_waves_obj[0],exit_wave_obj)
assert t.allclose(exit_waves_obj[0], exit_wave_obj)
# For the shifted probe, we should find 4 pixels with intensity
exit_waves_probe = t.as_tensor(exit_waves_probe)[0]
probe_shift = np.array([[0.3*0.8,0.3*0.2],
[0.7*0.8,0.7*0.2]])
obj_section = random_obj[128+46:128+48,
128+53:128+55]
exit_section = exit_waves_probe[128:130,128:130]
probe_shift = np.array([[0.3 * 0.8, 0.3 * 0.2],
[0.7 * 0.8, 0.7 * 0.2]])
obj_section = random_obj[128 + 46:128 + 48,
128 + 53:128 + 55]
exit_section = exit_waves_probe[128:130, 128:130]
assert np.allclose(probe_shift * obj_section, exit_section)
# For the shifted obj, we should find one pixel with intensity
exit_waves_obj = t.as_tensor(exit_waves_obj)[0]
obj_shift = np.array([[0.3*0.8,0.3*0.2],
[0.7*0.8,0.7*0.2]])
obj_section = random_obj[128+46:128+48,
128+53:128+55]
exit_pixel = exit_waves_obj[128,128]
assert np.isclose(np.sum(obj_shift * obj_section),exit_pixel)
obj_shift = np.array([[0.3 * 0.8, 0.3 * 0.2],
[0.7 * 0.8, 0.7 * 0.2]])
obj_section = random_obj[128 + 46:128 + 48,
128 + 53:128 + 55]
exit_pixel = exit_waves_obj[128, 128]
assert np.isclose(np.sum(obj_shift * obj_section), exit_pixel)
# Test for a single translation
def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
# For this one, I just want to check one translation, but
# I need to check both formats
translations = np.array([[46.7,53.2]])
translation = np.array([46.7,53.2])
translations = np.array([[46.7, 53.2]])
translation = np.array([46.7, 53.2])
exit_waves_probe = interactions.ptycho_2D_sinc(
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
@@ -228,32 +217,31 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
shift_probe=True)
# Check that the outputs match
assert t.allclose(exit_waves_probe[0],exit_wave_probe)
assert t.allclose(exit_waves_probe[0], exit_wave_probe)
# Now we explicitly define what the sinc interpolated array should
# look like
xs = np.arange(256) - 128
Ys,Xs = np.meshgrid(xs,xs)
Ys, Xs = np.meshgrid(xs, xs)
sinc_probe = np.sinc(Xs) * np.sinc(Ys)
# Just check that the unshifted probe is correct
assert np.allclose(single_pixel_probe, sinc_probe)
sinc_shifted_probe = np.sinc(Xs-0.7) * np.sinc(Ys-0.2)
obj_section = random_obj[46:46+256,
53:53+256]
sinc_shifted_probe = np.sinc(Xs - 0.7) * np.sinc(Ys - 0.2)
obj_section = random_obj[46:46 + 256,
53:53 + 256]
exit_wave_np = sinc_shifted_probe * obj_section
exit_wave_torch = exit_wave_probe.numpy()
# The fidelity isn't great due to the FFT-based approach, so we need
# a pretty relaxed condition
assert np.max(np.abs(exit_wave_np-exit_wave_torch)) < 0.005
assert np.max(np.abs(exit_wave_np - exit_wave_torch)) < 0.005
def test_RPI_interaction(random_probe, random_obj):
random_obj1 = random_obj[:79,:68] * 0 + 1
random_obj1 = random_obj[:79, :68] * 0 + 1
random_probe1 = random_probe * 0 + 1
t_random_obj1 = t.as_tensor(random_obj1)
t_random_probe1 = t.as_tensor(random_probe1)
@@ -261,36 +249,32 @@ def test_RPI_interaction(random_probe, random_obj):
obj1_fourier = fftshift(fft.fft2(ifftshift(random_obj1), norm='ortho'))
obj1_ups = np.zeros(random_probe1.shape[:2]).astype(np.complex128)
obj1_ups[random_probe1.shape[0]//2 - 79//2:
-(random_probe1.shape[0]-79 - (random_probe1.shape[0]//2 - 79//2)),
(random_probe1.shape[1]-68)//2:
(random_probe1.shape[1]-68)//2 + 68] = obj1_fourier
obj1_ups[random_probe1.shape[0] // 2 - 79 // 2:
-(random_probe1.shape[0] - 79 - (random_probe1.shape[0] // 2 - 79 // 2)),
(random_probe1.shape[1] - 68) // 2:
(random_probe1.shape[1] - 68) // 2 + 68] = obj1_fourier
output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups),
norm='ortho'))
norm='ortho'))
output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1]
/ (random_obj1.shape[-2] * random_obj1.shape[-1]))
output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1] / (random_obj1.shape[-2] * random_obj1.shape[-1]))
assert np.allclose(t_output1, output1)
random_obj2 = np.stack([random_obj[:64,:89]]*3)
random_probe2 = random_probe[3:,5:]
random_obj2 = np.stack([random_obj[:64, :89]] * 3)
random_probe2 = random_probe[3:, 5:]
t_random_obj2 = t.as_tensor(random_obj2)
t_random_probe2 = t.as_tensor(random_probe2)
t_output2 = interactions.RPI_interaction(t_random_probe2, t_random_obj2)
obj2_fourier = fftshift(fft.fft2(ifftshift(random_obj2), norm='ortho'))
obj2_ups = np.zeros((3,)+random_probe2.shape[:2]).astype(np.complex128)
obj2_ups[:,(random_probe2.shape[0]-64)//2:
(random_probe2.shape[0]-64)//2 + 64,
(random_probe2.shape[1]-89)//2:
(random_probe2.shape[1]-89)//2 + 89] = obj2_fourier
obj2_ups = np.zeros((3,) + random_probe2.shape[:2]).astype(np.complex128)
obj2_ups[:, (random_probe2.shape[0] - 64) // 2:
(random_probe2.shape[0] - 64) // 2 + 64,
(random_probe2.shape[1] - 89) // 2:
(random_probe2.shape[1] - 89) // 2 + 89] = obj2_fourier
output2 = random_probe2 * fftshift(fft.ifft2(ifftshift(obj2_ups),
norm='ortho'))
norm='ortho'))
output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1]
/ (random_obj2.shape[-2] * random_obj2.shape[-1]))
output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1] / (random_obj2.shape[-2] * random_obj2.shape[-1]))
assert np.allclose(t_output2, output2)
tests/tools/test_losses.py
+32 -38
View File
@@ -1,79 +1,73 @@
from cdtools.tools import losses
import numpy as np
import torch as t
from cdtools.tools import losses
# The idea here is to use a simple numpy calculation of the various
# objective functions to check the torch implementations and make sure
# that any optimizations in the future don't change the results
def test_amplitude_mse():
# Make some fake data
data = np.random.rand(10,100,100)
data = np.random.rand(10, 100, 100)
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
sim = data + 0.1 * np.random.rand(10, 100, 100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(bool)
mask = (np.random.rand(100, 100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2)
#np_result /= data.size
torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
# np_result /= data.size
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2)
#np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim),
mask = t.from_numpy(mask))
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
# np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
def test_intensity_mse():
# Make some fake data
data = np.random.rand(10,100,100)
data = np.random.rand(10, 100, 100)
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
sim = data + 0.1 * np.random.rand(10, 100, 100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(bool)
mask = (np.random.rand(100, 100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum((data - sim)**2)
np_result /= data.size
torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
np_result /= data.size
torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (data - sim)**2)
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim),
mask = t.from_numpy(mask))
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
def test_poisson_nll():
# Make some fake data
data = np.random.rand(10,100,100)
data = np.random.rand(10, 100, 100)
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
sim = data + 0.1 * np.random.rand(10, 100, 100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(bool)
mask = (np.random.rand(100, 100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum(sim - data * np.log(sim))
np_result /= data.size
torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
np_result /= data.size
torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (sim - data * np.log(sim)))
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim),
mask = t.from_numpy(mask), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(),0))
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim),
mask=t.from_numpy(mask), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
tests/tools/test_measurements.py
+37 -46
View File
@@ -1,99 +1,90 @@
from cdtools.tools import measurements
import torch as t
import numpy as np
from cdtools.tools import measurements
def test_intensity():
wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
epsilon=1e-6
wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10))
epsilon = 1e-6
np_result = np.abs(wavefields.numpy())**2 + epsilon
assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon),
assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon),
t.as_tensor(np_result))
# Test single field case
assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon),
assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon),
t.as_tensor(np_result[0]))
det_slice = np.s_[3:,5:8]
assert t.allclose(measurements.intensity(wavefields,det_slice,epsilon=epsilon),
t.as_tensor(np_result[(np.s_[:],)+det_slice]))
det_slice = np.s_[3:, 5:8]
assert t.allclose(measurements.intensity(wavefields, det_slice, epsilon=epsilon),
t.as_tensor(np_result[(np.s_[:],) + det_slice]))
# Test single field case
assert t.allclose(measurements.intensity(wavefields[0],det_slice,epsilon=epsilon),
assert t.allclose(measurements.intensity(wavefields[0], det_slice, epsilon=epsilon),
t.as_tensor(np_result[0][det_slice]))
# With oversampling on
np_oversampling_result = (np_result[:,::2,::2] + \
np_result[:,1::2,::2] + \
np_result[:,::2,1::2] + \
np_result[:,1::2,1::2]) / 4
np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4
# With multiple fields
assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon, oversampling=2),
assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result,))
# With a single field
assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon, oversampling=2),
assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result[0],))
def test_incoherent_sum():
# With no explicit slice given
wavefields = t.rand((5,4,10,10)) + 1j * t.rand((5,4,10,10))
epsilon=1e-6
np_result = np.sum(np.abs(wavefields.numpy())**2,axis=-3) + epsilon
assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon),
wavefields = t.rand((5, 4, 10, 10)) + 1j * t.rand((5, 4, 10, 10))
epsilon = 1e-6
np_result = np.sum(np.abs(wavefields.numpy())**2, axis=-3) + epsilon
assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon),
t.as_tensor(np_result))
# Test single field case
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon),
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon),
t.as_tensor(np_result[0]))
# With a slice given
det_slice = np.s_[3:,5:8]
assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon),
t.as_tensor(np_result[(np.s_[:],)+det_slice]))
det_slice = np.s_[3:, 5:8]
assert t.allclose(measurements.incoherent_sum(wavefields, det_slice, epsilon=epsilon),
t.as_tensor(np_result[(np.s_[:],) + det_slice]))
# Test single field case
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],det_slice,epsilon=epsilon),
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], det_slice, epsilon=epsilon),
t.as_tensor(np_result[0][det_slice]))
# With oversampling on
np_oversampling_result = (np_result[:,::2,::2] + \
np_result[:,1::2,::2] + \
np_result[:,::2,1::2] + \
np_result[:,1::2,1::2]) / 4
np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4
# With multiple fields
assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon, oversampling=2),
assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result,))
# With a single field
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon, oversampling=2),
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result[0],))
def test_quadratic_background():
# test with intensity
wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
epsilon=1e-6
background = t.rand((10,10))
wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10))
epsilon = 1e-6
background = t.rand((10, 10))
np_result = np.abs(wavefields.numpy())**2 + background.numpy()**2 + epsilon
det_slice = np.s_[3:,5:8]
det_slice = np.s_[3:, 5:8]
result = measurements.quadratic_background(wavefields,background[det_slice],
result = measurements.quadratic_background(wavefields, background[det_slice],
detector_slice=det_slice,
epsilon=epsilon,
measurement=measurements.intensity)
assert t.allclose(result, t.tensor(np_result[(np.s_[:],)+det_slice]))
assert t.allclose(result, t.tensor(np_result[(np.s_[:],) + det_slice]))
# test with incoherent sum but no slice and no stack
wavefields = t.rand((4,10,10)) + 1j * t.rand((4,10,10))
np_result = np.sum(np.abs(wavefields.numpy())**2,axis=0)
wavefields = t.rand((4, 10, 10)) + 1j * t.rand((4, 10, 10))
np_result = np.sum(np.abs(wavefields.numpy())**2, axis=0)
np_result += background.numpy()**2
result = measurements.quadratic_background(wavefields, background,
epsilon=epsilon,
tests/tools/test_plotting.py
+15 -12
View File
@@ -1,47 +1,50 @@
from cdtools.tools import plotting
from cdtools.tools import initializers
import numpy as np
import torch as t
import scipy.datasets
import matplotlib.pyplot as plt
from cdtools.tools import plotting
from cdtools.tools import initializers
def test_plot_amplitude(show_plot):
# Test with tensor
im = t.as_tensor(scipy.datasets.ascent(),dtype=t.complex128)
plotting.plot_amplitude(im, basis = np.array([[0,-1], [-1,0], [0,0]]), title = 'Test Amplitude')
im = t.as_tensor(scipy.datasets.ascent(), dtype=t.complex128)
plotting.plot_amplitude(im, basis=np.array([[0, -1], [-1, 0], [0, 0]]), title='Test Amplitude')
if show_plot:
plt.show()
# Test with numpy array
im = scipy.datasets.ascent().astype(np.complex128)
plotting.plot_amplitude(im, title = 'Test Amplitude')
plotting.plot_amplitude(im, title='Test Amplitude')
if show_plot:
plt.show()
def test_plot_phase(show_plot):
# Test with tensor
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
plotting.plot_phase(im, title = 'Test Phase')
im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
plotting.plot_phase(im, title='Test Phase')
if show_plot:
plt.show()
# Test with numpy array
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]).numpy()
plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[0,-1], [-1,0], [0,0]]))
im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1]).numpy()
plotting.plot_phase(im, title='Test Phase', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
if show_plot:
plt.show()
def test_plot_colorized(show_plot):
# Test with tensor
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
gaussian = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
im = gaussian * t.as_tensor(scipy.datasets.ascent(), dtype=t.complex64)
plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[0,-1], [-1,0], [0,0]]))
plotting.plot_colorized(im, title='Test Colorize', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
if show_plot:
plt.show()
# Test with numpy array
im = im.numpy()
plotting.plot_colorized(im, title = 'Test Colorize')
plotting.plot_colorized(im, title='Test Colorize')
if show_plot:
plt.show()
tests/tools/test_propagators.py
+194 -211
View File
@@ -1,7 +1,3 @@
from cdtools.tools import initializers
from cdtools.tools import propagators
from cdtools.tools import image_processing
import numpy as np
import torch as t
import pytest
@@ -9,35 +5,37 @@ import scipy.datasets
from scipy import stats
from matplotlib import pyplot as plt
from cdtools.tools import initializers
from cdtools.tools import propagators
from cdtools.tools import image_processing
@pytest.fixture(scope='module')
def exit_waves_1():
# Import scipy test image and add a random phase
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 = 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)
# Construct wavefront from image
probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3)
return probe * obj
def test_far_field(exit_waves_1):
# Far field diffraction patterns calculated by numpy with zero frequency in center
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho'))
assert(np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()), norm='ortho'))
assert (np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
def test_inverse_far_field(exit_waves_1):
# We want the inverse far field to map back to the exit waves with no intensity corrections
# Far field result for exit waves calculated with numpy
far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho')))
far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()), norm='ortho')))
assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
assert (np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
def test_generate_high_NA_k_intensity_map():
@@ -45,9 +43,9 @@ def test_generate_high_NA_k_intensity_map():
# We need to generate a plausible scenario. I will start
# by using the initializer to generate a reasonable exit wave geometry
# and detector pair
basis = t.Tensor([[0,-30e-6,0],
[-20e-6,0,0]]).transpose(0,1)
shape = t.Size([478,573])
basis = t.Tensor([[0, -30e-6, 0],
[-20e-6, 0, 0]]).transpose(0, 1)
shape = t.Size([478, 573])
wavelength = 1e-9
distance = 1
rs_basis = \
@@ -60,7 +58,7 @@ def test_generate_high_NA_k_intensity_map():
# generate a good test exit wave
i = (np.arange(478) - 240)
j = (np.arange(573) - 270)
Is,Js = np.meshgrid(i,j,indexing='ij')
Is, Js = np.meshgrid(i, j, indexing='ij')
wavefield = ((np.abs(Is) < 20) * (np.abs(Js) < 25)).astype(np.complex128)
t_wavefield = t.as_tensor(wavefield, dtype=t.complex64)
@@ -74,37 +72,36 @@ def test_generate_high_NA_k_intensity_map():
# Checking first that for a low-NA propagation they give the same result
# 1e-4 tolerance seems to be reasonable in this comparison given my
# exploration with the code
#assert np.max(np.abs(high_NA-low_NA))/np.max(np.abs(low_NA)) < 1e-4
# assert np.max(np.abs(high_NA-low_NA))/np.max(np.abs(low_NA)) < 1e-4
# Now I will explore some results with a tilted sample
#print(rs_basis)
#print(rs_basis_tilted)
distance = 0.01#6e-3
rs_basis = \
initializers.exit_wave_geometry(basis, shape, wavelength, distance)
# print(rs_basis)
# print(rs_basis_tilted)
distance = 0.01 # 6e-3
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
rs_basis_tilted = rs_basis.clone()
rs_basis_tilted[2,1] = rs_basis_tilted[0,1]
rs_basis_tilted[2, 1] = rs_basis_tilted[0, 1]
k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
rs_basis_tilted, basis, shape, distance, wavelength,
dtype=t.float32)
high_NA_propagated = propagators.high_NA_far_field(
t_wavefield, k_map, intensity_map=intensity_map)
low_NA_propagated = propagators.far_field(t_wavefield)
low_NA = low_NA_propagated.numpy()
high_NA = high_NA_propagated.numpy()
#plt.close('all')
#plt.imshow(np.abs(low_NA))
#plt.figure()
#plt.imshow(np.abs(high_NA))
#plt.colorbar()
#plt.imshow(np.abs(wavefield))
#plt.show()
print(low_NA.shape, high_NA.shape)
# plt.close('all')
# plt.imshow(np.abs(low_NA))
# plt.figure()
# plt.imshow(np.abs(high_NA))
# plt.colorbar()
# plt.imshow(np.abs(wavefield))
# plt.show()
# Now I want to test that it doesn't crash for wavefields of various shapes
propagators.high_NA_far_field(t_wavefield.unsqueeze(0),
@@ -114,8 +111,9 @@ def test_generate_high_NA_k_intensity_map():
# I believe this works, but I still would like to get a second method for
# simulating at least one diffraction pattern as an independent check
#assert 0
# assert 0
def test_near_field_direction():
#
@@ -135,19 +133,19 @@ def test_near_field_direction():
x = (np.arange(901) - 400)
y = (np.arange(1200) - 500)
Ys,Xs = np.meshgrid(y,x)
Rs = np.sqrt(Xs**2+Ys**2)
E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)),dtype=t.complex64)
Ys, Xs = np.meshgrid(y, x)
Rs = np.sqrt(Xs**2 + Ys**2)
E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)), dtype=t.complex64)
E0_real = propagators.inverse_far_field(E0_fourier)
# This is in the top-left corner in Fourier space
wavelength = 3e-9 #nm
z = 1000e-9
wavelength = 3e-9 # nm
z = 1000e-9 # nm
asp = propagators.generate_angular_spectrum_propagator(
E0_real.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex64)
E0_real.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex64)
Ez_real = propagators.near_field(E0_real,asp)
Ez_real = propagators.near_field(E0_real, asp)
centroid = image_processing.centroid(t.abs(Ez_real))
@@ -155,35 +153,35 @@ def test_near_field_direction():
assert centroid[0] < Ez_real.shape[0] // 2
# Assert it's in left half
assert centroid[1] < Ez_real.shape[1] // 2
#plt.imshow(t.abs(E0_fourier))
#plt.figure()
#plt.imshow(t.abs(E0_real))
#plt.figure()
#plt.imshow(t.abs(Ez_real))
#plt.show()
# plt.imshow(t.abs(E0_fourier))
# plt.figure()
# plt.imshow(t.abs(E0_real))
# plt.figure()
# plt.imshow(t.abs(Ez_real))
# plt.show()
def test_near_field():
# The strategy is to compare the propagation of a gaussian beam to
# the propagation in the paraxial approximation.
x = (np.arange(901) - 450) * 1.5e-9
y = (np.arange(1200) - 600) * 1e-9
Ys,Xs = np.meshgrid(y,x)
Rs = np.sqrt(Xs**2+Ys**2)
Ys, Xs = np.meshgrid(y, x)
Rs = np.sqrt(Xs**2 + Ys**2)
wavelength = 3e-9 #nm
sigma = 20e-9 #nm
z = 1000e-9 #nm
wavelength = 3e-9 # nm
sigma = 20e-9 # nm
z = 1000e-9 # nm
k = 2 * np.pi / wavelength
w0 = np.sqrt(2)*sigma
w0 = np.sqrt(2) * sigma
zr = np.pi * w0**2 / wavelength
wz = w0 * np.sqrt(1 + (z / zr)**2)
Rz = z * (1 + (zr / z)**2)
Rz = z * (1 + (zr / z)**2)
E0 = np.exp(-Rs**2 / w0**2)
# The analytical expression for propagation of a gaussian beam in the
@@ -208,52 +206,49 @@ def test_near_field():
# If we choose e^(-ikx) to represent light propagating along K, the
# answer is no, and we find we have to use the inverse FT instead.
# Thus, e^(ikx) is the right choice here.
Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * ( z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr))
Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * (z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr))
Ez_nozphase = Ez * np.exp(-1j * k * z)
# First we check it normally
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
assert np.max(np.abs(Ez_nozphase-Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
assert np.max(np.abs(Ez_nozphase - Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
Emz = np.conj(Ez_nozphase)
Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy()
Emz_t = propagators.inverse_near_field(t.as_tensor(E0), asp).numpy()
# Again, 10^-3 is about all the accuracy we can expect
assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz))
assert np.max(np.abs(Emz - Emz_t)) < 1e-3 * np.max(np.abs(Emz))
# Then, we check that the bandlimiting at least does something
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,
E0.shape, (1.5e-9, 1e-9), wavelength, z,
dtype=t.complex128, bandlimit=0.3)
assert asp[140,0] == 0
assert asp[0,180] == 0
assert asp[130,0] != 0
assert asp[0,175] != 0
assert asp[140, 0] == 0
assert asp[0, 180] == 0
assert asp[130, 0] != 0
assert asp[0, 175] != 0
# Then, we check that automatic differentiation works
z = t.tensor([z],requires_grad=True)
spacing = t.tensor((1.5e-9,1e-9), requires_grad=True)
z = t.tensor([z], requires_grad=True)
spacing = t.tensor((1.5e-9, 1e-9), requires_grad=True)
wavelength = t.tensor([wavelength], requires_grad=True)
asp = propagators.generate_angular_spectrum_propagator(
E0.shape, spacing, wavelength, z)
t.real(asp[10,10]).backward()
t.real(asp[10, 10]).backward()
assert z.grad != 0
assert spacing.grad[0] != 0
assert wavelength.grad !=0
assert wavelength.grad != 0
def test_generalized_near_field():
@@ -264,26 +259,25 @@ def test_generalized_near_field():
# First, we should do a test with the phase ramp along the z direction
# explicitly included
basis= np.array([[0,-1.5e-9],[-1e-9,0],[0,0]])
i_vec,j_vec = np.arange(901) - 450 ,np.arange(1200)-600
Is, Js = np.meshgrid(i_vec,j_vec,indexing='ij')
Xs_0,Ys_0,Zs_0 = np.tensordot(basis,np.stack([Is,Js]),axes=1)
#x = (np.arange(901) - 450) * 1.5e-9
#y = (np.arange(1200) - 600) * 1e-9
#Xs_0,Ys_0 = np.meshgrid(x,y)
#Zs_0 = np.zeros(Xs_0.shape)
basis = np.array([[0, -1.5e-9], [-1e-9, 0], [0, 0]])
i_vec, j_vec = np.arange(901) - 450, np.arange(1200) - 600
Is, Js = np.meshgrid(i_vec, j_vec, indexing='ij')
Xs_0, Ys_0, Zs_0 = np.tensordot(basis, np.stack([Is, Js]), axes=1)
# x = (np.arange(901) - 450) * 1.5e-9
# y = (np.arange(1200) - 600) * 1e-9
# Xs_0,Ys_0 = np.meshgrid(x,y)
# Zs_0 = np.zeros(Xs_0.shape)
Positions = np.stack([Xs_0, Ys_0, Zs_0])
Positions = np.stack([Xs_0,Ys_0,Zs_0])
# assert 0
wavelength = 3e-9 #nm
sigma = 20e-9 #nm
z = 1000e-9 #nm
wavelength = 3e-9 # nm
sigma = 20e-9 # nm
z = 1000e-9 # nm
k = 2 * np.pi / wavelength
w0 = np.sqrt(2)*sigma
w0 = np.sqrt(2) * sigma
zr = np.pi * w0**2 / wavelength
# The analytical expression for propagation of a gaussian beam in the
# paraxial approx
@@ -292,140 +286,132 @@ def test_generalized_near_field():
def get_inv_R(Zs):
return Zs / (Zs**2 + zr**2)
def get_E(Xs, Ys, Zs, correct=True):
# Again, this follows the convention opposite from Wikipedia. See
# the note in test_angular_spectrum_propagator.
# if correct is True, remove the e^(ikz) dependence
Rs_sq = Xs**2 + Ys**2
Wzs = get_w(Zs)
E = w0 / Wzs * np.exp(-Rs_sq / Wzs**2) *\
np.exp(1j * k * ( Zs + Rs_sq * get_inv_R(Zs) / 2) + \
- 1j * np.arctan(Zs / zr))
np.exp(1j * k * (Zs + Rs_sq * get_inv_R(Zs) / 2) - 1j * np.arctan(Zs / zr))
# This removes the z-dependence of the phase
if correct:
if correct:
E = E * np.exp(-1j * k * Zs)
return E
def check_equiv(analytical, numerical):
phase = np.angle(np.mean(numerical.conj()*analytical))
comp = np.exp(1j*phase) * numerical
return (np.max(np.abs(analytical-comp))
< 1e-3 * np.max(np.abs(analytical)))
phase = np.angle(np.mean(numerical.conj() * analytical))
comp = np.exp(1j * phase) * numerical
return (np.max(np.abs(analytical - comp)) < 1e-3 * np.max(np.abs(analytical)))
# We make some rotation matrices to test
# This tests the straight ahead case
I = np.eye(3)
IdentityMatrix = np.eye(3)
# This tests a rotation about the y axis
th = np.deg2rad(5)
Ry = np.array([[np.cos(th),0,np.sin(th)],
[0,1,0],
[-np.sin(th),0,np.cos(th)]])
Ry = np.array([[np.cos(th), 0, np.sin(th)],
[0, 1, 0],
[-np.sin(th), 0, np.cos(th)]])
# This tests a rotation about two axes
phi = np.deg2rad(2)
Rx = np.array([[1,0,0],
[0,np.cos(phi),-np.sin(phi)],
[0,np.sin(phi),np.cos(phi)]])
Rboth = np.matmul(Rx,Ry)
Rx = np.array([[1, 0, 0],
[0, np.cos(phi), -np.sin(phi)],
[0, np.sin(phi), np.cos(phi)]])
Rboth = np.matmul(Rx, Ry)
# This tests a shearing
shear = 0.23
Rshear = np.array([[1,shear,0],
[0,1,0],
[0,0,1]])
Rshear = np.array([[1, shear, 0],
[0, 1, 0],
[0, 0, 1]])
# This tests an inversion of the axes
Rinv = np.array([[-1,0,0],
[0,-1,0],
[0,0,-1]])
Rinv = np.array([[-1, 0, 0],
[0, -1, 0],
[0, 0, -1]])
# This tests a reflection about the y-z plane
Rrefl = np.array([[-1,0,0],
[0,1,0],
[0,0,-1]])
# This tests a shearing and a rotation together
Rall = np.matmul(Rrefl,np.matmul(Rboth,Rshear))
Rrefl = np.array([[-1, 0, 0],
[0, 1, 0],
[0, 0, -1]])
# This tests a shearing and a rotation together
Rall = np.matmul(Rrefl, np.matmul(Rboth, Rshear))
# And we make some propagation vectors to test:
# This is along the z direction
z_dir = np.array([0,0,1])
z_dir = np.array([0, 0, 1])
# This checks that it's not sensitive to the magnitude
z_dir_large = np.array([0,0,10])
z_dir_large = np.array([0, 0, 10])
# And finally some offset vectors
# This checks straight ahead
z_offset = np.array([0,0,z])
z_offset = np.array([0, 0, z])
# This checks with an offset in x and y
shear_offset = np.array([0.1*z,-0.03*z,z])
shear_offset = np.array([0.1 * z, -0.03 * z, z])
# This checks with an offset in x and y, with negative z
shear_back_offset = np.array([0.1*z,-0.03*z,-z])
shear_back_offset = np.array([0.1 * z, -0.03 * z, -z])
rot_mats = [Rrefl,I,Rinv, Rboth, Rboth,Rboth, Rall, Rall, Rall, I, Rall]
offset_vecs = [z_offset]*8 + [shear_offset] + [shear_back_offset]*2
propagation_vecs = ['perp','offset',z_dir,
'perp','offset',z_dir_large,
'perp','offset',z_dir_large,
rot_mats = [Rrefl, IdentityMatrix, Rinv, Rboth, Rboth, Rboth, Rall, Rall, Rall, IdentityMatrix, Rall]
offset_vecs = [z_offset] * 8 + [shear_offset] + [shear_back_offset] * 2
propagation_vecs = ['perp', 'offset', z_dir,
'perp', 'offset', z_dir_large,
'perp', 'offset', z_dir_large,
z_dir, z_dir_large]
purposes = ['standard']*3 + ['both-rot']*3 + ['shear-rot']*3 + ['backward']*2
#rot_mats = [Ry.transpose()]
#offset = np.cross(np.dot(Ry.transpose(),basis)[:,0],
purposes = ['standard'] * 3 + ['both-rot'] * 3 + ['shear-rot'] * 3 + ['backward'] * 2
# rot_mats = [Ry.transpose()]
# offset = np.cross(np.dot(Ry.transpose(),basis)[:,0],
# np.dot(Ry.transpose(),basis)[:,1])
#offset /= np.linalg.norm(offset) / 3e-6
#offset_vecs = [-offset]#[shear_offset]
#propagation_vecs = [z_dir]
#purposes=['meh']
for purpose,rot_mat,offset_vec, propagation_vec \
in zip(purposes,rot_mats,offset_vecs,propagation_vecs):
# offset /= np.linalg.norm(offset) / 3e-6
# offset_vecs = [-offset]#[shear_offset]
# propagation_vecs = [z_dir]
# purposes=['meh']
for purpose, rot_mat, offset_vec, propagation_vec in zip(purposes, rot_mats, offset_vecs, propagation_vecs):
print('Testing', purpose)
Xs,Ys,Zs_0 = np.tensordot(rot_mat,Positions,axes=1)
Xs, Ys, Zs_0 = np.tensordot(rot_mat, Positions, axes=1)
new_basis = np.dot(rot_mat, basis)
Xs_prop, Ys_prop, Zs_prop = np.stack([Xs,Ys,Zs_0]) \
+ offset_vec[:,None,None]
Xs_prop, Ys_prop, Zs_prop = np.stack([Xs, Ys, Zs_0]) \
+ offset_vec[:, None, None]
print('Propagate Along', propagation_vec)
print('Propagate Along',propagation_vec)
if str(propagation_vec) == 'perp':
E0 = get_E(Xs,Ys,Zs_0, correct=False)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=False)
E0 = get_E(Xs, Ys, Zs_0, correct=False)
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=False)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,dtype=t.complex128)
E0.shape, new_basis, wavelength, offset_vec, dtype=t.complex128)
elif str(propagation_vec) == 'offset':
E0 = get_E(Xs,Ys,Zs_0, correct=True)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
E0 = get_E(Xs, Ys, Zs_0, correct=True)
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,
E0.shape, new_basis, wavelength, offset_vec,
dtype=t.complex128, propagate_along_offset=True)
else:
E0 = get_E(Xs,Ys,Zs_0, correct=True)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
E0 = get_E(Xs, Ys, Zs_0, correct=True)
Ez = get_E(Xs_prop, Ys_prop, Zs_prop, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,
dtype=t.complex128, propagation_vector=propagation_vec)
E0.shape, new_basis, wavelength, offset_vec, dtype=t.complex128,
propagation_vector=propagation_vec)
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
if not check_equiv(Ez, Ez_t):
#if True:
# if True:
plt.close('all')
plt.figure()
plt.imshow(np.angle(E0))
@@ -443,70 +429,67 @@ def test_generalized_near_field():
plt.imshow(np.angle(Ez_t))
plt.title('Angle of numerically calculated Ez')
plt.figure()
plt.imshow(np.abs(Ez-Ez_t))#/np.max(np.abs(Ez)))
plt.imshow(np.abs(Ez - Ez_t))
plt.title('Magnitude of difference')
plt.show()
assert check_equiv(Ez, Ez_t)
Em0_t = propagators.inverse_near_field(t.as_tensor(Ez),asp).numpy()
assert check_equiv(E0,Em0_t)
Em0_t = propagators.inverse_near_field(t.as_tensor(Ez), asp).numpy()
assert check_equiv(E0, Em0_t)
print('Test Successful')
# One final test, to see if any of a few arbitrary rotations will
# change the predicted propagation if everything else is kept
# constant
Rrands = [stats.ortho_group.rvs(3) for i in range(3)]
Xs,Ys,Zs_0 = np.tensordot(Rboth,Positions,axes=1)
Xs, Ys, Zs_0 = np.tensordot(Rboth, Positions, axes=1)
new_basis = np.dot(rot_mat, basis)
offset_vec = shear_back_offset
propagation_vec = z_dir
E0 = get_E(Xs,Ys,Zs_0, correct=True)
E0 = get_E(Xs, Ys, Zs_0, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape, new_basis, wavelength,offset_vec,
dtype=t.complex128, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
E0.shape, new_basis, wavelength, offset_vec,
dtype=t.complex128, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
for Rrand in Rrands:
Xs,Ys,Zs_0 = np.tensordot(Rrand, np.tensordot(Rboth,Positions,axes=1),axes=1)
Xs, Ys, Zs_0 = np.tensordot(Rrand, np.tensordot(Rboth, Positions, axes=1), axes=1)
rot_offset = np.dot(Rrand, offset_vec)
rot_basis = np.dot(Rrand, new_basis)
rot_prop = np.dot(Rrand, propagation_vec)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape, rot_basis, wavelength, rot_offset,
dtype=t.complex128, propagation_vector=rot_prop)
Ez_rot_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
E0.shape, rot_basis, wavelength, rot_offset,
dtype=t.complex128, propagation_vector=rot_prop)
Ez_rot_t = propagators.near_field(t.as_tensor(E0), asp).numpy()
assert np.max(np.abs(Ez_t - Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
assert np.max(np.abs(Ez_t-Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
def test_inverse_near_field():
x = (np.arange(800) - 400) * 1.5e-9
y = (np.arange(1200) - 600) * 1e-9
Ys,Xs = np.meshgrid(y,x)
Rs = np.sqrt(Xs**2+Ys**2)
wavelength = 3e-9 #nm
sigma = 20e-9 #nm
z = 1000e-9 #nm
Ys, Xs = np.meshgrid(y, x)
Rs = np.sqrt(Xs**2 + Ys**2)
w0 = np.sqrt(2)*sigma
wavelength = 3e-9 # nm
sigma = 20e-9 # nm
z = 1000e-9 # nm
w0 = np.sqrt(2) * sigma
E0 = np.exp(-Rs**2 / w0**2)
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
E0.shape, (1.5e-9, 1e-9), wavelength, z, dtype=t.complex128)
E0 = t.as_tensor(E0, dtype=t.complex128)
E_prop = propagators.near_field(E0, asp)
E0 = t.as_tensor(E0,dtype=t.complex128)
E_prop = propagators.near_field(E0,asp)
E_backprop = propagators.inverse_near_field(E_prop, asp)
# We just want to check that it actually is the inverse
assert t.all(t.isclose(E0,E_backprop))
assert t.all(t.isclose(E0, E_backprop))