Skip to content

Commit 7faa192

Browse files
AreWeDreamingAreWeDreaming
committed
Fix errors in test data and handle 2d ragged data
1 parent 4228cad commit 7faa192

File tree

3 files changed

+28
-19
lines changed

3 files changed

+28
-19
lines changed

‎imas/backends/netcdf/ids_tensorizer.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,7 @@ def tensorize(self, path, fillvalue):
208208
def recursively_convert_to_list(self, path: str, inactive_index:Tuple,
209209
shape:Tuple, i_dim: int):
210210
entry = []
211-
for index in path:
211+
for index in range(shape[i_dim]):
212212
new_index = inactive_index + (index,)
213213
if i_dim == len(shape) - 1:
214214
entry.append(self.filled_data[path][new_index].value)
@@ -231,12 +231,16 @@ def awkward_tensorize(self, path:str):
231231
"""
232232
if path in self.shapes:
233233
shape = self.shapes[path]
234+
if shape.ndim > 2:
235+
raise NotImplementedError("Dimensions higher than 2 are not yet implemented.")
236+
shape = shape.shape
237+
hdf5_dim = 1
234238
else:
235239
dimensions = self.ncmeta.get_dimensions(path, self.homogeneous_time)
236240
shape = tuple(self.dimension_size[dim] for dim in dimensions)
237-
# Get the split between HDF5 indices and stored matrices
238-
# i.e. equilibrium.time_slice.profiles_2d <-> psi
239-
hdf5_dim = len(list(self.filled_data[path].keys())[0])
241+
# Get the split between HDF5 indices and stored matrices
242+
# i.e. equilibrium.time_slice.profiles_2d <-> psi
243+
hdf5_dim = len(list(self.filled_data[path].keys())[0])
240244
if hdf5_dim == 0:
241245
return self.filled_data[path][()].value
242246
else:

‎imas/test/test_wrangle.py‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -67,23 +67,28 @@ def flat(test_data):
6767
flat["equilibrium.time_slice.profiles_2d.psi"][:] = test_data["equilibrium"]["psi_2d"][None, ...]
6868

6969
# Thomson scattering test data (ragged)
70+
N = test_data["thomson_scattering"]["N_ch"][0] + test_data["thomson_scattering"]["N_ch"][1]
7071
flat["thomson_scattering.ids_properties.homogeneous_time"] = 0
7172
flat["thomson_scattering.channel.t_e.time"] = ak.concatenate([np.tile(test_data["thomson_scattering"]["time"][0],
72-
(test_data["thomson_scattering"]["N_ch"][0],1)),
73+
(test_data["thomson_scattering"]["N_ch"][0],
74+
1)),
7375
np.tile(test_data["thomson_scattering"]["time"][1],
74-
(test_data["thomson_scattering"]["N_ch"][1],1))])
75-
flat["thomson_scattering.channel.t_e.data"] = ak.concatenate([np.tile(test_data["thomson_scattering"]["t_e"][0],
76-
(test_data["thomson_scattering"]["N_ch"][0],1)),
77-
np.tile(test_data["thomson_scattering"]["t_e"][1],
78-
(test_data["thomson_scattering"]["N_ch"][1],1))])
76+
(test_data["thomson_scattering"]["N_ch"][1],
77+
1))])
78+
flat["thomson_scattering.channel.t_e.data"] = ak.concatenate([np.repeat(test_data["thomson_scattering"]["t_e"][:test_data["thomson_scattering"]["N_ch"][0],None],
79+
test_data["thomson_scattering"]["N_time"][0], axis=1),
80+
np.repeat(test_data["thomson_scattering"]["t_e"][test_data["thomson_scattering"]["N_ch"][0]:,None],
81+
test_data["thomson_scattering"]["N_time"][1], axis=1)])
7982
flat["thomson_scattering.channel.n_e.time"] = ak.concatenate([np.tile(test_data["thomson_scattering"]["time"][0],
80-
(test_data["thomson_scattering"]["N_ch"][0],1)),
83+
(test_data["thomson_scattering"]["N_ch"][0],
84+
1)),
8185
np.tile(test_data["thomson_scattering"]["time"][1],
82-
(test_data["thomson_scattering"]["N_ch"][1],1))])
83-
flat["thomson_scattering.channel.n_e.data"] = ak.concatenate([np.tile(test_data["thomson_scattering"]["n_e"][0],
84-
(test_data["thomson_scattering"]["N_ch"][0],1)),
85-
np.tile(test_data["thomson_scattering"]["n_e"][1],
86-
(test_data["thomson_scattering"]["N_ch"][1],1))])
86+
(test_data["thomson_scattering"]["N_ch"][1],
87+
1))])
88+
flat["thomson_scattering.channel.n_e.data"] = ak.concatenate([np.repeat(test_data["thomson_scattering"]["n_e"][:test_data["thomson_scattering"]["N_ch"][0],None],
89+
test_data["thomson_scattering"]["N_time"][0], axis=1),
90+
np.repeat(test_data["thomson_scattering"]["n_e"][test_data["thomson_scattering"]["N_ch"][0]:,None],
91+
test_data["thomson_scattering"]["N_time"][1], axis=1)])
8792
flat["thomson_scattering.channel.position.r"] = test_data["thomson_scattering"]["r"]
8893
flat["thomson_scattering.channel.position.z"] = test_data["thomson_scattering"]["z"]
8994
return flat
@@ -115,7 +120,7 @@ def test_ids_dict(test_data):
115120
thomson_scattering.channel[i].t_e.data = np.tile(test_data["thomson_scattering"]["t_e"][i],
116121
test_data["thomson_scattering"]["N_time"][index])
117122
thomson_scattering.channel[i].n_e.time = test_data["thomson_scattering"]["time"][index]
118-
thomson_scattering.channel[i].n_e.data = np.tile(test_data["thomson_scattering"]["t_e"][i],
123+
thomson_scattering.channel[i].n_e.data = np.tile(test_data["thomson_scattering"]["n_e"][i],
119124
test_data["thomson_scattering"]["N_time"][index])
120125
thomson_scattering.channel[i].position.r = test_data["thomson_scattering"]["r"][i]
121126
thomson_scattering.channel[i].position.z = test_data["thomson_scattering"]["z"][i]
@@ -132,4 +137,4 @@ def test_wrangle(test_ids_dict, flat):
132137
def test_unwrangle(test_ids_dict, flat):
133138
result = unwrangle(list(flat.keys()), test_ids_dict)
134139
for key in flat.keys():
135-
np.testing.assert_allclose(result[key], flat[key])
140+
assert ak.almost_equal(result[key], flat[key])

‎imas/wrangler.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ def recursively_put(location, value, ids):
1616
sub_ids.resize(N)
1717
elif sub_ids.size != N:
1818
raise ValueError(
19-
f"""Inconsistent size across flat entries {location}, {N} (flat) vs. ids {ids.size}!
19+
f"""Inconsistent size across flat entries {location}, {N} (flat) vs. ids {sub_ids.size}!
2020
"""
2121
)
2222
# Need to iterate over indices (e.g. equilibrium.time_slice[:].)

0 commit comments

Comments
 (0)