Coverage for gamdpy/tools/TrajectoryIO.py: 8%

1# This class is used for loading the output of a simulation as a dictionary following the formatting of sim.output or save sim.output to file

2# Can be also used to convert between output formats (rumd3 -> gamdpy supported so far)

4import sys

5import os

6import time

8import numpy as np

9import h5py

12# This class is a wrapper for several possible inputs/output

13class TrajectoryIO():

14 """

15 This class handles the loading and saving of simulation data.

16 When the class can be instantiated with an output from a previous simulation.

17 The data can be saved in self.h5 in the same format of the sim.output object.

18 When the class is instantiated without an input, self.h5 is None and can be assigned afterward (used to save output from memory simulation)

20 Parameters

21 ----------

23 name : str

24 Name of the file or folder to read output from. Can be a gamdpy .h5 output or a rumd3 TrajectoryFiles folder.

26 Examples

27 --------

29 >>> import gamdpy as gp

30 >>> import h5py

31 >>> output = gp.tools.TrajectoryIO("examples/Data/LJ_r0.973_T0.70_toread.h5").get_h5()

32 Found .h5 file (examples/Data/LJ_r0.973_T0.70_toread.h5), loading to gamdpy as output dictionary

33 >>> nblocks, nconfs, N, D = output['trajectory_saver/positions'].shape

34 >>> print(f"Output file examples/Data/LJ_r0.973_T0.70.h5 contains a simulation of {N} particles in {D} dimensions")

35 Output file examples/Data/LJ_r0.973_T0.70.h5 contains a simulation of 2048 particles in 3 dimensions

36 >>> print(f"The simulation output is divided into {nblocks} blocks, each of them with {nconfs} configurations")

37 The simulation output is divided into 32 blocks, each of them with 12 configurations

38 >>> output = gp.tools.TrajectoryIO("examples/Data/NVT_N4000_T2.0_rho1.2_KABLJ_rumd3/TrajectoryFiles").get_h5() # Read from rumd3

39 Found rumd3 TrajectoryFiles, loading to rumpdy as output dictionary

40 >>> nblocks, nconfs, N, D = output['trajectory_saver/positions'].shape

41 >>> print(f"File examples/Data/NVT_N4000_T2.0_rho1.2_KABLJ_rumd3/TrajectoryFiles contains a simulation of {N} particles in {D} dimensions")

42 File examples/Data/NVT_N4000_T2.0_rho1.2_KABLJ_rumd3/TrajectoryFiles contains a simulation of 4000 particles in 3 dimensions

43 >>> print(f"The simulation output is divided into {nblocks} blocks, each of them with {nconfs} configurations")

44 The simulation output is divided into 2 blocks, each of them with 24 configurations

46 """

48 def __init__(self, name=""):

49 import importlib.util # This python library can be used if module can be imported

51 ## Following lines defines the initialization behavior depending on given input

52 if name[-3:]==".h5":

53 modification_time = os.path.getmtime(name)

54 readable_time = time.ctime(modification_time)

55 #print(f"Found .h5 file ({name}, {readable_time}), loading to gamdpy as output dictionary")

56 print(f"Found .h5 file ({name}), loading to gamdpy as output dictionary") # Cant handle timestamp in doctest

57 self.h5 = self.load_h5(name)

58 elif "TrajectoryFiles" in name:

59 try:

60 assert os.path.isdir(name)==True

61 print("Found rumd3 TrajectoryFiles, loading to rumpdy as output dictionary")

62 self.h5 = self.load_rumd3(name)

63 except:

64 raise Exception(f"Folder {name} doesn't exists")

65 elif name=="":

66 print("Warning: class initialized without input data, set self.h5 manually")

67 self.h5 = None

68 else:

69 print("Input not recognized, unsupported format")

70 self.h5 = None

72 ## Following lines deals with the compression of the output

73 # The class is initialized with some default compression settings depending on what is available.

74 # These parameters can be changed by the user by changing the class attributes

76 # Checks if hdf5plugin lib is available

77 if importlib.util.find_spec("hdf5plugin")==None:

78 self.compression_type = "gzip"

79 self.compression_opts = 6 # seems the best option in terms of compression and timing

80 else:

81 import hdf5plugin

82 # if hdf5plugin is available use BZip2 compression (slightly better and faster)

83 self.compression_type = hdf5plugin.BZip2()

85 def get_h5(self) -> h5py.File:

86 """ Returns self.h5 """

87 return self.h5

89 def load_h5(self, name:str) -> h5py.File:

90 """ Makes self.h5 a view of the .h5 files """

91 return h5py.File(name, "r")

93 # Load from TrajectoryFiles (std rumd3 output)

94 # It assumes trajectories are spaced in log2

95 def load_rumd3(self, name:str) -> h5py.File:

96 """ Reads a rumd3 TrajectoryFiles folder and convert it into gamdpy .h5 output. Assumes RectangularSimulationBox, corresponding to gamdpy's Orthorhombic.

97 This function returns a memory .h5"""

98 import os, gzip, glob

99 import pandas as pd

100

101 # Rumd3 output is always in D=3

102 dim = 3

103 # Checks if energy output is present

104 if "LastComplete_energies.txt" not in os.listdir(name):

105 print("LastComplete_energies.txt not present")

106 energy = False

107 else:

108 energy = True

109

110 # Checks if trajectory output is present

111 if "LastComplete_trajectory.txt" not in os.listdir(name):

112 print("LastComplete_trajectory.txt not present")

113 traj = False

114 else:

115 traj = True

116

117 # Exit if no output is found in TrajectoryFiles folder

118 if not energy and not traj:

119 print(f"The folder {name} has no LastComplete_*.txt files, returning None")

120 return None

121

122 # Reads block information from LastComplete_*

123 if energy: nblocks, blocksize = np.loadtxt(f"{name}/LastComplete_energies.txt", dtype=np.int32)

124 else : nblocks, blocksize = np.loadtxt(f"{name}/LastComplete_trajectory.txt", dtype=np.int32)

125

126 # Defining output memory .h5 file

127 fullpath = f"{name}"

128 output = h5py.File(f"{id(fullpath)}.h5", "w", driver='core', backing_store=False)

129 assert isinstance(output, h5py.File), "Error creating memory h5 file in TrajectoryIO.load_rumd3"

130

131 # Read trajectories

132 if traj:

133 traj_files = sorted(glob.glob(f"{name}/trajectory*"))

134 # Remove last block if incomplete

135 if traj_files[-1]==f"{name}/trajectory{nblocks+1:04d}.xyz.gz": traj_files = traj_files[:-1]

136 # Read metadata from first file in the list

137 with gzip.open(f"{traj_files[0]}", "r") as f:

138 npart = int(f.readline())

139 cmt_line = f.readline().decode().split()

140 meta_data = dict()

141 for item in cmt_line:

142 key, val = item.split("=")

143 meta_data[key] = val

144 ntrajinblock = int(meta_data['logLin'].split(",")[3])

145 num_types = int(meta_data['numTypes'])

146 masses = [float(x) for x in meta_data['mass'].split(',')]

147 assert len(masses) == num_types

148 if meta_data['ioformat'] == '1':

149 lengths = np.array([float(x) for x in meta_data['boxLengths'].split(',')], dtype=np.float32)

150 else:

151 assert meta_data['ioformat'] == '2'

152 sim_box_data = meta_data['sim_box'].split(',')

153 sim_box_type = sim_box_data[0]

154 sim_box_params = [float(x) for x in sim_box_data[1:]]

155 assert sim_box_type == 'RectangularSimulationBox'

156 lengths = np.array(sim_box_params, dtype=np.float32)

157 integrator_data = meta_data['integrator'].split(',')

158 timestep = integrator_data[1]

159 # Loop over the files and read them assuming each line is type, x, y, z, imx, imy, imz

160 #ntrajinblock = int(1+np.log2(blocksize))

161 toskip1 = np.array([ (npart+2)*x for x in range(1+ntrajinblock)])

162 toskip2 = np.array([1+(npart+2)*x for x in range(1+ntrajinblock)])

163 toskip = sorted(list(np.concatenate((toskip1, toskip2))))

164 #print(f"Found informations about trajectory files: {nblocks=} {blocksize=} {ntrajinblock=}")

165 positions = list()

166 images = list()

167 for trajectory in traj_files:

168 tmp_data = pd.read_csv(trajectory, skiprows = toskip, usecols=[0,1,2,3,4,5,6], names=["type", "x", "y", "z", "imx", "imy", "imz"], delimiter=" ")

169 type_array = tmp_data['type'].to_numpy()

170 pos_array = np.c_[tmp_data['x'].to_numpy(), tmp_data['y'].to_numpy(), tmp_data['z'].to_numpy()]

171 img_array = np.c_[tmp_data['imx'].to_numpy(), tmp_data['imy'].to_numpy(), tmp_data['imz'].to_numpy()]

172 positions.append(pos_array.reshape((-1,npart,dim)))

173 images.append(img_array.reshape((-1,npart,dim)))

174 # Saving data in output h5py

175 output.attrs['dt'] = timestep

176 output.attrs['simbox_initial'] = lengths

177 #output.attrs['vectors_names'] = ["r", "r_im"]

178 output.create_group('trajectory_saver')

179 output['trajectory_saver'].create_dataset('positions', shape=(len(traj_files), 1+ntrajinblock, npart, dim), dtype=np.float32)

180 output['trajectory_saver/positions'][:,:,:,:] = np.array(positions)

181 output['trajectory_saver'].create_dataset('images', shape=(len(traj_files), 1+ntrajinblock, npart, dim), dtype=np.int32)

182 output['trajectory_saver/images'][:,:,:,:] = np.array(images)

183

184 #output.create_dataset("ptype", data=type_array[:npart], shape=(npart), dtype=np.int32)

185 output.create_group('initial_configuration')

186 output['initial_configuration'].create_dataset("ptype", data=type_array[:npart], shape=(npart), dtype=np.int32)

187

188 # Read energies

189 if energy:

190 energy_files = sorted(glob.glob(f"{name}/energies*"))

191 # Remove last block if incomplete

192 if energy_files[-1]==f"{name}/energies{nblocks+1:04d}.dat.gz": energy_files = energy_files[:-1]

193 # Read metadata from first file in the list

194 with gzip.open(f"{energy_files[0]}", "r") as f:

195 # ioformat=2 N=4096 Dt=147.266830 columns=ke,pe,p,T,Etot,W (example of lin saving)

196 # ioformat=2 N=4096 timeStepIndex=0 logLin=0,2,0,17,0 columns=ke,pe,p,T,Etot,W (example of log saving)

197 cmt_line = f.readline().decode().split()[1:]

198 meta_data = dict()

199 for item in cmt_line:

200 key, val = item.split("=")

201 meta_data[key] = val

202 npart = meta_data['N']

203 if 'Dt' in meta_data: save_interval = meta_data['Dt']

204 elif 'logLin' in meta_data: save_interval = f'log{meta_data["logLin"][1]}'

205 col_names = meta_data['columns'].split(",")

206 all_energies = list()

207 for energies in energy_files:

208 tmp_data = pd.read_csv(energies, skiprows=1, names=col_names, usecols = [i for i in range(len(col_names))], delimiter=" ")

209 all_energies.append(tmp_data.to_numpy())

210 # Saving data in output h5py

211 grp = output.create_group('scalar_saver')

212 if 'dt' in output.attrs.keys() and 'Dt' in meta_data:

213 output['scalar_saver'].attrs['steps_between_output'] = float(save_interval)/float(output.attrs['dt'])

214 output['scalar_saver'].attrs['time_between_output'] = save_interval

215 output['scalar_saver'].attrs['scalar_names'] = list(col_names)

216 output.create_dataset('scalar_saver/scalars', data=np.vstack(all_energies))

217

218 return output

219

220 def save_h5(self, name:str):

221 """

222 This method saves self.h5 to disk.

223 It can be used to save sim.output to file if class is initialized without arguments and then self.h5 = sim.output .

224 By default output is compressed. This can be avoided setting self.compression_type = "gzip" and self.compression_opts=0 .

225 """

226

227 # LC: This monster needs to be re-structured

228 import h5py

229 import importlib.util

230 if importlib.util.find_spec("hdf5plugin")!=None:

231 import hdf5plugin

232 fout = h5py.File(name, "w")

233 #fout.attrs.update(self.h5.attrs)

234 print(self.h5.keys())

235 for key in self.h5.keys():

236 fout.copy(source=self.h5[key], dest="/")

237 print(f"{key} {fout.keys()}")

238 fout.close()

239 return

240 for key in self.h5.keys():

241 if isinstance(self.h5[key], h5py.Group):

242 print(f"Writing Group {key} to {name}")

243 fout.create_group(f"/{key}")

244 fout.copy(self.h5[key], f"/{key}")

245 continue

246 #print(self.h5.keys(), key, fout.keys())

247 for subkey in self.h5[key].keys():

248 if isinstance(self.h5[key][subkey], h5py.Group):

249 #print(self.h5[key].keys(), subkey, self.h5[key][subkey], fout[key].keys())

250 #print(f"Writing dataset {key}/{subkey} to {name}")

251 fout.create_group(f"/{key}/{subkey}")

252 for subsubkey in self.h5[key][subkey].keys():

253 print(key, subkey, subsubkey)

254 if importlib.util.find_spec("hdf5plugin")!=None:

255 if self.compression_type == "gzip":

256 fout[f"{key}/{subkey}"].create_dataset(subsubkey, data=self.h5[f"{key}/{subkey}"][subsubkey], chunks=True,

257 compression=self.compression_type, compression_opts=self.compression_opts, shuffle=True)

258 else:

259 fout[f"{key}/{subkey}"].create_dataset(subsubkey, data=self.h5[f"{key}/{subkey}"][subsubkey], chunks=True,

260 compression=self.compression_type, shuffle=True)

261 elif self.compression_type == "gzip":

262 fout[f"{key}/{subkey}"].create_dataset(subsubkey, data=self.h5[f"{key}/{subkey}"][subsubkey], chunks=True,

263 compression=self.compression_type, compression_opts=self.compression_opts, shuffle=True)

264 else:

265 fout[f"{key}/{subkey}"].create_dataset(subsubkey, data=self.h5[f"{key}/{subkey}"][subsubkey], chunks=True,

266 compression=self.compression_type)

267 #print(f"Writing attributes {key}/{subkey} to {name}")

268 fout[f"{key}/{subkey}"][subsubkey].attrs.update(self.h5[key][subkey].attrs)

269 #print(f"Done dataset {key}/{subkey} to {name}")

270 else:

271 #print(self.h5[key].keys(), subkey, self.h5[key][subkey], fout[key].keys())

272 #print(f"Writing dataset {key}/{subkey} to {name}")

273 if importlib.util.find_spec("hdf5plugin")!=None:

274 if self.compression_type == "gzip":

275 fout[key].create_dataset(subkey, data=self.h5[key][subkey], chunks=True,

276 compression=self.compression_type, compression_opts=self.compression_opts, shuffle=True)

277 else:

278 fout[key].create_dataset(subkey, data=self.h5[key][subkey], chunks=True,

279 compression=self.compression_type, shuffle=True)

280 elif self.compression_type == "gzip":

281 fout[key].create_dataset(subkey, data=self.h5[key][subkey], chunks=True,

282 compression=self.compression_type, compression_opts=self.compression_opts, shuffle=True)

283 else:

284 fout[key].create_dataset(subkey, data=self.h5[key][subkey], chunks=True,

285 compression=self.compression_type)

286 #print(f"Writing attributes {key}/{subkey} to {name}")

287 fout[key][subkey].attrs.update(self.h5[key][subkey].attrs)

288 #print(f"Done dataset {key}/{subkey} to {name}")

289 fout[key].attrs.update(self.h5[key].attrs)

290 #print(f"Done Group {key} to {name}")

291 elif isinstance(self.h5[key], h5py.Dataset):

292 print(f"Writing dataset {key} to {name}")

293 if importlib.util.find_spec("hdf5plugin")!=None:

294 if self.compression_type == "gzip":

295 fout.create_dataset(key, data=self.h5[key], chunks=True, compression=self.compression_type, compression_opts=self.compression_opts, shuffle=True)

296 else:

297 fout.create_dataset(key, data=self.h5[key], chunks=True, compression=self.compression_type, shuffle=True)

298 elif self.compression_type == "gzip":

299 fout.create_dataset(key, data=self.h5[key], chunks=True, compression=self.compression_type, compression_opts=self.compression_opts, shuffle=True)

300 else:

301 fout.create_dataset(key, data=self.h5[key], chunks=True, compression=self.compression_type)

302 fout[key].attrs.update(self.h5[key].attrs)

303 else:

304 print("Problem")

305

306 print(f"All written using {self.compression_type}")

307 fout.close()

308 print("Output file closed")

309 return

310