Commit 0695290f authored by Jason Reneuve's avatar Jason Reneuve
Browse files

speedup load_time_series

parent 71aebd52c49a
...@@ -329,6 +329,12 @@ class SpatioTemporalSpectra(SpecificOutput): ...@@ -329,6 +329,12 @@ class SpatioTemporalSpectra(SpecificOutput):
times = [] times = []
for path in paths_1st_rank: for path in paths_1st_rank:
with h5py.File(path, "r") as file: with h5py.File(path, "r") as file:
tmax_file = file["times"][-1]
if tmax_file < tmin:
continue
tmin_file = file["times"][0]
if tmin_file > tmax:
break
times_file = file["times"][:] times_file = file["times"][:]
cond_times = (times_file >= tmin) & (times_file <= tmax) cond_times = (times_file >= tmin) & (times_file <= tmax)
times.append(times_file[cond_times]) times.append(times_file[cond_times])
...@@ -358,34 +364,41 @@ class SpatioTemporalSpectra(SpecificOutput): ...@@ -358,34 +364,41 @@ 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 all files # loop on all files, rank by rank
for path_file in track(paths, description="Rearranging..."): for rank in track(ranks, description="Rearranging..."):
with h5py.File(path_file, "r") as file: paths_rank = [
# check if the file contains the time we're looking for p for p in paths if p.name.startswith(f"rank{rank:05}")
tmin_file = file["times"][0] ]
tmax_file = file["times"][-1] for path_file in paths_rank:
if (tmin_file > tmax) or (tmax_file < tmin): # for a given rank, paths are sorted by time
continue with h5py.File(path_file, "r") as file:
# check if the file contains the time we're looking for
# time indices tmax_file = file["times"][-1]
times_file = file["times"][:] if tmax_file < tmin:
cond_file = (times_file >= tmin) & (times_file <= tmax) continue
its_file = np.where(cond_file)[0] tmin_file = file["times"][0]
its = np.where(times == times_file[cond_file])[0] if tmin_file > tmax:
break
# k_adim_loc = global probes indices!
ik0 = file["probes_k0adim_loc"][:] # time indices
ik1 = file["probes_k1adim_loc"][:] times_file = file["times"][:]
ik2 = file["probes_k2adim_loc"][:] cond_file = (times_file >= tmin) & (times_file <= tmax)
its_file = np.where(cond_file)[0]
# load data at desired times for all keys_fields its = np.where(times == times_file[cond_file])[0]
for key in self.keys_fields:
skey = f"spect_{key}" # k_adim_loc = global probes indices!
data = file[skey + "_loc"][:, its_file] ik0 = file["probes_k0adim_loc"][:]
for i in range(its.size): ik1 = file["probes_k1adim_loc"][:]
series[skey][ik0, ik1, ik2, its[i]] = data[ ik2 = file["probes_k2adim_loc"][:]
:, i
].transpose() # load data at desired times for all keys_fields
for key in self.keys_fields:
skey = f"spect_{key}"
data = file[skey + "_loc"][:, its_file]
for i in range(its.size):
series[skey][ik0, ik1, ik2, its[i]] = data[
:, i
].transpose()
# 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