diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index b72542aaa..80e1527ec 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -240,45 +240,6 @@ def __init__(self, self.verbose = verbose self.restart = restart - self.set_distance_range(rmin=rmin, rmax=rmax, verbose = False) - - self.set_unit_system(MU=MU, DU=DU, TU=TU, - MU2KG=MU2KG, DU2M=DU2M, TU2S=TU2S, - MU_name=MU_name, DU_name=DU_name, TU_name=TU_name, - recompute_values=True, - verbose = False) - - self.set_init_cond_files(init_cond_file_type=init_cond_file_type, - init_cond_file_name=init_cond_file_name, - init_cond_format=init_cond_format, - verbose = False) - - self.set_output_files(output_file_type=output_file_type, - output_file_name=output_file_name, - output_format=output_format, - verbose = False) - - self.set_feature(close_encounter_check=close_encounter_check, - general_relativity=general_relativity, - fragmentation=fragmentation, - rotation=rotation, - compute_conservation_values=compute_conservation_values, - extra_force=extra_force, - big_discard=big_discard, - rhill_present=rhill_present, - restart=restart, - verbose = False) - - self.set_simulation_time(t0=t0, - tstart=tstart, - tstop=tstop, - dt=dt, - tstep_out=tstep_out, - istep_out=istep_out, - istep_dump=istep_dump, - verbose = False - ) - # If the parameter file is in a different location than the current working directory, we will need # to use it to properly open bin files self.sim_dir = os.path.dirname(os.path.realpath(param_file)) @@ -288,6 +249,37 @@ def __init__(self, else: print(f"{param_file} not found.") + self.set_parameters(t0=t0, + tstart=tstart, + tstop=tstop, + dt=dt, + tstep_out=tstep_out, + istep_out=istep_out, + istep_dump=istep_dump, + rmin=rmin, rmax=rmax, + MU=MU, DU=DU, TU=TU, + MU2KG=MU2KG, DU2M=DU2M, TU2S=TU2S, + MU_name=MU_name, DU_name=DU_name, TU_name=TU_name, + recompute_unit_values=False, + init_cond_file_type=init_cond_file_type, + init_cond_file_name=init_cond_file_name, + init_cond_format=init_cond_format, + output_file_type=output_file_type, + output_file_name=output_file_name, + output_format=output_format, + close_encounter_check=close_encounter_check, + general_relativity=general_relativity, + fragmentation=fragmentation, + rotation=rotation, + compute_conservation_values=compute_conservation_values, + extra_force=extra_force, + big_discard=big_discard, + rhill_present=rhill_present, + restart=restart, + verbose = False) + + + if read_old_output_file: binpath = os.path.join(self.sim_dir,self.param['BIN_OUT']) if os.path.exists(binpath): @@ -296,6 +288,48 @@ def __init__(self, print(f"BIN_OUT file {binpath} not found.") return + def _get_valid_arg_list(self, arg_list: str | List[str] | None = None, valid_var: Dict | None = None): + """ + Internal function for getters that extracts subset of arguments that is contained in the dictionary of valid + argument/parameter variable pairs. + + Parameters + ---------- + arg_list : str | List[str], optional + A single string or list of strings containing the Simulation argument. If none are supplied, + then it will create the arg_list out of all keys in the valid_var dictionary. + valid_var : valid_var: Dict + A dictionary where the key is the argument name and the value is the equivalent key in the Simulation + parameter dictionary (i.e. the left-hand column of a param.in file) + + Returns + ------- + valid_arg : [str] + The list of valid arguments that match the keys in valid_var + param : dict + A dictionary that is the subset of the Simulation parameter dictionary corresponding to the arguments listed + in arg_list. + + """ + + if arg_list is None: + valid_arg = None + else: + valid_arg = arg_list.copy() + + if valid_arg is None: + valid_arg = list(valid_var.keys()) + elif type(arg_list) is str: + valid_arg = [arg_list] + else: + # Only allow arg_lists to be checked if they are valid. Otherwise ignore. + valid_arg = [k for k in arg_list if k in list(valid_var.keys())] + + # Extract the arg_list dictionary + param = {valid_var[arg]: self.param[valid_var[arg]] for arg in valid_arg if valid_var[arg] in self.param} + + return valid_arg, param + def set_simulation_time(self, t0: float | None = None, tstart: float | None = None, @@ -335,7 +369,8 @@ def set_simulation_time(self, Returns ------- - Sets the appropriate parameter variables in the parameter dictionary. + time_dict : dict + A dictionary containing the requested parameters """ @@ -409,16 +444,11 @@ def set_simulation_time(self, self.param['ISTEP_DUMP'] = istep_dump update_list.append("istep_dump") + time_dict = self.get_simulation_time(update_list,verbose=verbose) - if verbose is None: - verbose = self.verbose - if verbose: - if len(update_list) > 0: - time_dict = self.get_simulation_time(update_list) - - return + return time_dict - def get_simulation_time(self, arg_list: str | List[str] | None = None): + def get_simulation_time(self, arg_list: str | List[str] | None = None, verbose: bool | None = None): """ Returns a subset of the parameter dictionary containing the current simulation time parameters. @@ -430,6 +460,8 @@ def get_simulation_time(self, arg_list: str | List[str] | None = None): A single string or list of strings containing the names of the simulation time parameters to extract. Default is all of: ["t0", "tstart", "tstop", "dt", "istep_out", "tstep_out", "istep_dump"] + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag Returns ------- @@ -463,7 +495,10 @@ def get_simulation_time(self, arg_list: str | List[str] | None = None): valid_arg, time_dict = self._get_valid_arg_list(arg_list, valid_var) - if self.verbose: + if verbose is None: + verbose = self.verbose + + if verbose: for arg in valid_arg: key = valid_var[arg] if key in time_dict: @@ -506,6 +541,8 @@ def set_feature(self, rhill_present: bool | None = None, restart: bool | None = None, tides: bool | None = None, + interaction_loops: Literal["TRIANGULAR", "FLAT", "ADAPTIVE"] | None = None, + encounter_check_loops: Literal["TRIANGULAR", "SORTSWEEP", "ADAPTIVE"] | None = None, verbose: bool | None = None, **kwargs: Any ): @@ -535,6 +572,28 @@ def set_feature(self, Includes big bodies when performing a discard (Swifter only) rhill_present: bool, optional Include the Hill's radius with the input files. + interaction_loops : {"TRIANGULAR","FLAT","ADAPTIVE"}, default "TRIANGULAR" + *Swiftest Experimental feature* + Specifies which algorithm to use for the computation of body-body gravitational forces. + * "TRIANGULAR" : Upper-triangular double-loops . + * "FLAT" : Body-body interation pairs are flattened into a 1-D array. + * "ADAPTIVE" : Periodically times the TRIANGULAR and FLAT methods and determines which one to use based on + the wall time to complete the loop. *Note:* Using ADAPTIVE means that bit-identical repeatability cannot + be assured, as the choice of algorithm depends on possible external factors that can influence the wall + time calculation. The exact floating-point results of the interaction will be different between the two + algorithm types. + encounter_check_loops : {"TRIANGULAR","SORTSWEEP","ADAPTIVE"}, default "TRIANGULAR" + *Swiftest Experimental feature* + Specifies which algorithm to use for checking whether bodies are in a close encounter state or not. + * "TRIANGULAR" : Upper-triangular double-loops. + * "SORTSWEEP" : A Sort-Sweep algorithm is used to reduce the population of potential close encounter bodies. + This algorithm is still in development, and does not necessarily speed up the encounter checking. + Use with caution. + * "ADAPTIVE" : Periodically times the TRIANGULAR and SORTSWEEP methods and determines which one to use based + on the wall time to complete the loop. *Note:* Using ADAPTIVE means that bit-identical repeatability cannot + be assured, as the choice of algorithm depends on possible external factors that can influence the wall + time calculation. The exact floating-point results of the interaction will be different between the two + algorithm types. tides: bool, optional Turns on tidal model (IN DEVELOPMENT - IGNORED) Yarkovsky: bool, optional @@ -553,7 +612,8 @@ def set_feature(self, Returns ------- - Sets the appropriate parameters in the self.param dictionary + feature_dict : dict + A dictionary containing the requested features. """ @@ -597,17 +657,31 @@ def set_feature(self, self.param["RESTART"] = restart update_list.append("restart") + if interaction_loops is not None: + valid_vals = ["TRIANGULAR","FLAT","ADAPTIVE"] + if interaction_loops not in valid_vals: + print(f"{interaction_loops} is not a valid option for interaction loops.") + print(f"Must be one of {valid_vals}") + self.param["INTERACTION_LOOPS"] = interaction_loops + else: + update_list.append("interaction_loops") + + if encounter_check_loops is not None: + valid_vals = ["TRIANGULAR", "SORTSWEEP", "ADAPTIVE"] + if encounter_check_loops not in valid_vals: + print(f"{encounter_check_loops} is not a valid option for interaction loops.") + print(f"Must be one of {valid_vals}") + else: + self.param["ENCOUNTER_CHECK"] = encounter_check_loops + update_list.append("encounter_check_loops") + self.param["TIDES"] = False - if verbose is None: - verbose = self.verbose - if verbose: - if len(update_list) > 0: - feature_dict = self.get_feature(update_list) - return + feature_dict = self.get_feature(update_list, verbose) + return feature_dict - def get_feature(self, arg_list: str | List[str] | None = None): + def get_feature(self, arg_list: str | List[str] | None = None, verbose: bool | None = None): """ Returns a subset of the parameter dictionary containing the current value of the feature boolean values. @@ -618,6 +692,8 @@ def get_feature(self, arg_list: str | List[str] | None = None): arg_list: str | List[str], optional A single string or list of strings containing the names of the features to extract. Default is all of: ["close_encounter_check", "general_relativity", "fragmentation", "rotation", "compute_conservation_values"] + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag Returns ------- @@ -634,12 +710,17 @@ def get_feature(self, arg_list: str | List[str] | None = None): "extra_force": "EXTRA_FORCE", "big_discard": "BIG_DISCARD", "rhill_present": "RHILL_PRESENT", - "restart" : "RESTART" + "restart" : "RESTART", + "interaction_loops" : "INTERACTION_LOOPS", + "encounter_check_loops" : "ENCOUNTER_CHECK_LOOPS" } valid_arg, feature_dict = self._get_valid_arg_list(arg_list, valid_var) - if self.verbose: + if verbose is None: + verbose = self.verbose + + if verbose: for arg in valid_arg: key = valid_var[arg] print(f"{arg:<32} {feature_dict[key]}") @@ -688,7 +769,8 @@ def set_init_cond_files(self, Returns ------- - Sets the appropriate initial conditions file parameters inside the self.param dictionary. + init_cond_file_dict : dict + A dictionary containing the requested parameters """ @@ -774,16 +856,12 @@ def ascii_file_input_error_msg(codename): self.param["NC_IN"] = init_cond_file_name - if verbose is None: - verbose = self.verbose - if verbose: - if len(update_list) > 0: - init_cond_file_dict = self.get_init_cond_files(update_list) + init_cond_file_dict = self.get_init_cond_files(update_list,verbose) - return + return init_cond_file_dict - def get_init_cond_files(self, arg_list: str | List[str] | None = None): + def get_init_cond_files(self, arg_list: str | List[str] | None = None, verbose: bool | None = None): """ Returns a subset of the parameter dictionary containing the current initial condition file parameters @@ -795,6 +873,8 @@ def get_init_cond_files(self, arg_list: str | List[str] | None = None): A single string or list of strings containing the names of the simulation time parameters to extract. Default is all of: ["init_cond_file_type", "init_cond_file_name", "init_cond_format"] + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag Returns ------- @@ -840,7 +920,10 @@ def get_init_cond_files(self, arg_list: str | List[str] | None = None): valid_arg, init_cond_file_dict = self._get_valid_arg_list(valid_arg, valid_var) - if self.verbose: + if verbose is None: + verbose = self.verbose + + if verbose: for arg in valid_arg: key = valid_var[arg] print(f"{arg:<32} {init_cond_file_dict[key]}") @@ -857,7 +940,7 @@ def set_output_files(self, **kwargs: Any ): """ - Sets the output file parameters in the parameters dictionary. + Sets the output file parameters in the parameter dictionary. Parameters ---------- @@ -881,7 +964,8 @@ def set_output_files(self, Returns ------- - Sets the appropriate initial conditions file parameters inside the self.param dictionary. + output_file_dict : dict + A dictionary containing the requested parameters """ update_list = [] @@ -936,16 +1020,12 @@ def set_output_files(self, else: self.param["OUT_STAT"] = "REPLACE" - if verbose is None: - verbose = self.verbose - if verbose: - if len(update_list) > 0: - output_file_dict = self.get_output_files(update_list) + output_file_dict = self.get_output_files(update_list, verbose=verbose) - return + return output_file_dict - def get_output_files(self, arg_list: str | List[str] | None = None): + def get_output_files(self, arg_list: str | List[str] | None = None, verbose: bool | None = None): """ Returns a subset of the parameter dictionary containing the current output file parameters @@ -957,6 +1037,8 @@ def get_output_files(self, arg_list: str | List[str] | None = None): A single string or list of strings containing the names of the simulation time parameters to extract. Default is all of: ["output_file_type", "output_file_name", "output_format"] + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag Returns ------- @@ -972,7 +1054,10 @@ def get_output_files(self, arg_list: str | List[str] | None = None): valid_arg, output_file_dict = self._get_valid_arg_list(arg_list, valid_var) - if self.verbose: + if verbose is None: + verbose = self.verbose + + if verbose: for arg in valid_arg: key = valid_var[arg] print(f"{arg:<32} {output_file_dict[key]}") @@ -990,7 +1075,7 @@ def set_unit_system(self, MU_name: str | None = None, DU_name: str | None = None, TU_name: str | None = None, - recompute_values: bool = True, + recompute_unit_values: bool = True, verbose: bool | None = None, **kwargs: Any): """ @@ -1043,7 +1128,7 @@ def set_unit_system(self, TU_name : str, optional The name of the time unit. When setting one of the standard units via `TU` a name will be automatically set for the unit, so this argument will override the automatic name. - recompute_values : bool, default True + recompute_unit_values : bool, default True Recompute all values into the new unit system. >*Note:* This is a destructive operation, however if not executed then the values contained in the parameter > file and input/output data files computed previously may not be consistent with the new unit conversion @@ -1056,8 +1141,8 @@ def set_unit_system(self, Returns ---------- - Sets the values of MU2KG, DU2M, and TU2S in the param dictionary to the appropriate units. Also computes the - gravitational constant, self.GU and recomput + unit_dict : dict + A dictionary containing the requested unit conversion parameters """ MU2KG_old = None @@ -1149,18 +1234,14 @@ def set_unit_system(self, self.VU_name = f"{self.DU_name}/{self.TU_name}" self.GU = constants.GC * self.param['TU2S']**2 * self.param['MU2KG'] / self.param['DU2M']**3 - if recompute_values: + if recompute_unit_values: self.update_param_units(MU2KG_old, DU2M_old, TU2S_old) - if verbose is None: - verbose = self.verbose - - if verbose: - unit_dict = self.get_unit_system(update_list) + unit_dict = self.get_unit_system(update_list, verbose) - return + return unit_dict - def get_unit_system(self, arg_list: str | List[str] | None = None): + def get_unit_system(self, arg_list: str | List[str] | None = None, verbose: bool | None = None): """ Returns a subset of the parameter dictionary containing the current simulation unit system. @@ -1175,8 +1256,8 @@ def get_unit_system(self, arg_list: str | List[str] | None = None): Returns ------- - time_dict : dict - A dictionary containing the requested parameters + unit_dict : dict + A dictionary containing the requested unit conversion parameters """ @@ -1212,7 +1293,10 @@ def get_unit_system(self, arg_list: str | List[str] | None = None): valid_arg, unit_dict = self._get_valid_arg_list(arg_list, valid_var) - if self.verbose: + if verbose is None: + verbose = self.verbose + + if verbose: for arg in valid_arg: key = valid_var[arg] print(f"{arg}: {units1[arg]:<28} {unit_dict[key]} {units2[arg]}") @@ -1268,6 +1352,7 @@ def update_param_units(self,MU2KG_old,DU2M_old,TU2S_old): def set_distance_range(self, rmin: float | None = None, rmax: float | None = None, + verbose: bool | None = None, **kwargs: Any): """ Sets the minimum and maximum distances of the simulation. @@ -1284,9 +1369,12 @@ def set_distance_range(self, Returns ------- - Sets the appropriate param dictionary values. + range_dict : dict + A dictionary containing the requested parameters. """ + + update_list = [] CHK_QMIN_RANGE = self.param.pop('CHK_QMIN_RANGE',None) if CHK_QMIN_RANGE is None: CHK_QMIN_RANGE = [-1, -1] @@ -1296,58 +1384,21 @@ def set_distance_range(self, self.param['CHK_QMIN'] = rmin self.param['CHK_RMIN'] = rmin CHK_QMIN_RANGE[0] = rmin + update_list.append("rmin") if rmax is not None: self.param['CHK_RMAX'] = rmax self.param['CHK_EJECT'] = rmax CHK_QMIN_RANGE[1] = rmax + update_list.append("rmax") self.param['CHK_QMIN_RANGE'] =f"{CHK_QMIN_RANGE[0]} {CHK_QMIN_RANGE[1]}" - return - - def _get_valid_arg_list(self, arg_list: str | List[str] | None = None, valid_var: Dict | None = None): - """ - Internal function for getters that extracts subset of arguments that is contained in the dictionary of valid - argument/parameter variable pairs. - - Parameters - ---------- - arg_list : str | List[str], optional - A single string or list of strings containing the Simulation argument. If none are supplied, - then it will create the arg_list out of all keys in the valid_var dictionary. - valid_var : valid_var: Dict - A dictionary where the key is the argument name and the value is the equivalent key in the Simulation - parameter dictionary (i.e. the left-hand column of a param.in file) - - Returns - ------- - valid_arg : [str] - The list of valid arguments that match the keys in valid_var - param : dict - A dictionary that is the subset of the Simulation parameter dictionary corresponding to the arguments listed - in arg_list. - - """ - - if arg_list is None: - valid_arg = None - else: - valid_arg = arg_list.copy() + range_dict = self.get_distance_range(update_list,verbose=verbose) - if valid_arg is None: - valid_arg = list(valid_var.keys()) - elif type(arg_list) is str: - valid_arg = [arg_list] - else: - # Only allow arg_lists to be checked if they are valid. Otherwise ignore. - valid_arg = [k for k in arg_list if k in list(valid_var.keys())] + return range_dict - # Extract the arg_list dictionary - param = {valid_var[arg]:self.param[valid_var[arg]] for arg in valid_arg if valid_var[arg] in self.param} - return valid_arg, param - - def get_distance_range(self, arg_list: str | List[str] | None = None): + def get_distance_range(self, arg_list: str | List[str] | None = None, verbose: bool | None = None): """ Returns a subset of the parameter dictionary containing the current values of the distance range parameters. @@ -1388,7 +1439,10 @@ def get_distance_range(self, arg_list: str | List[str] | None = None): valid_arg, range_dict = self._get_valid_arg_list(arg_list, valid_var) - if self.verbose: + if verbose is None: + verbose = self.verbose + + if verbose: if "rmin" in valid_arg: key = valid_var["rmin"] print(f"{'rmin':<32} {range_dict[key]} {units['rmin']}")