import numpy as np
import h5py
[docs]def read(fname,req_quants=("x1","x2")):
"""Reads an OSIRIS output data file.
:param fname: The path to the data file
:type fname: str
:param req_quants: The quantities required when reading RAW or TRACKS files, defaults to ("x1","x2")
:type req_quants: tuple of str, optional
:return: An os_data object
:rtype: :class:`os_data`
"""
#retrieve the data type of the file
f=h5py.File(fname,"r")
datatype=f.attrs["TYPE"][0].decode('UTF-8')
f.close()
data=None
#initialize the data object according to dataype
match datatype:
case "grid":
data=grid(fname)
case "particles":
data=raw(fname,req_quants)
case "tracks-2":
data=track(fname,req_quants)
case _:
exit()
return data
[docs]class os_data:
""" A class to represent a generic OSIRIS output data object.
:param datatype: The type of OSIRIS data ('grid', 'particles' our 'tracks')
:type datatype: str
"""
def __init__(self,datatype):
"""Constructs all the necessary attributes for the os_data object.
:param datatype: The type of OSIRIS data ('grid', 'particles' our 'tracks')
:type datatype: str
"""
self._datatype=datatype
@property
def datatype(self):
return self._datatype
[docs]class axis:
""" A class used to represent an axis object
:param nx : The number of points along the axis
:type nx: int
:param ax_arr : The array containing the axis poins
:type ax_arr : :class:`np.array`
:param label: The axis label
:type label : str
"""
[docs] def __init__(self,nx,lims,label):
"""Constructs all the necessary attributes for the os_data object.
:param nx: The number of points along the axis
:type nx: int
:param ax_arr: The array containing the axis poins
:type ax_arr: :class:`np.array`
:param label: The axis label
:type label: str
"""
self._nx=nx
self._label=label
self._ax_arr=np.linspace(lims[0],lims[1],nx)
@property
def nx(self):
return self._nx
@property
def label(self):
return self._label
@property
def ax_arr(self):
return self._ax_arr
[docs]class grid(os_data):
"""A class used to represent a grid data object
:param dims: A formatted string to print out what the animal says
:type dims: int
:param axis: A list of axis objects containg the spatial limits of the grid
:type axis: list of :class:`axis`
:param data: The grid data
:type data: :class:`np.array`
:param label: The label of the quantity in the grid
:type label: str
:param time_s: The timestamp of the grid file
:type time_s: :class:`np.float`
"""
[docs] def __init__(self,fname):
"""Constructs all the necessary attributes for the grid object.
:param fname: The path to the osiris data file
:type fname: str
"""
#initialize the parent class
os_data.__init__(self,"grid")
#initialize the axis list
self._axis=[]
#open the file
f=h5py.File(fname,"r")
#read the dataset
objs=f.keys()
for name in objs:
if isinstance(f[name], h5py.Dataset):
dat=np.array(f[name]) #charge for density, e2 for OSIRIS fields, e3_x2_slice for slices
#loop through the axis and fill the axis array
objs=f["AXIS"].keys()
for i,axis_n in enumerate(objs):
ax1=f["AXIS/"+axis_n]
axis1=np.array(ax1)
#retrieve number of point along axis
np.shape(dat)[-(i+1)]
#retrive axis label
ax1name=ax1.attrs["NAME"][0].decode('UTF-8')+" ["+ax1.attrs["UNITS"][0].decode('UTF-8')+"]"
self._axis.append(axis(np.shape(dat)[-(i+1)],axis1,ax1name))
#retrive data label
dataname=f.attrs["NAME"][0].decode('UTF-8')+" ["+f.attrs["UNITS"][0].decode('UTF-8')+"]"
#retrive data timestamp
time_s=f.attrs["TIME"][0]
#close the files
f.close()
#fill in the attributes
self._dims=len(self._axis)
self._data=dat
self._label=dataname
self._time_s=time_s
@property
def axis(self):
return self._axis
@property
def dims(self):
return self._dims
@property
def data(self):
return self._data
@property
def label(self):
return self._label
@property
def data(self):
return self._data
@property
def time_s(self):
return self._time_s
[docs]class raw(os_data):
"""A class used to represent a particles data object
:param data: a dictionary containing the required quantities
:type data: dict of (str,np.arrays)
:param label: a dictionary containing the labels of required quantities
:type label: dict of (str,str)
:param time_s: the timestamp of the grid file
:type time_s: np.float
"""
[docs] def __init__(self,fname,req_quants):
"""Constructs all the necessary attributes for the raw object.
Loads the required quantities (if available) into a dictionary of np.arrays whose keys are the requeired quantities.
:param fname: The path to the osiris data file
:type fname: str
:param req_quants: A list containing the required quantities
:type req_quants: list of str
:raises KeyError: If a required quantity is not available on the file
"""
# initialize the parent class
os_data.__init__(self,"particles")
#initialize the dictionaries
self._data=dict.fromkeys(req_quants,None)
self._label=dict.fromkeys(req_quants,None)
#open the file
f=h5py.File(fname,"r")
#retrieve the available quantities
quants=[i.decode('UTF-8') for i in f.attrs["QUANTS"]]
#retrieve the corresponding labels
labels=[i.decode('UTF-8') for i in f.attrs["LABELS"]]
#retrieve the corresponding units
units=[i.decode('UTF-8') for i in f.attrs["UNITS"]]
#build the aux dictionaries
labels=dict(zip(quants, labels))
units=dict(zip(quants, units))
#try to build the class dictionaries
try:
for quant in req_quants:
self._data[quant]=np.array(f[quant])
self._label[quant]=labels[quant]+" ["+units[quant]+"]"
except KeyError as ke:
err_str="Available Objects: "+str(quants)
raise KeyError(err_str)
#retrieve data timestamp
self._time_s=f.attrs["TIME"][0]
#close the files
f.close()
@property
def label(self):
return self._label
@property
def data(self):
return self._data
@property
def time_s(self):
return self._time_s
[docs]class track(os_data):
"""A class used to represent a tracks data object
:param data: a dictionary containing the required quantities for each particle
:type data: dict of (str,list of np.arrays)
:param label: a dictionary containing the labels of required quantities
:type label: dict of (str,str)
"""
[docs] def __init__(self,fname,req_quants):
"""Constructs all the necessary attributes for the tracks object.
Reads fname and loads the required quantities (if available) into a dictionary of lists np.arrays whose keys are the requeired quantities.
:param fname: The path to the osiris data file
:type fname: str
:param req_quants: A list containing the required quantities
:type req_quants: list of str
:raises KeyError: If a required quantity is not available on the file
"""
# initialize the parent class
os_data.__init__(self,"tracks-v2")
#initialize the dictionaries
self._data=dict.fromkeys(req_quants,None)
self._label=dict.fromkeys(req_quants,None)
#open the file
f=h5py.File(file,"r")
#retrieve the available quantities, labels and units
quants=[i.decode('UTF-8') for i in f.attrs["QUANTS"]][1:]
labels=[i.decode('UTF-8') for i in f.attrs["LABELS"]][1:]
units=[i.decode('UTF-8') for i in f.attrs["UNITS"]][1:]
#build auxiliary dictionaries
labels=dict(zip(quants, labels))
units=dict(zip(quants, units))
#retrieve the itermap array
itermap=np.array(f["itermap"])
#retrieve the number of tracks in the file
ntracks=f.attrs["NTRACKS"][0]
#modify the itermap array to contain the actual starting index for each particle instead of the relative index
for i in range(len(itermap)):
itermap[i,2]=np.sum(itermap[:i,1])
#initialize the itermaps array
itermaps=[]
#the itermaps array is a list containing the itermap bounds for each particle (actual start and end idx for each particle in the data array)
for i in range(1,ntracks+1):
itermaps.append(itermap[itermap[:,0]==i,1:])
#initialize the itermaps array
itermaps_tracks=[]
#the itermaps_tracks array is a list containing the itermap for each particle (all the idx for each particle in the data array)
for i in range(len(itermaps)):
itermap_=[]
if len(itermaps[i])>1:
for bound in itermaps[i]:
itermap_.append(np.arange(bound[1],bound[1]+bound[0]))
itermaps_tracks.append(np.concatenate(itermap_))
#retrieve the data array
data=np.array(f["data"])
print(f.attrs["QUANTS"])
f.close()
#try to build the class dictionaries
try:
for quant in req_quants:
idx=quants.index(quant)
print(idx)
self._data[quant]=[]
#loop through the particles and retireve the data usinf the arrays in itermaps_tracks
for track_idx in itermaps_tracks:
#append the data for particle x to the list in the dictionary
self._data[quant].append(data[track_idx,idx])
self._label[quant]=labels[quant]+" ["+units[quant]+"]"
except ValueError as ke:
err_str="Available Objects: "+str(quants)
raise ValueError(err_str)
@property
def label(self):
return self._label
@property
def data(self):
return self._data