diff --git a/exptool/io/psp_io.py b/exptool/io/psp_io.py index f18535e..430dbec 100644 --- a/exptool/io/psp_io.py +++ b/exptool/io/psp_io.py @@ -30,6 +30,33 @@ except ImportError: raise ImportError("You will need to 'pip install pyyaml' to use this reader.") +def _to_python(obj): + """ + Recursively convert numpy types to native Python types. + + This helper function ensures that objects containing numpy types (such as numpy scalars or arrays) + are converted to their native Python equivalents. This is particularly important for YAML serialization, + which may not handle numpy types correctly. + + Parameters + ---------- + obj : any + Any Python object, potentially containing numpy types (e.g., numpy scalars, arrays, or nested structures). + + Returns + ------- + out : any + The input object with all numpy types converted to native Python types. + """ + if isinstance(obj, dict): + return {k: _to_python(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [_to_python(v) for v in obj] + elif hasattr(obj, 'item'): # catches numpy scalars + return obj.item() + else: + return obj + class Input: """Input class to adaptively handle OUT. format specifically @@ -79,8 +106,7 @@ def __init__(self, filename,comp=None,verbose=0): # do an initial read of the header self.primary_header = dict() - # initialise dictionaries - self.comp_map = dict() + # initialise dictionary for headers self.header = dict() self._read_primary_header() @@ -88,18 +114,68 @@ def __init__(self, filename,comp=None,verbose=0): self.comp = comp _comps = list(self.header.keys()) + # if a component is defined, retrieve data if comp != None: - if comp not in _comps: - raise IOError('The specified component does not exist.') + + # or check if we are reading all components + if comp == 'all': + self.data = dict() + for c in _comps: + self.data[c] = self._read_component_data(self.filename, + c, + self.header[c]['nbodies'], + int(self.header[c]['data_start'])) else: - self.data = self._read_component_data(self.filename, - self.header[self.comp]['nbodies'], - int(self.header[self.comp]['data_start'])) + if comp not in _comps: + raise IOError(f'The specified component, {comp}, does not exist.') + + else: + self.data = self._read_component_data(self.filename, + self.comp, + self.header[self.comp]['nbodies'], + int(self.header[self.comp]['data_start'])) + # wrapup self.f.close() + def write(self, filename): + """ + Write the current data to a PSP/OUT. file. + + Parameters + ---------- + filename : str + The output filename to which the data will be written. + + Behavior + -------- + Writes all components to the specified file if `comp='all'` was used when reading. + Writing of single components is not implemented and will raise an exception. + + Exceptions + ---------- + NotImplementedError + Raised if attempting to write when `comp` is not 'all'. + + Example + ------- + >>> inp = Input("input.OUT", comp="all") + >>> inp.write("output.OUT") + """ + if self.comp != 'all': + raise NotImplementedError("Writing single components is not implemented yet. Use comp='all'.") + + with open(filename, 'wb') as f: + self._write_primary_header(f) + + # Now write all component headers and data sequentially. + for comp in self.header: + self._write_component_header(f, self.header[comp]) + self._write_component_data(f, comp, self.data[comp]) + + def _read_primary_header(self): """read the primary header from an OUT. file""" @@ -117,6 +193,25 @@ def _read_primary_header(self): next_comp = self._read_out_component_header() data_start = next_comp + def _write_primary_header(self, f): + # time is always