diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..9214cb9 --- /dev/null +++ b/.flake8 @@ -0,0 +1,2 @@ +[flake8] +ignore = E501, W503 \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index 781dd9e..850935d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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' diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index 3f78a73..e5422d7 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -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 - - diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py index f77c589..b6b1868 100644 --- a/tests/models/test_simple_ptycho.py +++ b/tests/models/test_simple_ptycho.py @@ -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 diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 8999117..4d17f0b 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -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, :]) diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index 0c84a27..d293c9f 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -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)) + diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index 39b64dd..7f2571c 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -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]) diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index 83c34d8..a6385cd 100644 --- a/tests/tools/test_image_processing.py +++ b/tests/tools/test_image_processing.py @@ -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) diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index 53c284a..5430271 100644 --- a/tests/tools/test_initializers.py +++ b/tests/tools/test_initializers.py @@ -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 diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index ae3d6bd..563d8f6 100644 --- a/tests/tools/test_interactions.py +++ b/tests/tools/test_interactions.py @@ -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) - diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index 633027d..b36b85d 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -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)) diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py index bfa06de..93ac72c 100644 --- a/tests/tools/test_measurements.py +++ b/tests/tools/test_measurements.py @@ -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, diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py index b377ef6..baca33d 100644 --- a/tests/tools/test_plotting.py +++ b/tests/tools/test_plotting.py @@ -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() diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index ec55140..d571b41 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -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))