Skip to content

Commit

Permalink
Merge pull request #371 from RemiLehe/prepare_v2
Browse files Browse the repository at this point in the history 
Support reading of 'rt' lasy files
  • Loading branch information
@ax3l
ax3l authored Oct 12, 2023
2 parents 88ab6e1 + dbd1cd5 commit b36dc43
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 26 deletions.
29 changes: 14 additions & 15 deletions openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/field_reader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -198,24 +198,23 @@ def read_field_circ( filename, iteration, field, coord,
field_path = join_infile_path( field, coord )
group, dset = find_dataset( dfile, iteration, field_path )

# Current simulation time
time = (it.attrs['time'] + group.attrs['timeOffset']) * it.attrs['timeUnitSI']

# Extract the metainformation
coord_labels = {ii: coord.decode() for (ii,coord) in
enumerate(group.attrs['axisLabels'])}
coord_label_str = ''.join(coord_labels.values())
coord_label_str = 'm' + coord_label_str
coord_order = RZorder[coord_label_str]

# Current simulation time
time = (it.attrs['time'] + group.attrs['timeOffset']) * it.attrs['timeUnitSI']

if coord_order == RZorder.mrz:
if coord_labels[0] == 'r':
coord_order = RZorder.mrz
Nm, Nr, Nz = get_shape( dset )
N_pair = (Nr, Nz)
elif coord_order == RZorder.mzr:
elif coord_labels[1] == 'r':
Nm, Nz, Nr = get_shape( dset )
N_pair = (Nz, Nr)
coord_order = RZorder.mzr
else:
raise Exception(order_error_msg)

info = FieldMetaInformation( coord_labels, N_pair,
group.attrs['gridSpacing'], group.attrs['gridGlobalOffset'],
group.attrs['gridUnitSI'], dset.attrs['position'], time,
Expand Down Expand Up @@ -287,12 +286,6 @@ def read_field_circ( filename, iteration, field, coord,
else:

# Extract the modes and recombine them properly
if coord_order is RZorder.mrz:
F_total = np.zeros( (2 * Nr, Nz ), dtype=dset.dtype )
elif coord_order is RZorder.mzr:
F_total = np.zeros( (Nz, 2 * Nr ), dtype=dset.dtype )
else:
raise Exception(order_error_msg)
if m == 'all':
# Sum of all the modes
# - Prepare the multiplier arrays
Expand All @@ -308,11 +301,13 @@ def read_field_circ( filename, iteration, field, coord,
# - Sum the modes
F = get_data( dset ) # (Extracts all modes)
if coord_order is RZorder.mrz:
F_total = np.zeros( (2 * Nr, Nz ), dtype=F.dtype )
F_total[Nr:, :] = np.tensordot( mult_above_axis,
F, axes=(0, 0) )[:, :]
F_total[:Nr, :] = np.tensordot( mult_below_axis,
F, axes=(0, 0) )[::-1, :]
elif coord_order is RZorder.mzr:
F_total = np.zeros( (Nz, 2 * Nr ), dtype=F.dtype )
F_total[:, Nr:] = np.tensordot( mult_above_axis,
F, axes=(0, 0) )[:, :]
F_total[:, :Nr] = np.tensordot( mult_below_axis,
Expand All @@ -321,9 +316,11 @@ def read_field_circ( filename, iteration, field, coord,
# Extract mode 0
F = get_data( dset, 0, 0 )
if coord_order is RZorder.mrz:
F_total = np.zeros( (2 * Nr, Nz ), dtype=F.dtype )
F_total[Nr:, :] = F[:, :]
F_total[:Nr, :] = F[::-1, :]
elif coord_order is RZorder.mzr:
F_total = np.zeros( (Nz, 2 * Nr ), dtype=F.dtype )
F_total[:, Nr:] = F[:, :]
F_total[:, :Nr] = F[:, ::-1]
else:
Expand All @@ -337,9 +334,11 @@ def read_field_circ( filename, iteration, field, coord,
F_sin = get_data( dset, 2 * m, 0 )
F = cos * F_cos + sin * F_sin
if coord_order is RZorder.mrz:
F_total = np.zeros( (2 * Nr, Nz ), dtype=F.dtype )
F_total[Nr:, :] = F[:, :]
F_total[:Nr, :] = (-1) ** m * F[::-1, :]
elif coord_order is RZorder.mzr:
F_total = np.zeros( (Nz, 2 * Nr ), dtype=F.dtype )
F_total[:, Nr:] = F[:, :]
F_total[:, :Nr] = (-1) ** m * F[:, ::-1]
else:
Expand Down
22 changes: 11 additions & 11 deletions openpmd_viewer/openpmd_timeseries/data_reader/io_reader/field_reader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -203,15 +203,15 @@ def read_field_circ( series, iteration, field_name, component_name,
# grid spacing/offset/position

coord_labels = {ii: coord for (ii, coord) in enumerate(field.axis_labels)}
coord_label_str = ''.join(coord_labels.values())
coord_label_str = 'm' + coord_label_str
coord_order = RZorder[coord_label_str]
if coord_order is RZorder.mrz:

if coord_labels[0] == 'r':
coord_order = RZorder.mrz
Nm, Nr, Nz = component.shape
N_pair = (Nr, Nz)
elif coord_order is RZorder.mzr:
elif coord_labels[1] == 'r':
Nm, Nz, Nr = component.shape
N_pair = (Nz, Nr)
coord_order = RZorder.mzr
else:
raise Exception(order_error_msg)
time = (it.time + field.time_offset) * it.time_unit_SI
Expand Down Expand Up @@ -285,12 +285,6 @@ def read_field_circ( series, iteration, field_name, component_name,
else:

# Extract the modes and recombine them properly
if coord_order is RZorder.mrz:
F_total = np.zeros( (2 * Nr, Nz ), dtype=component.dtype )
elif coord_order is RZorder.mzr:
F_total = np.zeros( (Nz, 2 * Nr ), dtype=component.dtype )
else:
raise Exception(order_error_msg)
if m == 'all':
# Sum of all the modes
# - Prepare the multiplier arrays
Expand All @@ -306,11 +300,13 @@ def read_field_circ( series, iteration, field_name, component_name,
# - Sum the modes
F = get_data( series, component ) # (Extracts all modes)
if coord_order is RZorder.mrz:
F_total = np.zeros( (2 * Nr, Nz ), dtype=F.dtype )
F_total[Nr:, :] = np.tensordot( mult_above_axis,
F, axes=(0, 0) )[:, :]
F_total[:Nr, :] = np.tensordot( mult_below_axis,
F, axes=(0, 0) )[::-1, :]
elif coord_order is RZorder.mzr:
F_total = np.zeros( (Nz, 2 * Nr ), dtype=F.dtype )
F_total[:, Nr:] = np.tensordot( mult_above_axis,
F, axes=(0, 0) )[:, :]
F_total[:, :Nr] = np.tensordot( mult_below_axis,
Expand All @@ -321,9 +317,11 @@ def read_field_circ( series, iteration, field_name, component_name,
# Extract mode 0
F = get_data( series, component, 0, 0 )
if coord_order is RZorder.mrz:
F_total = np.zeros( (2 * Nr, Nz ), dtype=F.dtype )
F_total[Nr:, :] = F[:, :]
F_total[:Nr, :] = F[::-1, :]
elif coord_order is RZorder.mzr:
F_total = np.zeros( (Nz, 2 * Nr ), dtype=F.dtype )
F_total[:, Nr:] = F[:, :]
F_total[:, :Nr] = F[:, ::-1]
else:
Expand All @@ -336,9 +334,11 @@ def read_field_circ( series, iteration, field_name, component_name,
F_sin = get_data( series, component, 2 * m, 0 )
F = cos * F_cos + sin * F_sin
if coord_order is RZorder.mrz:
F_total = np.zeros( (2 * Nr, Nz ), dtype=F.dtype )
F_total[Nr:, :] = F[:, :]
F_total[:Nr, :] = (-1) ** m * F[::-1, :]
elif coord_order is RZorder.mzr:
F_total = np.zeros( (Nz, 2 * Nr ), dtype=F.dtype )
F_total[:, Nr:] = F[:, :]
F_total[:, :Nr] = (-1) ** m * F[:, ::-1]
else:
Expand Down

0 comments on commit b36dc43

Please sign in to comment.