diff --git a/CDTools/tools/data.py b/CDTools/tools/data.py index 4ef1796..63bf9d6 100644 --- a/CDTools/tools/data.py +++ b/CDTools/tools/data.py @@ -2,6 +2,11 @@ from __future__ import division, print_function, absolute_import import h5py import numpy as np +import numbers +import datetime +import dateutil.parser +import torch as t +from contextlib import contextmanager __all__ = ['get_entry_info', 'get_sample_info', @@ -9,7 +14,15 @@ __all__ = ['get_entry_info', 'get_detector_geometry', 'get_mask', 'get_data', - 'get_ptycho_translations'] + 'get_ptycho_translations', + 'create_cxi', + 'add_entry_info', + 'add_sample_info', + 'add_source', + 'add_detector', + 'add_mask', + 'add_data', + 'add_ptycho_translations'] # # @@ -56,13 +69,19 @@ __all__ = ['get_entry_info', # + # # Functions to inspect the basic attributes of a cxi file represented as an # h5 file object # + def get_entry_info(cxi_file): """Returns a dictionary with the basic metadata from the cxi file's entry_1 attribute + + String type metadata is read out as a string, and datetime metadata + is converted to python datetime objects if the string is properly + formatted. Args: cxi_file (h5py.File) : a file object to be read @@ -75,11 +94,17 @@ def get_entry_info(cxi_file): metadata_attrs = ['title', 'experiment_identifier', 'experiment_description', - 'program_name', - 'start_time', - 'end_time'] + 'program_name'] metadata = {attr: str(e1[attr][()].decode()) for attr in metadata_attrs if attr in e1} + datetime_attrs = ['start_time', + 'end_time'] + for attr in datetime_attrs: + if attr in e1: + try: + metadata[attr] = dateutil.parser.parse(str(e1[attr][()].decode())) + except ValueError: + metadata[attr] = str(e1[attr][()].decode()) return metadata @@ -317,28 +342,241 @@ def get_ptycho_translations(cxi_file): return translations + # # It might be useful to make some helper functions to help write cxi files # -# -# A function to place the skeleton of a cxi file down -# -# -# A function to define the source attributes -# +def create_cxi(filename): + """Creates a new cxi file with a single entry group -# -# A function to define the detector geometry -# + Args: + filename (str) : The path at which to create the file + """ + file_obj = h5py.File(filename,'w') + file_obj.create_dataset('cxi_version', data=160) + file_obj.create_dataset('number_of_entries',data=1) + e1f = file_obj.create_group('entry_1') + return file_obj -# -# A function to save out a mask, converting it to the correct format -# -# -# Perhaps a function to store the data and link it correctly? But this might -# have to change too much situation to situation -# +def add_entry_info(cxi_file, metadata): + """Adds a dictionary of entry metadata to the entry_1 group of a cxi file object + + Args: + cxi_file (h5py.File) : The file to add the info to + metadata (dict) : A dictionary containing all the metadata to be stored + """ + # Just the string and datetime types should be relevant but all are + # included in case the cxi spec becomes more permissive + for key, value in metadata.items(): + if isinstance(value,(str,bytes)): + cxi_file['entry_1'][key] = np.string_(value) + elif isinstance(value, datetime.datetime): + cxi_file['entry_1'][key] = np.string_(value.isoformat()) + elif isinstance(value, numbers.Number): + si[key] = value + elif isinstance(value, (np.ndarray,list,tuple)): + s1.create_dataset(key, data=np.asarray(value)) + elif isinstance(value, t.Tensor): + asnumpy = value.detach().cpu().numpy() + cxi_file['entry_1'].create_dataset(key, data=asnumpy) + +def add_sample_info(cxi_file, metadata): + """Adds a dictionary of entry metadata to the entry_1/sample_1 group of a cxi file object + + This function will create the sample_1 attribute if it doesn't already exist + + Args: + cxi_file (h5py.File) : The file to add the info to + metadata (dict) : A dictionary containing all the metadata to be stored + """ + if 'entry_1/sample_1' not in cxi_file: + cxi_file['entry_1'].create_group('sample_1') + s1 = cxi_file['entry_1/sample_1'] + + for key, value in metadata.items(): + if isinstance(value,(str,bytes)): + s1[key] = np.string_(value) + elif isinstance(value, datetime.datetime): + s1[key] = np.string_(value.isoformat()) + elif isinstance(value, numbers.Number): + s1[key] = value + elif isinstance(value, (np.ndarray,list,tuple)): + s1.create_dataset(key, data=np.asarray(value)) + elif isinstance(value, t.Tensor): + asnumpy = value.detach().cpu().numpy() + s1.create_dataset(key, data=asnumpy) + + +def add_source(cxi_file, wavelength): + """Adds the entry_1/source_1 group to a cxi file object + + It stores the energy and wavelength attributes in the source_1 group, + given a wavelength to define them from. + + Args: + cxi_file (h5py.File) : The file to add the source to + wavelength (float) : The wavelength of light + """ + if 'entry_1/instrument_1' not in cxi_file: + cxi_file['entry_1'].create_group('instrument_1') + i1 = cxi_file['entry_1/instrument_1'] + if 'source_1' not in i1: + i1.create_group('source_1') + s1 = i1['source_1'] + s1['wavelength'] = np.float32(wavelength) + s1['energy'] = np.float32(1.9864459e-25 / wavelength) + + + +def add_detector(cxi_file, distance, basis, corner=None): + """Adds the entry_1/instrument_1/detector_1 group to a cxi file object + + It will define all the relevant parameters - distance, pixel size, + detector basis, and corner position (if relevant) based on the provided + information + + Args: + cxi_file (h5py.File) : The file to add the detector to + distance (float) : The sample to detector distance + basis (array_like) : The detector basis + corner (array_like) : Optional, the corner position of the detector + + """ + if 'entry_1/instrument_1' not in cxi_file: + cxi_file['entry_1'].create_group('instrument_1') + i1 = cxi_file['entry_1/instrument_1'] + if 'detector_1' not in i1: + i1.create_group('detector_1') + d1 = i1['detector_1'] + + d1['distance'] = np.float32(distance) + d1['x_pixel_size'] = np.linalg.norm(basis[:,1]) + d1['y_pixel_size'] = np.linalg.norm(basis[:,0]) + if isinstance(basis, t.Tensor): + basis = basis.detach().cpu().numpy() + d1.create_dataset('basis_vectors', data=basis) + if corner is not None: + if isinstance(corner, t.Tensor): + corner = corner.detach().cpu().numpy() + d1.create_dataset('corner_position',data=corner) + + +def add_mask(cxi_file, mask): + """Adds the specified mask to the cxi file + + It places the mask into the mask dataset under + entry_1/instrument_1/detector_1. The internal mask is defined + simply as a 1 for an "on" pixel and a 0 for an "off" pixel, and + the saved mask is exactly the opposite. This is simpler than the + most general mask allowed by the cxi file format but it captures the + distinction between pixels to be used and pixels not to be used. + + Args: + cxi_file (h5py.File) : The file to add the mask to + mask (array_like) : The mask to save out to the file + """ + + if 'entry_1/instrument_1' not in cxi_file: + cxi_file['entry_1'].create_group('instrument_1') + i1 = cxi_file['entry_1/instrument_1'] + if 'detector_1' not in i1: + i1.create_group('detector_1') + d1 = i1['detector_1'] + if isinstance(mask, t.Tensor): + mask = mask.detach().cpu().numpy() + + mask_to_save = np.zeros(mask.shape).astype(np.uint32) + mask_to_save[mask == 0] = 1 + d1.create_dataset('mask',data=mask_to_save) + + +def add_data(cxi_file, data, axes=None): + """Adds the specified data to the cxi file + + It will add the data unchanged to the file, placing it in two spots: + + 1) The entry_1/instrument_1/detector_1/data path + 2) A softlink at entry_1/data_1/data + + Args: + cxi_file (h5py.File) : The file to add the data to + data (array_like) : The data to be saved + axes (list) : Optional, a list of axis names to be saved in the axes attribute + """ + if 'entry_1/data_1' not in cxi_file: + cxi_file['entry_1'].create_group('data_1') + data1 = cxi_file['entry_1/data_1'] + + if 'entry_1/instrument_1' not in cxi_file: + cxi_file['entry_1'].create_group('instrument_1') + i1 = cxi_file['entry_1/instrument_1'] + if 'detector_1' not in i1: + i1.create_group('detector_1') + det1 = i1['detector_1'] + + if isinstance(data, t.Tensor): + data = data.detach().cpu().numpy() + + det1.create_dataset('data', data=data) + data1['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data') + + if axes is not None: + if isinstance(axes, list): + axes_str = ':'.join(axes) + else: + axes_str = str(axes) + det1['data'].attrs['axes'] = np.string_(axes_str) + + +def add_ptycho_translations(cxi_file, translations): + """Adds the specified translations to the cxi file + + It will add the translations to the file, negating them to conform to + the standard in cxi files that the translations refer to the object's + translation, and also transposing them to match the cxi file specification. + + It will generally store them in 3 places: + + 1) The entry_1/sample_1/geometry_1/translation path + 2) A softlink at entry_1/data_1/translation + 3) A softlink at entry_1/instrument_1/detector_1/translation + + Args: + cxi_file (h5py.File) : The file to add the translations to + translations (array_like) : The translations to be saved + """ + + if 'entry_1/sample_1' not in cxi_file: + cxi_file['entry_1'].create_group('sample_1') + s1 = cxi_file['entry_1/sample_1'] + + if 'geometry_1' not in s1: + s1.create_group('geometry_1') + g1 = s1['geometry_1'] + + if 'entry_1/data_1' not in cxi_file: + cxi_file['entry_1'].create_group('data_1') + data1 = cxi_file['entry_1/data_1'] + + if 'entry_1/instrument_1' not in cxi_file: + cxi_file['entry_1'].create_group('instrument_1') + i1 = cxi_file['entry_1/instrument_1'] + if 'detector_1' not in i1: + i1.create_group('detector_1') + det1 = i1['detector_1'] + + + if isinstance(translations, t.Tensor): + translations = translations.detach().cpu().numpy() + + # accounting for the different definition between cxi files and + # CDTools + translations = -translations.transpose() + + g1.create_dataset('translation', data=translations) + data1['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') + det1['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') diff --git a/setup.py b/setup.py index c543375..de2f1ef 100644 --- a/setup.py +++ b/setup.py @@ -16,6 +16,7 @@ setuptools.setup( "numpy", "scipy", "matplotlib", + "dateutil", #"pytorch", "h5py"], packages=setuptools.find_packages(), diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index ed8b508..78630c3 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -7,6 +7,7 @@ import h5py import pytest import os import datetime +import numbers from pathlib import Path @@ -54,10 +55,10 @@ def ptycho_cxi_1(): e1f = f.create_group('entry_1') expected['entry metadata'] = {} e1e = expected['entry metadata'] - e1e['start_time'] = datetime.datetime.now().isoformat() - e1f['start_time'] = np.string_(e1e['start_time']) - e1e['end_time'] = datetime.datetime.now().isoformat() - e1f['end_time'] = np.string_(e1e['end_time']) + e1e['start_time'] = datetime.datetime.now() + e1f['start_time'] = np.string_(e1e['start_time'].isoformat()) + e1e['end_time'] = datetime.datetime.now() + e1f['end_time'] = np.string_(e1e['end_time'].isoformat()) e1e['experiment_identifier'] = 'Fake Experiment 1' e1f['experiment_identifier'] = np.string_(e1e['experiment_identifier']) e1e['experiment_description'] = 'A fully defined ptychography experiment to test the data loading' @@ -134,7 +135,9 @@ def ptycho_cxi_1(): d1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') expected['translations'] = -translations.transpose() - return f, expected + yield f, expected + + f.close() @pytest.fixture(scope='module') @@ -208,7 +211,9 @@ def ptycho_cxi_2(): g1f.create_dataset('translation',data=translations) expected['translations'] = -translations.transpose() - return f, expected + yield f, expected + + f.close() @pytest.fixture(scope='module') @@ -236,10 +241,10 @@ def ptycho_cxi_3(): e1f = f.create_group('entry_1') expected['entry metadata'] = {} e1e = expected['entry metadata'] - e1e['start_time'] = datetime.datetime.now().isoformat() - e1f['start_time'] = np.string_(e1e['start_time']) - e1e['end_time'] = datetime.datetime.now().isoformat() - e1f['end_time'] = np.string_(e1e['end_time']) + e1e['start_time'] = datetime.datetime.now() + e1f['start_time'] = np.string_(e1e['start_time'].isoformat()) + e1e['end_time'] = datetime.datetime.now() + e1f['end_time'] = np.string_(e1e['end_time'].isoformat()) # Set up the sample info expected['sample info'] = None @@ -280,9 +285,11 @@ def ptycho_cxi_3(): data1f.create_dataset('translation',data=translations) expected['translations'] = -translations.transpose() - return f, expected + yield f, expected + + f.close() + -# # 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 # ensure compatibility with many beamlines @@ -373,3 +380,156 @@ 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') 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', + 'title': 'my cool experiment', + 'start_time': datetime.datetime.now(), + 'end_time': datetime.datetime.now()} + + 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') as f: + read_entry_info = data.get_entry_info(f) + + for key in entry_info: + if isinstance(entry_info[key], np.ndarray): + assert np.allclose(entry_info[key], read_entry_info[key]) + else: + assert entry_info[key] == read_entry_info[key] + + +def test_add_sample_info(tmp_path): + 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.])} + + 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') as f: + read_sample_info = data.get_sample_info(f) + + for key in sample_info: + if isinstance(sample_info[key], np.ndarray): + assert np.allclose(sample_info[key], read_sample_info[key]) + elif isinstance(sample_info[key], numbers.Number): + 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 + energy = 1.9864459e-25 / wavelength + + 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') 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) + + +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) + + 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') as f: + # Check this directly since we want to make sure it saved + # the pixel sizes + d1 = f['entry_1/instrument_1/detector_1'] + read_basis = np.array(d1['basis_vectors']) + read_x_pix = np.float32(d1['x_pixel_size']) + read_y_pix = np.float32(d1['y_pixel_size']) + read_distance = np.float32(d1['distance']) + read_corner = np.array(d1['corner_position']) + + 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) + + +def test_add_mask(tmp_path): + 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') as f: + read_mask = data.get_mask(f) + + assert np.all(mask == read_mask) + + +def test_add_data(tmp_path): + # First test from numpy, with axes + 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') as f: + # Check this directly since we want to make sure it saved + # it in all the places it should have + read_data_1 = np.array(f['entry_1/data_1/data']) + read_data_2 = np.array(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') as f: + read_data, axes = data.get_data(f) + + assert np.allclose(fake_data.numpy(),read_data) + + +def test_add_ptycho_translations(tmp_path): + + 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') as f: + # Check this directly since we want to make sure it saved + # it in all the places it should have + read_translations_1 = np.array(f['entry_1/data_1/translation']) + read_translations_2 = np.array(f['entry_1/instrument_1/detector_1/translation']) + read_translations_3 = np.array(f['entry_1/sample_1/geometry_1/translation']) + + assert np.allclose(-translations.transpose(), read_translations_1) + assert np.allclose(-translations.transpose(), read_translations_2) + assert np.allclose(-translations.transpose(), read_translations_3)