From 31b05def1ea711a54a027dce86cca745befe634f Mon Sep 17 00:00:00 2001 From: Carlisle Wishard Date: Tue, 11 Oct 2022 13:57:10 -0400 Subject: [PATCH] Added docstring headers to python scripts --- python/swiftest/swiftest/init_cond.py | 50 +++++- python/swiftest/swiftest/io.py | 176 ++++++++++++++++++- python/swiftest/swiftest/simulation_class.py | 151 ++++++++++++++-- python/swiftest/swiftest/tool.py | 12 ++ 4 files changed, 371 insertions(+), 18 deletions(-) diff --git a/python/swiftest/swiftest/init_cond.py b/python/swiftest/swiftest/init_cond.py index 75c48377b..89e1c8dee 100644 --- a/python/swiftest/swiftest/init_cond.py +++ b/python/swiftest/swiftest/init_cond.py @@ -22,7 +22,7 @@ def solar_system_horizons(plname, idval, param, ephemerides_start_date, ds): Returns ------- - xarray dataset + ds : xarray dataset """ # Planet ids planetid = { @@ -255,7 +255,53 @@ def solar_system_horizons(plname, idval, param, ephemerides_start_date, ds): return ds def vec2xr(param, idvals, namevals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, rhill=None, Ip1=None, Ip2=None, Ip3=None, rotx=None, roty=None, rotz=None, t=0.0): - + """ + Converts and stores the variables of all bodies in an xarray dataset. + + Parameters + ---------- + param : dict + Swiftest paramuration parameters. + idvals : integer + Array of body index values. + namevals : + + v1 : array of floats + xh + v2 : array of floats + yh + v3 : array of floats + zh + v4 : array of floats + vhxh + v5 : array of floats + vhyh + v6 : array of floats + vhzh + GMpl : array of floats + G*mass + Rpl : array of floats + radius + rhill : array of floats + Hill Radius + Ip1 : array of floats + Principal axes moments of inertia + Ip2 : array of floats + Principal axes moments of inertia + Ip3 : array of floats + Principal axes moments of inertia + rox : array of floats + Rotation rate vector + roty : array of floats + Rotation rate vector + rotz : array of floats + Rotation rate vector + t : array of floats + Time at start of simulation + Returns + ------- + ds : xarray dataset + """ if param['ROTATION'] == 'YES': if Ip1 is None: Ip1 = np.full_like(v1, 0.4) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index ff456d0c1..138943e6d 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -271,6 +271,20 @@ def read_swift_param(param_file_name, startfile="swift.in", verbose=True): def write_swift_param(param, param_file_name): + """ + Writes a Swift param.in file. + + Parameters + ---------- + param : dictionary + The entries in the user parameter file + param_file_name : string + File name of the input parameter file + + Returns + ------- + Prints a text file containing the parameter information. + """ outfile = open(param_file_name, 'w') print(param['T0'], param['TSTOP'], param['DT'], file=outfile) print(param['DTOUT'], param['DTDUMP'], file=outfile) @@ -283,6 +297,20 @@ def write_swift_param(param, param_file_name): def write_labeled_param(param, param_file_name): + """ + Writes a Swifter/Swiftest param.in file. + + Parameters + ---------- + param : dictionary + The entries in the user parameter file + param_file_name : string + File name of the input parameter file + + Returns + ------- + Prints a text file containing the parameter information. + """ outfile = open(param_file_name, 'w') keylist = ['! VERSION', 'T0', @@ -404,6 +432,29 @@ def swifter_stream(f, param): def make_swiftest_labels(param): + """ + Creates the lables for the variables to be included in the output file. + + Parameters + ---------- + param : dictionary + The entries in the user parameter file + + Returns + ------- + clab : string list + Labels for the cvec data + plab : string list + Labels for the pvec data + tlab : string list + Labels for the tvec data + infolab_float : string list + Labels for floating point data + infolab_int : + Labels for integer data + infolab_str : + Labels for string data + """ tlab = [] if param['OUT_FORM'] == 'XV' or param['OUT_FORM'] == 'XVEL': tlab.append('xhx') @@ -746,10 +797,32 @@ def swiftest2xr(param, verbose=True): return ds def xstrip(a): + """ + Cleans up the string values in the DataSet to remove extra white space + + Parameters + ---------- + a : xarray dataset + + Returns + ------- + da : xarray dataset with the strings cleaned up + """ func = lambda x: np.char.strip(x) return xr.apply_ufunc(func, a.str.decode(encoding='utf-8'),dask='parallelized') def string_converter(da): + """ + Converts a string to a unicode string + + Parameters + ---------- + da : xarray dataset + + Returns + ------- + da : xarray dataset with the strings cleaned up + """ if da.dtype == np.dtype(object): da = da.astype(' ") if plname == '': @@ -1605,6 +1749,21 @@ def swift2swiftest(swift_param, plname="", tpname="", cbname="", conversion_ques return swiftest_param def swiftest2swifter_param(swiftest_param, J2=0.0, J4=0.0): + """ + Converts from a Swiftest run to a Swifter run + + Parameters + ---------- + swiftest_param : dictionary + Swiftest input parameters. + J2 : float + Central body oblateness term. Default spherical. + J4 : float + Central body oblateness term. Default spherical. + Returns + ------- + swifter_param : A set of input files for a new Swifter run + """ swifter_param = swiftest_param CBIN = swifter_param.pop("CB_IN", None) GMTINY = swifter_param.pop("GMTINY", None) @@ -1639,6 +1798,21 @@ def swiftest2swifter_param(swiftest_param, J2=0.0, J4=0.0): def swifter2swift_param(swifter_param, J2=0.0, J4=0.0): + """ + Converts from a Swifter run to a Swift run + + Parameters + ---------- + swifter_param : dictionary + Swifter input parameters. + J2 : float + Central body oblateness term. Default spherical. + J4 : float + Central body oblateness term. Default spherical. + Returns + ------- + swift_param : A set of input files for a new Swift run + """ swift_param = { '! VERSION': f"Swift parameter input file converted from Swifter", 'T0': 0.0, diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 60a9036ea..a62723bbf 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -80,7 +80,7 @@ def add(self, plname, date=date.today().isoformat(), idval=None): Date to use when obtaining the ephemerides in the format YYYY-MM-DD. Defaults to "today" Returns ------- - self.ds : xarray dataset + self.ds : xarray dataset """ self.ds = init_cond.solar_system_horizons(plname, idval, self.param, date, self.ds) return @@ -93,21 +93,35 @@ def addp(self, idvals, namevals, t1, t2, t3, t4, t5, t6, GMpl=None, Rpl=None, rh Parameters ---------- - idvals : Integer array of body index values. - t1 : xh for param['IN_FORM'] == "XV"; a for param['IN_FORM'] == "EL" - t2 : yh for param['IN_FORM'] == "XV"; e for param['IN_FORM'] == "EL" - t3 : zh for param['IN_FORM'] == "XV"; inc for param['IN_FORM'] == "EL" - t4 : vhxh for param['IN_FORM'] == "XV"; capom for param['IN_FORM'] == "EL" - t5 : vhyh for param['IN_FORM'] == "XV"; omega for param['IN_FORM'] == "EL" - t6 : vhzh for param['IN_FORM'] == "XV"; capm for param['IN_FORM'] == "EL" - Gmass : Optional: Array of G*mass values if these are massive bodies - radius : Optional: Array radius values if these are massive bodies - rhill : Optional: Array rhill values if these are massive bodies - Ip1,y,z : Optional: Principal axes moments of inertia - rotx,y,z: Optional: Rotation rate vector components + idvals : integer + Array of body index values. + t1 : float + xh for param['IN_FORM'] == "XV"; a for param['IN_FORM'] == "EL" + t2 : float + yh for param['IN_FORM'] == "XV"; e for param['IN_FORM'] == "EL" + t3 : float + zh for param['IN_FORM'] == "XV"; inc for param['IN_FORM'] == "EL" + t4 : float + vhxh for param['IN_FORM'] == "XV"; capom for param['IN_FORM'] == "EL" + t5 : float + vhyh for param['IN_FORM'] == "XV"; omega for param['IN_FORM'] == "EL" + t6 : float + vhzh for param['IN_FORM'] == "XV"; capm for param['IN_FORM'] == "EL" + Gmass : float + Optional: Array of G*mass values if these are massive bodies + radius : float + Optional: Array radius values if these are massive bodies + rhill : float + Optional: Array rhill values if these are massive bodies + Ip1,y,z : float + Optional: Principal axes moments of inertia + rotx,y,z: float + Optional: Rotation rate vector components + t : float + Optional: Time at start of simulation Returns ------- - self.ds : xarray dataset + self.ds : xarray dataset """ if t is None: t = self.param['T0'] @@ -122,6 +136,19 @@ def addp(self, idvals, namevals, t1, t2, t3, t4, t5, t6, GMpl=None, Rpl=None, rh def read_param(self, param_file, codename="Swiftest", verbose=True): + """ + Reads in a param.in file and determines whether it is a Swift/Swifter/Swiftest parameter file. + + Parameters + ---------- + param_file : string + File name of the input parameter file + codename : string + Type of parameter file, either "Swift", "Swifter", or "Swiftest" + Returns + ------- + self.ds : xarray dataset + """ if codename == "Swiftest": self.param = io.read_swiftest_param(param_file, self.param, verbose=verbose) self.codename = "Swiftest" @@ -138,6 +165,17 @@ def read_param(self, param_file, codename="Swiftest", verbose=True): def write_param(self, param_file, param=None): + """ + Writes to a param.in file and determines whether the output format needs to be converted between Swift/Swifter/Swiftest. + + Parameters + ---------- + param_file : string + File name of the input parameter file + Returns + ------- + self.ds : xarray dataset + """ if param is None: param = self.param # Check to see if the parameter type matches the output type. If not, we need to convert @@ -153,7 +191,27 @@ def write_param(self, param_file, param=None): def convert(self, param_file, newcodename="Swiftest", plname="pl.swiftest.in", tpname="tp.swiftest.in", cbname="cb.swiftest.in", conversion_questions={}): """ - Converts simulation input files from one code type to another (Swift, Swifter, or Swiftest). Returns the old parameter configuration. + Converts simulation input files from one format to another (Swift, Swifter, or Swiftest). + + Parameters + ---------- + param_file : string + File name of the input parameter file + newcodename : string + Name of the desired format (Swift/Swifter/Swiftest) + plname : string + File name of the massive body input file + tpname : string + File name of the test particle input file + cbname : string + File name of the central body input file + conversion_questions : dictronary + Dictionary of additional parameters required to convert between formats + + Returns + ------- + oldparam : xarray dataset + The old parameter configuration. """ oldparam = self.param if self.codename == newcodename: @@ -186,6 +244,16 @@ def convert(self, param_file, newcodename="Swiftest", plname="pl.swiftest.in", t def bin2xr(self): + """ + Converts simulation output files from a flat binary file to a xarray dataset. + + Parameters + ---------- + + Returns + ------- + self.ds : xarray dataset + """ if self.codename == "Swiftest": self.ds = io.swiftest2xr(self.param, verbose=self.verbose) if self.verbose: print('Swiftest simulation data stored as xarray DataSet .ds') @@ -200,6 +268,18 @@ def bin2xr(self): def follow(self, codestyle="Swifter"): + """ + An implementation of the Swift tool_follow algorithm. Under development. Currently only for Swift simulations. + + Parameters + ---------- + codestyle : string + Name of the desired format (Swift/Swifter/Swiftest) + + Returns + ------- + fol : xarray dataset + """ if self.ds is None: self.bin2xr() if codestyle == "Swift": @@ -227,6 +307,23 @@ def follow(self, codestyle="Swifter"): def save(self, param_file, framenum=-1, codename="Swiftest"): + """ + Saves an xarray dataset to a set of input files. + + Parameters + ---------- + param_file : string + Name of the parameter input file + framenum : integer (default=-1) + Time frame to use to generate the initial conditions. If this argument is not passed, the default is to use the last frame in the dataset. + codename : string + Name of the desired format (Swift/Swifter/Swiftest) + + Returns + ------- + self.ds : xarray dataset + """ + if codename == "Swiftest": io.swiftest_xr2infile(self.ds, self.param, framenum) self.write_param(param_file) @@ -243,6 +340,30 @@ def save(self, param_file, framenum=-1, codename="Swiftest"): return def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_file="param.new.in", new_initial_conditions_file="bin_in.nc", restart=False, codename="Swiftest"): + """ + Generates a set of input files from a old output file. + + Parameters + ---------- + framenum : integer (default=-1) + Time frame to use to generate the initial conditions. If this argument is not passed, the default is to use the last frame in the dataset. + new_param : string + File to copy parameters from. Default is the old parameter file. + new_param_file : string + Name of the new parameter file. + new_initial_conditions_file : string + Name of the new NetCDF file containing the new initial conditions. + restart : True or False + If True, overwrite the old output file. If False, generate a new output file. + codename : string + Name of the desired format (Swift/Swifter/Swiftest) + + Returns + ------- + frame : NetCDF dataset + """ + + if codename != "Swiftest": self.save(new_param_file, framenum, codename) return diff --git a/python/swiftest/swiftest/tool.py b/python/swiftest/swiftest/tool.py index efdc82a8b..41742061a 100644 --- a/python/swiftest/swiftest/tool.py +++ b/python/swiftest/swiftest/tool.py @@ -8,6 +8,18 @@ """ def wrap_angle(angle): + """ + Converts angles to be between 0 and 360 degrees. + + Parameters + ---------- + angle : float + Angle to be converted + Returns + ------- + angle : float + Converted angle + """ while np.any(angle >= 360.0 ): angle[angle >= 360.0] -= 360.0 while np.any(angle < 0.0):