Write the data saving tools and finish the data tests

This commit is contained in:
Abe Levitan
2019-03-25 19:11:02 -04:00
parent a398151700
commit c73caa3a06
3 changed files with 431 additions and 32 deletions
+258 -20
View File
@@ -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')
+1
View File
@@ -16,6 +16,7 @@ setuptools.setup(
"numpy",
"scipy",
"matplotlib",
"dateutil",
#"pytorch",
"h5py"],
packages=setuptools.find_packages(),
+172 -12
View File
@@ -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)