Commit b1cac9b7 authored by Jason Reneuve's avatar Jason Reneuve
Browse files

speed up load_time_series

parent 09928674fc32
...@@ -358,39 +358,34 @@ class SpatioTemporalSpectra(SpecificOutput): ...@@ -358,39 +358,34 @@ class SpatioTemporalSpectra(SpecificOutput):
f"spect_{k}": np.empty(spect_shape, dtype="complex") f"spect_{k}": np.empty(spect_shape, dtype="complex")
for k in self.keys_fields for k in self.keys_fields
} }
# loop on times # loop on all files
for time in track(times, description="Rearranging..."): for path_file in track(paths, description="Rearranging..."):
# loop on ranks with h5py.File(path_file, "r") as file:
for rank in ranks: # check if the file contains the time we're looking for
for path_file in paths: tmin_file = file["times"][0]
if not path_file.name.startswith(f"rank{rank:05}"): tmax_file = file["times"][-1]
continue if (tmin_file > tmax) or (tmax_file < tmin):
with h5py.File(path_file, "r") as file: continue
# check if the file contains the time we're looking for
# add [:] in case file contains only one time # time indices
# or else : h5py TypeError! times_file = file["times"][:]
tmin_file, tmax_file = file["times"][:][[0, -1]] cond_file = (times_file >= tmin) & (times_file <= tmax)
if (time < tmin_file) or (time > tmax_file): its_file = np.where(cond_file)[0]
continue its = np.where(times == times_file[cond_file])[0]
# time indices # k_adim_loc = global probes indices!
it = np.where(times == time)[0] ik0 = file["probes_k0adim_loc"][:]
it_file = np.where(file["times"][:] == time)[0] ik1 = file["probes_k1adim_loc"][:]
ik2 = file["probes_k2adim_loc"][:]
# k_adim_loc = global probes indices!
ik0 = file["probes_k0adim_loc"][:] # load data at desired times for all keys_fields
ik1 = file["probes_k1adim_loc"][:] for key in self.keys_fields:
ik2 = file["probes_k2adim_loc"][:] skey = f"spect_{key}"
data = file[skey + "_loc"][:, its_file]
# load data at time t for all keys_fields for i in range(its.size):
for key in self.keys_fields: series[skey][ik0, ik1, ik2, its[i]] = data[
skey = f"spect_{key}" :, i
series[skey][ik0, ik1, ik2, it] = file[skey + "_loc"][ ].transpose()
:, it_file
].transpose()
# stop opening files when we've reached the right one
break
# add Ki_adim arrays, times and dims order # add Ki_adim arrays, times and dims order
k0_adim = np.r_[0 : ik0max + 1, ik0min:0] k0_adim = np.r_[0 : ik0max + 1, ik0min:0]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment