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,14 +364,21 @@ class SpatioTemporalSpectra(SpecificOutput): ...@@ -358,14 +364,21 @@ 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..."):
paths_rank = [
p for p in paths if p.name.startswith(f"rank{rank:05}")
]
for path_file in paths_rank:
# for a given rank, paths are sorted by time
with h5py.File(path_file, "r") as file: with h5py.File(path_file, "r") as file:
# check if the file contains the time we're looking for # check if the file contains the time we're looking for
tmin_file = file["times"][0]
tmax_file = file["times"][-1] tmax_file = file["times"][-1]
if (tmin_file > tmax) or (tmax_file < tmin): if tmax_file < tmin:
continue continue
tmin_file = file["times"][0]
if tmin_file > tmax:
break
# time indices # time indices
times_file = file["times"][:] times_file = file["times"][:]
......
Supports Markdown
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