From 1620fdf736b7324174be2e60eee473221937be3c Mon Sep 17 00:00:00 2001 From: David A Minton Date: Mon, 15 May 2023 15:06:51 -0400 Subject: [PATCH 1/2] Cleaned up bugs that were preventing Swifter input file from being read properly --- python/swiftest/swiftest/init_cond.py | 6 +- python/swiftest/swiftest/io.py | 304 ++++--------------- python/swiftest/swiftest/simulation_class.py | 245 +++++++-------- 3 files changed, 179 insertions(+), 376 deletions(-) diff --git a/python/swiftest/swiftest/init_cond.py b/python/swiftest/swiftest/init_cond.py index a4c47548c..8c0227d09 100644 --- a/python/swiftest/swiftest/init_cond.py +++ b/python/swiftest/swiftest/init_cond.py @@ -233,7 +233,7 @@ def vec2xr(param: Dict, **kwargs: Any): Parameters ---------- param : dict - Swiftest paramuration parameters. + Swiftest simulation parameters. name : str or array-like of str, optional Name or names of Bodies. If none passed, name will be "Body" id : int or array-like of int, optional @@ -252,9 +252,9 @@ def vec2xr(param: Dict, **kwargs: Any): capm : float or array-like of float, optional mean anomaly for param['IN_FORM'] == "EL" rh : (n,3) array-like of float, optional - Position vector array. This can be used instead of passing v1, v2, and v3 sepearately for "XV" input format + Position vector array. vh : (n,3) array-like of float, optional - Velocity vector array. This can be used instead of passing v4, v5, and v6 sepearately for "XV" input format + Velocity vector array. Gmass : float or array-like of float, optional G*mass values if these are massive bodies (only one of mass or Gmass can be passed) radius : float or array-like of float, optional diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index f7517e5af..8f4e12078 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -28,6 +28,7 @@ "YARKOVSKY", "YORP", "IN_FORM", + "NC_IN", "SEED", "INTERACTION_LOOPS", "ENCOUNTER_CHECK", @@ -483,6 +484,7 @@ def swifter_stream(f, param): tlab : string list Labels for the tvec data """ + while True: # Loop until you read the end of file try: # Read single-line header @@ -492,32 +494,52 @@ def swifter_stream(f, param): t = record[0] npl = record[1][0] - 1 ntp = record[2][0] - - pvec = np.empty((6, npl)) + + if param['OUT_FORM'] == 'XV': + rpvec = np.empty((npl,3)) + vpvec = np.empty((npl,3)) + rtvec = np.empty((ntp,3)) + vtvec = np.empty((ntp,3)) + elif param['OUT_FORM'] == 'EL': + elempvec = np.empty((npl,6)) + elemtvec = np.empty((ntp,6)) plid = np.empty(npl, dtype='int') - tvec = np.empty((6, ntp)) tpid = np.empty(ntp, dtype='int') if npl > 0: GMpl = np.empty(npl) Rpl = np.empty(npl) for i in range(npl): # Read single-line pl frame for - record = f.read_record(' 0: for i in range(ntp): - record = f.read_record(' 0: - plid = f.read_ints() - dtstr = f'({npl[0]},)a{NAMELEN}' - names = f.read_record(np.dtype(dtstr)) - plnames = [] - for i in range(npl[0]): - plnames.append(np.char.strip(str(names[i], encoding='utf-8'))) - p1 = f.read_reals(np.float64) - p2 = f.read_reals(np.float64) - p3 = f.read_reals(np.float64) - p4 = f.read_reals(np.float64) - p5 = f.read_reals(np.float64) - p6 = f.read_reals(np.float64) - if param['OUT_FORM'] == 'XVEL': - p7 = f.read_reals(np.float64) - p8 = f.read_reals(np.float64) - p9 = f.read_reals(np.float64) - p10 = f.read_reals(np.float64) - p11 = f.read_reals(np.float64) - p12 = f.read_reals(np.float64) - GMpl = f.read_reals(np.float64) - if param['RHILL_PRESENT']: - rhill = f.read_reals(np.float64) - Rpl = f.read_reals(np.float64) - if param['ROTATION']: - Ipplx = f.read_reals(np.float64) - Ipply = f.read_reals(np.float64) - Ipplz = f.read_reals(np.float64) - rotplx = f.read_reals(np.float64) - rotply = f.read_reals(np.float64) - rotplz = f.read_reals(np.float64) - if param['TIDES']: - k2pl = f.read_reals(np.float64) - Qpl = f.read_reals(np.float64) - if ntp[0] > 0: - tpid = f.read_ints() - dtstr = f'({ntp[0]},)a{NAMELEN}' - names = f.read_record(np.dtype(dtstr)) - tpnames = [] - for i in range(npl[0]): - tpnames.append(np.char.strip(str(names[i], encoding='utf-8'))) - t1 = f.read_reals(np.float64) - t2 = f.read_reals(np.float64) - t3 = f.read_reals(np.float64) - t4 = f.read_reals(np.float64) - t5 = f.read_reals(np.float64) - t6 = f.read_reals(np.float64) - if param['OUT_FORM'] == 'XVEL': - t7 = f.read_reals(np.float64) - t8 = f.read_reals(np.float64) - t9 = f.read_reals(np.float64) - t10 = f.read_reals(np.float64) - t11 = f.read_reals(np.float64) - t12 = f.read_reals(np.float64) - clab, plab, tlab, infolab_float, infolab_int, infolab_str = make_swiftest_labels(param) - - if npl > 0: - pvec = np.vstack([p1, p2, p3, p4, p5, p6]) - if param['OUT_FORM'] == 'XVEL': - pvec = np.vstack([pvec, p7, p8, p9, p10, p11, p12]) - pvec = np.vstack([pvec, GMpl, Rpl]) - else: - if param['OUT_FORM'] == 'XVEL': - pvec = np.empty((14, 0)) - else: - pvec = np.empty((8, 0)) - plid = np.empty(0) - plnames = np.empty(0) - if ntp > 0: - tvec = np.vstack([t1, t2, t3, t4, t5, t6]) - if param['OUT_FORM'] == 'XVEL': - tvec = np.vstack([tvec, t7, t8, t9, t10, t11, t12]) - else: - if param['OUT_FORM'] == 'XVEL': - tvec = np.empty((12, 0)) - else: - tvec = np.empty((6, 0)) - tpid = np.empty(0) - tpnames = np.empty(0) - cvec = np.array([Mcb, Rcb, J2cb, J4cb]) - if param['RHILL_PRESENT']: - if npl > 0: - pvec = np.vstack([pvec, rhill]) - if param['ROTATION']: - cvec = np.vstack([cvec, Ipcbx, Ipcby, Ipcbz, rotcbx, rotcby, rotcbz]) - if npl > 0: - pvec = np.vstack([pvec, Ipplx, Ipply, Ipplz, rotplx, rotply, rotplz]) - if param['TIDES']: - cvec = np.vstack([cvec, k2cb, Qcb]) - if npl > 0: - pvec = np.vstack([pvec, k2pl, Qpl]) - yield t, cbid, cbnames, cvec.T, clab, \ - npl, plid, plnames, pvec.T, plab, \ - ntp, tpid, tpnames, tvec.T, tlab - + if param['OUT_FORM'] == 'XV': + yield t, npl, [plid, rpvec, vpvec, GMpl, Rpl], plab, \ + ntp, [tpid, rtvec, vtvec], tlab + elif param['OUT_FORM'] == 'EL': + yield t, npl, [plid, elempvec, GMpl, Rpl], plab, \ + ntp, [tpid, elemtvec], tlab def swifter2xr(param, verbose=True): """ @@ -783,9 +573,16 @@ def swifter2xr(param, verbose=True): dims = ['time', 'id','vec'] pl = [] tp = [] + ds = [] with FortranFile(param['BIN_OUT'], 'r') as f: - for t, npl, plid, pvec, plab, \ - ntp, tpid, tvec, tlab in swifter_stream(f, param): + for t, npl, pvec, plab, \ + ntp, tvec, tlab in swifter_stream(f, param): + + pvec_args = dict(zip(plab,pvec)) + plds = swiftest.init_cond.vec2xr(param,**pvec_args) + if ntp > 0: + tvec_args = dict(zip(tlab,tvec)) + # Prepare frames by adding an extra axis for the time coordinate plframe = np.expand_dims(pvec, axis=0) tpframe = np.expand_dims(tvec, axis=0) @@ -806,6 +603,7 @@ def swifter2xr(param, verbose=True): tpds = tpda.to_dataset(dim='vec') if verbose: print('\nCreating Dataset') ds = xr.combine_by_coords([plds, tpds]) + #if param['OUT_FORM'] = 'XV': if verbose: print(f"Successfully converted {ds.sizes['time']} output frames.") return ds @@ -1848,9 +1646,9 @@ def swiftest2swifter_param(swiftest_param, J2=0.0, J4=0.0): DU2M = swifter_param.pop("DU2M", 1.0) TU2S = swifter_param.pop("TU2S", 1.0) GR = swifter_param.pop("GR", None) - if GR is not None: - if GR: - swifter_param['C'] = swiftest.einsteinC * np.longdouble(TU2S) / np.longdouble(DU2M) + # if GR is not None: + # if GR: + # swifter_param['C'] = swiftest.einsteinC * np.longdouble(TU2S) / np.longdouble(DU2M) for key in newfeaturelist: tmp = swifter_param.pop(key, None) if "ISTEP_DUMP" not in swifter_param: diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 59f4e6669..2ff9c8454 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -345,7 +345,6 @@ def __init__(self,read_param: bool = False, # Parameters are set in reverse priority order. First the defaults, then values from a pre-existing input file, # then using the arguments passed via **kwargs. - #-------------------------- # Lowest Priority: Defaults: #-------------------------- @@ -371,7 +370,7 @@ def __init__(self,read_param: bool = False, self.read_encounters = read_encounters if read_param or read_data: - if self.read_param(read_init_cond = True, dask=dask): + if self.read_param(read_init_cond = True, dask=dask, **kwargs): # We will add the parameter file to the kwarg list. This will keep the set_parameter method from # overriding everything with defaults when there are no arguments passed to Simulation() kwargs['param_file'] = self.param_file @@ -851,8 +850,8 @@ def set_parameter(self, verbose: bool = True, **kwargs): # Setters returning parameter dictionary values param_dict = {} - param_dict.update(self.set_unit_system(**kwargs)) param_dict.update(self.set_integrator(**kwargs)) + param_dict.update(self.set_unit_system(**kwargs)) param_dict.update(self.set_simulation_time(**kwargs)) param_dict.update(self.set_init_cond_files(**kwargs)) param_dict.update(self.set_output_files(**kwargs)) @@ -955,6 +954,11 @@ def set_integrator(self, self.binary_path = "NOT IMPLEMENTED FOR THIS CODE" self.driver_executable = None update_list.append("driver_executable") + + if self.codename == "Swifter": + J2 = self.param.pop("J2",0.0) + J4 = self.param.pop("J4",0.0) + self.param = io.swiftest2swifter_param(self.param, J2, J4) if integrator is not None: valid_integrator = ["symba","rmvs","whm","helio"] @@ -1172,113 +1176,21 @@ def set_feature(self, if close_encounter_check is not None: self.param["CHK_CLOSE"] = close_encounter_check update_list.append("close_encounter_check") - - if general_relativity is not None: - self.param["GR"] = general_relativity - update_list.append("general_relativity") - - fragmentation_models = ["FRAGGLE"] - if collision_model is not None: - collision_model = collision_model.upper() - fragmentation = collision_model in fragmentation_models - if self.codename != "Swiftest" and self.integrator != "symba" and fragmentation: - warnings.warn("Fragmentation is only available on Swiftest SyMBA.",stacklevel=2) - self.param['COLLISION_MODEL'] = "MERGE" - else: - self.param['COLLISION_MODEL'] = collision_model - update_list.append("collision_model") - if fragmentation: - if "MIN_GMFRAG" not in self.param and minimum_fragment_mass is None and minimum_fragment_gmass is None: - warnings.warn("Minimum fragment mass is not set. Set it using minimum_fragment_gmass or minimum_fragment_mass",stacklevel=2) - else: - update_list.append("minimum_fragment_gmass") - - if minimum_fragment_gmass is not None and minimum_fragment_mass is not None: - warnings.warn("Only set either minimum_fragment_mass or minimum_fragment_gmass, but not both!",stacklevel=2) - - if minimum_fragment_gmass is not None: - self.param["MIN_GMFRAG"] = minimum_fragment_gmass - if "minmum_fragment_gmass" not in update_list: - update_list.append("minimum_fragment_gmass") - elif minimum_fragment_mass is not None: - self.param["MIN_GMFRAG"] = minimum_fragment_mass * self.GU - if "minimum_fragment_gmass" not in update_list: - update_list.append("minimum_fragment_gmass") - - if nfrag_reduction is not None: - self.param["NFRAG_REDUCTION"] = nfrag_reduction - update_list.append("nfrag_reduction") - - if rotation is not None: - self.param['ROTATION'] = rotation - update_list.append("rotation") - - if self.param['COLLISION_MODEL'] == "FRAGGLE" and not self.param['ROTATION']: - self.param['ROTATION'] = True - update_list.append("rotation") - - if compute_conservation_values is not None: - self.param["ENERGY"] = compute_conservation_values - update_list.append("compute_conservation_values") - + + if rhill_present is not None: + self.param["RHILL_PRESENT"] = rhill_present + update_list.append("rhill_present") + if extra_force is not None: self.param["EXTRA_FORCE"] = extra_force update_list.append("extra_force") - if big_discard is not None: if self.codename != "Swifter": self.param["BIG_DISCARD"] = False else: self.param["BIG_DISCARD"] = big_discard - update_list.append("big_discard") - - if rhill_present is not None: - self.param["RHILL_PRESENT"] = rhill_present - update_list.append("rhill_present") - - if restart is not None: - self.param["RESTART"] = restart - update_list.append("restart") - - if interaction_loops is not None: - valid_vals = ["TRIANGULAR", "FLAT"] - interaction_loops = interaction_loops.upper() - if interaction_loops not in valid_vals: - msg = f"{interaction_loops} is not a valid option for interaction loops." - msg += f"\nMust be one of {valid_vals}" - warnings.warn(msg,stacklevel=2) - if "INTERACTION_LOOPS" not in self.param: - self.param["INTERACTION_LOOPS"] = valid_vals[0] - else: - self.param["INTERACTION_LOOPS"] = interaction_loops - update_list.append("interaction_loops") - - if encounter_check_loops is not None: - valid_vals = ["TRIANGULAR", "SORTSWEEP"] - encounter_check_loops = encounter_check_loops.upper() - if encounter_check_loops not in valid_vals: - msg = f"{encounter_check_loops} is not a valid option for interaction loops." - msg += f"\nMust be one of {valid_vals}" - warnings.warn(msg,stacklevel=2) - if "ENCOUNTER_CHECK" not in self.param: - self.param["ENCOUNTER_CHECK"] = valid_vals[0] - else: - self.param["ENCOUNTER_CHECK"] = encounter_check_loops - update_list.append("encounter_check_loops") - - if encounter_save is not None: - valid_vals = ["NONE", "TRAJECTORY", "CLOSEST", "BOTH"] - encounter_save = encounter_save.upper() - if encounter_save not in valid_vals: - msg = f"{encounter_save} is not a valid option for encounter_save." - msg += f"\nMust be one of {valid_vals}" - warnings.warn(msg,stacklevel=2) - if "ENCOUNTER_SAVE" not in self.param: - self.param["ENCOUNTER_SAVE"] = valid_vals[0] - else: - self.param["ENCOUNTER_SAVE"] = encounter_save - update_list.append("encounter_save") + update_list.append("big_discard") if simdir is not None: self.simdir = Path(simdir) @@ -1291,15 +1203,106 @@ def set_feature(self, self.driver_executable = self.binary_path / "swiftest_driver" self.param_file = Path(kwargs.pop("param_file","param.in")) - if coarray is not None: - if self.codename == "Swiftest": - self.param["COARRAY"] = coarray - update_list.append("coarray") - - - self.param["TIDES"] = False + if self.codename == "Swiftest": + if general_relativity is not None: + self.param["GR"] = general_relativity + update_list.append("general_relativity") + + fragmentation_models = ["FRAGGLE"] + if collision_model is not None: + collision_model = collision_model.upper() + fragmentation = collision_model in fragmentation_models + if self.codename != "Swiftest" and self.integrator != "symba" and fragmentation: + warnings.warn("Fragmentation is only available on Swiftest SyMBA.",stacklevel=2) + self.param['COLLISION_MODEL'] = "MERGE" + else: + self.param['COLLISION_MODEL'] = collision_model + update_list.append("collision_model") + if fragmentation: + if "MIN_GMFRAG" not in self.param and minimum_fragment_mass is None and minimum_fragment_gmass is None: + warnings.warn("Minimum fragment mass is not set. Set it using minimum_fragment_gmass or minimum_fragment_mass",stacklevel=2) + else: + update_list.append("minimum_fragment_gmass") + + if minimum_fragment_gmass is not None and minimum_fragment_mass is not None: + warnings.warn("Only set either minimum_fragment_mass or minimum_fragment_gmass, but not both!",stacklevel=2) + + if minimum_fragment_gmass is not None: + self.param["MIN_GMFRAG"] = minimum_fragment_gmass + if "minmum_fragment_gmass" not in update_list: + update_list.append("minimum_fragment_gmass") + elif minimum_fragment_mass is not None: + self.param["MIN_GMFRAG"] = minimum_fragment_mass * self.GU + if "minimum_fragment_gmass" not in update_list: + update_list.append("minimum_fragment_gmass") + + if nfrag_reduction is not None: + self.param["NFRAG_REDUCTION"] = nfrag_reduction + update_list.append("nfrag_reduction") + + if rotation is not None: + self.param['ROTATION'] = rotation + update_list.append("rotation") + + if self.param['COLLISION_MODEL'] == "FRAGGLE" and not self.param['ROTATION']: + self.param['ROTATION'] = True + update_list.append("rotation") + + if compute_conservation_values is not None: + self.param["ENERGY"] = compute_conservation_values + update_list.append("compute_conservation_values") + + if restart is not None: + self.param["RESTART"] = restart + update_list.append("restart") + + if interaction_loops is not None: + valid_vals = ["TRIANGULAR", "FLAT"] + interaction_loops = interaction_loops.upper() + if interaction_loops not in valid_vals: + msg = f"{interaction_loops} is not a valid option for interaction loops." + msg += f"\nMust be one of {valid_vals}" + warnings.warn(msg,stacklevel=2) + if "INTERACTION_LOOPS" not in self.param: + self.param["INTERACTION_LOOPS"] = valid_vals[0] + else: + self.param["INTERACTION_LOOPS"] = interaction_loops + update_list.append("interaction_loops") + + if encounter_check_loops is not None: + valid_vals = ["TRIANGULAR", "SORTSWEEP"] + encounter_check_loops = encounter_check_loops.upper() + if encounter_check_loops not in valid_vals: + msg = f"{encounter_check_loops} is not a valid option for interaction loops." + msg += f"\nMust be one of {valid_vals}" + warnings.warn(msg,stacklevel=2) + if "ENCOUNTER_CHECK" not in self.param: + self.param["ENCOUNTER_CHECK"] = valid_vals[0] + else: + self.param["ENCOUNTER_CHECK"] = encounter_check_loops + update_list.append("encounter_check_loops") + + if encounter_save is not None: + valid_vals = ["NONE", "TRAJECTORY", "CLOSEST", "BOTH"] + encounter_save = encounter_save.upper() + if encounter_save not in valid_vals: + msg = f"{encounter_save} is not a valid option for encounter_save." + msg += f"\nMust be one of {valid_vals}" + warnings.warn(msg,stacklevel=2) + if "ENCOUNTER_SAVE" not in self.param: + self.param["ENCOUNTER_SAVE"] = valid_vals[0] + else: + self.param["ENCOUNTER_SAVE"] = encounter_save + update_list.append("encounter_save") - + if coarray is not None: + if self.codename == "Swiftest": + self.param["COARRAY"] = coarray + update_list.append("coarray") + + self.param["TIDES"] = False + + feature_dict = self.get_feature(update_list, verbose) return feature_dict @@ -1902,11 +1905,12 @@ def set_unit_system(self, if all(key in self.param for key in ["MU2KG","DU2M","TU2S"]): self.GU = constants.GC * self.param["TU2S"] ** 2 * self.param["MU2KG"] / self.param["DU2M"] ** 3 - if recompute_unit_values and \ - MU2KG_old != self.param['MU2KG'] or \ - DU2M_old != self.param['DU2M'] or \ - TU2S_old != self.param['TU2S']: - self.update_param_units(MU2KG_old, DU2M_old, TU2S_old) + if "MU2KG" in self.param and "DU2M" in self.param and "TU2S" in self.param: + if recompute_unit_values and \ + MU2KG_old != self.param['MU2KG'] or \ + DU2M_old != self.param['DU2M'] or \ + TU2S_old != self.param['TU2S']: + self.update_param_units(MU2KG_old, DU2M_old, TU2S_old) unit_dict = self.get_unit_system(update_list, verbose) @@ -2348,14 +2352,14 @@ def _get_instance_var(self, arg_list: str | List[str], valid_arg: Dict, verbose: Parameters ---------- arg_list: str | List[str] - A single string or list of strings containing the names of the the instance variable to get. + A single string or list of strings containing the names of the the instance variable to get. valid_arg: dict - A dictionary where the key is the parameter argument and the value is the equivalent instance variable value. + A dictionary where the key is the parameter argument and the value is the equivalent instance variable value. verbose: bool, optional - If passed, it will override the Simulation object's verbose flag + If passed, it will override the Simulation object's verbose flag **kwargs - A dictionary of additional keyword argument. This allows this method to be called by the more general - get_parameter method, which takes all possible Simulation parameters as arguments, so these are ignored. + A dictionary of additional keyword argument. This allows this method to be called by the more general + get_parameter method, which takes all possible Simulation parameters as arguments, so these are ignored. Returns ------- @@ -2622,7 +2626,8 @@ def read_param(self, codename: Literal["Swiftest", "Swifter", "Swift"] | None = None, read_init_cond : bool | None = None, verbose: bool | None = None, - dask: bool = False): + dask: bool = False, + **kwargs : Any): """ Reads in an input parameter file and stores the values in the param dictionary. From 12bc34ee1425ffb43ccd36ccf65a0c3f5e332a78 Mon Sep 17 00:00:00 2001 From: David A Minton Date: Mon, 15 May 2023 15:25:14 -0400 Subject: [PATCH 2/2] Improved the reading of Swifter files using the vec2xr method --- python/swiftest/swiftest/init_cond.py | 10 +++++---- python/swiftest/swiftest/io.py | 31 +++++++++------------------ 2 files changed, 16 insertions(+), 25 deletions(-) diff --git a/python/swiftest/swiftest/init_cond.py b/python/swiftest/swiftest/init_cond.py index 8c0227d09..efaa6429f 100644 --- a/python/swiftest/swiftest/init_cond.py +++ b/python/swiftest/swiftest/init_cond.py @@ -282,10 +282,12 @@ def vec2xr(param: Dict, **kwargs: Any): # Check for valid keyword arguments kwargs = {k:kwargs[k] for k,v in kwargs.items() if v is not None} - if "rot" not in kwargs and "Gmass" in kwargs: - kwargs['rot'] = np.zeros((len(kwargs['Gmass']),3)) - if "Ip" not in kwargs and "Gmass" in kwargs: - kwargs['Ip'] = np.full((len(kwargs['Gmass']),3), 0.4) + + if "rotation" in param and param['rotation'] == True: + if "rot" not in kwargs and "Gmass" in kwargs: + kwargs['rot'] = np.zeros((len(kwargs['Gmass']),3)) + if "Ip" not in kwargs and "Gmass" in kwargs: + kwargs['Ip'] = np.full((len(kwargs['Gmass']),3), 0.4) if "time" not in kwargs: kwargs["time"] = np.array([0.0]) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 8f4e12078..32a10f34d 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -573,37 +573,26 @@ def swifter2xr(param, verbose=True): dims = ['time', 'id','vec'] pl = [] tp = [] - ds = [] with FortranFile(param['BIN_OUT'], 'r') as f: for t, npl, pvec, plab, \ ntp, tvec, tlab in swifter_stream(f, param): + sys.stdout.write('\r' + f"Reading in time {t[0]:.3e}") + sys.stdout.flush() + pvec_args = dict(zip(plab,pvec)) - plds = swiftest.init_cond.vec2xr(param,**pvec_args) + pl.append(swiftest.init_cond.vec2xr(param,time=t,**pvec_args)) if ntp > 0: tvec_args = dict(zip(tlab,tvec)) + tp.append(swiftest.init_cond.vec2xr(param,time=t,**tvec_args)) - # Prepare frames by adding an extra axis for the time coordinate - plframe = np.expand_dims(pvec, axis=0) - tpframe = np.expand_dims(tvec, axis=0) - - # Create xarray DataArrays out of each body type - plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab}) - tpxr = xr.DataArray(tpframe, dims=dims, coords={'time': t, 'id': tpid, 'vec': tlab}) - - pl.append(plxr) - tp.append(tpxr) - sys.stdout.write('\r' + f"Reading in time {t[0]:.3e}") - sys.stdout.flush() - - plda = xr.concat(pl, dim='time') - tpda = xr.concat(tp, dim='time') + plds = xr.concat(pl, dim='time') + if len(tp) > 0: + tpds = xr.concat(tp, dim='time') - plds = plda.to_dataset(dim='vec') - tpds = tpda.to_dataset(dim='vec') if verbose: print('\nCreating Dataset') - ds = xr.combine_by_coords([plds, tpds]) - #if param['OUT_FORM'] = 'XV': + if len(tp) > 0: + ds = xr.combine_by_coords([plds, tpds]) if verbose: print(f"Successfully converted {ds.sizes['time']} output frames.") return ds