Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

""" 

Helper functions for reading and writing objects to tar files 

""" 

 

import pickle 

import tempfile 

import warnings 

import ase.io as aseIO 

with warnings.catch_warnings(): 

warnings.filterwarnings('ignore', category=FutureWarning) 

import h5py 

 

 

def add_ase_atoms_to_tarfile(tar_file, atoms, arcname, format='json'): 

""" Adds an ase.Atoms object to tar_file """ 

temp_file = tempfile.NamedTemporaryFile() 

aseIO.write(temp_file.name, atoms, format=format) 

temp_file.seek(0) 

tar_info = tar_file.gettarinfo(arcname=arcname, fileobj=temp_file) 

tar_file.addfile(tar_info, temp_file) 

 

 

def read_ase_atoms_from_tarfile(tar_file, arcname, format='json'): 

""" Reads ase.Atoms from tar file """ 

temp_file = tempfile.NamedTemporaryFile() 

temp_file.write(tar_file.extractfile(arcname).read()) 

temp_file.seek(0) 

atoms = aseIO.read(temp_file.name, format=format) 

return atoms 

 

 

def add_items_to_tarfile_hdf5(tar_file, items, arcname): 

""" Add items to one hdf5 file """ 

temp_file = tempfile.NamedTemporaryFile() 

hf = h5py.File(temp_file.name) 

for key, value in items.items(): 

hf.create_dataset(key, data=value, compression='gzip') 

hf.close() 

temp_file.seek(0) 

tar_info = tar_file.gettarinfo(arcname=arcname, fileobj=temp_file) 

tar_file.addfile(tar_info, temp_file) 

temp_file.close() 

 

 

def add_items_to_tarfile_pickle(tar_file, items, arcname): 

""" Add items by pickling them """ 

temp_file = tempfile.TemporaryFile() 

pickle.dump(items, temp_file) 

temp_file.seek(0) 

tar_info = tar_file.gettarinfo(arcname=arcname, fileobj=temp_file) 

tar_file.addfile(tar_info, temp_file) 

temp_file.close() 

 

 

def add_items_to_tarfile_custom(tar_file, items): 

""" Add items assuming they have a custom write function """ 

for key, value in items.items(): 

temp_file = tempfile.TemporaryFile() 

value.write(temp_file) 

temp_file.seek(0) 

tar_info = tar_file.gettarinfo(arcname=key, fileobj=temp_file) 

tar_file.addfile(tar_info, temp_file) 

temp_file.close() 

 

 

def add_list_to_tarfile_custom(tar_file, objects, arcname): 

""" Add list of objects assuming they have a custom write function """ 

for i, obj in enumerate(objects): 

temp_file = tempfile.TemporaryFile() 

obj.write(temp_file) 

temp_file.seek(0) 

fname = '{}_{}'.format(arcname, i) 

tar_info = tar_file.gettarinfo(arcname=fname, fileobj=temp_file) 

tar_file.addfile(tar_info, temp_file) 

temp_file.close() 

 

 

def read_items_hdf5(tar_file, arcname): 

""" Read items from hdf5file inside tar_file """ 

 

# read hdf5 

temp_file = tempfile.NamedTemporaryFile() 

temp_file.write(tar_file.extractfile(arcname).read()) 

temp_file.seek(0) 

hf = h5py.File(temp_file.name, 'r') 

items = {key: value[:] for key, value in hf.items()} 

hf.close() 

return items 

 

 

def read_items_pickle(tar_file, arcname): 

items = dict() 

items = pickle.load(tar_file.extractfile(arcname)) 

return items 

 

 

def read_list_custom(tar_file, arcname, read_function, **kwargs): 

objects = [] 

i = 0 

while True: 

try: 

fname = '{}_{}'.format(arcname, i) 

f = tar_file.extractfile(fname) 

obj = read_function(f, **kwargs) 

objects.append(obj) 

f.close() 

except KeyError: 

break 

i += 1 

return objects