Skip to content

Commit

Permalink
sorting of tsteps
Browse files Browse the repository at this point in the history
  • Loading branch information
haykh committed Jun 10, 2024
1 parent e2f5e40 commit c248359
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 21 deletions.
Binary file added dist/nt2py-0.4.1.tar.gz
Binary file not shown.
2 changes: 1 addition & 1 deletion nt2/__init__.py
Original file line number Diff line number Diff line change
@@ -1 +1 @@
__version__ = "0.4.0"
__version__ = "0.4.1"
61 changes: 41 additions & 20 deletions nt2/read.py
Original file line number Diff line number Diff line change
Expand Up @@ -630,11 +630,13 @@ def __init__(self, fname):
else:
self.metric = MinkowskiMetric()
coords = list(CoordinateDict[coordinates].values())[::-1][-dimension:]

for s in self.file.keys():
if any([k.startswith("X") for k in self.file[s].keys()]):
# cell-centered coords
cc_coords = {c: self.file[s][f"X{i+1}"] for i, c in enumerate(coords[::-1])}
cc_coords = {
c: self.file[s][f"X{i+1}"] for i, c in enumerate(coords[::-1])
}
# cell edges
cell_1 = {
f"{c}_1": (
Expand Down Expand Up @@ -664,7 +666,7 @@ def __init__(self, fname):
)

self.dataset = xr.Dataset()

# -------------------------------- load fields ------------------------------- #
fields = None
f_outsteps = []
Expand All @@ -678,6 +680,10 @@ def __init__(self, fname):
f_times.append(self.file[s]["Time"][()])
f_steps.append(self.file[s]["Step"][()])

f_outsteps = sorted(f_outsteps, key=lambda x: int(x.replace("Step", "")))
f_steps = sorted(f_steps)
f_times = np.array(sorted(f_times), dtype=np.float64)

for k in self.file.attrs.keys():
if (
type(self.file.attrs[k]) == bytes
Expand Down Expand Up @@ -722,7 +728,7 @@ def __init__(self, fname):
},
)
self.dataset[k_] = x

# ------------------------------ load particles ------------------------------ #
particles = None
p_outsteps = []
Expand All @@ -735,19 +741,27 @@ def __init__(self, fname):
p_outsteps.append(s)
p_times.append(self.file[s]["Time"][()])
p_steps.append(self.file[s]["Step"][()])


p_outsteps = sorted(p_outsteps, key=lambda x: int(x.replace("Step", "")))
p_steps = sorted(p_steps)
p_times = np.array(sorted(p_times), dtype=np.float64)

self._particles = {}


if len(p_outsteps) > 0:
species = np.unique(
[int(pq.split("_")[1]) for pq in self.file[p_outsteps[0]].keys() if pq.startswith("p")]
[
int(pq.split("_")[1])
for pq in self.file[p_outsteps[0]].keys()
if pq.startswith("p")
]
)

def list_to_ragged(arr):
max_len = np.max([len(a) for a in arr])
return map(
lambda a: np.concatenate([a, np.full(max_len - len(a), np.nan)]), arr
lambda a: np.concatenate([a, np.full(max_len - len(a), np.nan)]),
arr,
)

for s in species:
Expand All @@ -773,11 +787,16 @@ def list_to_ragged(arr):
if "p" + q in self.file[step_k].keys():
prtl_data[q_].append(self.file[step_k]["p" + q])
else:
prtl_data[q_].append(np.full_like(prtl_data[q_][-1], np.nan))
prtl_data[q_].append(
np.full_like(prtl_data[q_][-1], np.nan)
)
prtl_data[q_] = list_to_ragged(prtl_data[q_])
prtl_data[q_] = da.from_array(list(prtl_data[q_]))
prtl_data[q_] = xr.DataArray(
prtl_data[q_], dims=["t", "id"], name=q_, coords={"t": p_times, "s": ("t", p_steps)}
prtl_data[q_],
dims=["t", "id"],
name=q_,
coords={"t": p_times, "s": ("t", p_steps)},
)
if coordinates == "sph":
prtl_data["x"] = (
Expand All @@ -794,7 +813,7 @@ def list_to_ragged(arr):
prtl_data[PrtlDict[coordinates]["X2"]]
)
self._particles[s] = xr.Dataset(prtl_data)

# ------------------------------- load spectra ------------------------------- #
spectra = None
s_outsteps = []
Expand All @@ -807,13 +826,21 @@ def list_to_ragged(arr):
s_outsteps.append(s)
s_times.append(self.file[s]["Time"][()])
s_steps.append(self.file[s]["Step"][()])


s_outsteps = sorted(s_outsteps, key=lambda x: int(x.replace("Step", "")))
s_steps = sorted(s_steps)
s_times = np.array(sorted(s_times), dtype=np.float64)

self._spectra = xr.Dataset()
log_bins = self.file.attrs["output.spectra.log_bins"]

if len(s_outsteps) > 0:
species = np.unique(
[int(pq.split("_")[1]) for pq in self.file[s_outsteps[0]].keys() if pq.startswith("sN")]
[
int(pq.split("_")[1])
for pq in self.file[s_outsteps[0]].keys()
if pq.startswith("sN")
]
)
e_bins = self.file[s_outsteps[0]]["sEbn"]
if log_bins:
Expand All @@ -839,8 +866,6 @@ def list_to_ragged(arr):
)
self._spectra[f"n_{sp}"] = x



def __del__(self):
self.file.close()

Expand All @@ -855,15 +880,11 @@ def __enter__(self):

def __exit__(self, exc_type, exc_value, traceback):
self.file.close()
self.close()
for _, v in self._particles.items():
del v
del self

@property
def particles(self):
return self._particles

@property
def spectra(self):
return self._spectra
Expand Down

0 comments on commit c248359

Please sign in to comment.