From 4301f6b2dc1502538f79cff991178c56846dfb35 Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 20 Dec 2023 11:27:33 -0500 Subject: [PATCH] Restructured project and improved docstrings for the documentation page. Changed the structure of the imports so that the _bindings module is only imported when it needs to be run, so that the documentation can be built without it. --- docs/api.rst | 118 +++++++++- docs/conf.py | 3 +- pyproject.toml | 31 +-- src/globals/globals_module.f90 | 2 +- swiftest/init_cond.py | 5 +- swiftest/io.py | 319 ++++++++++++-------------- swiftest/simulation_class.py | 407 +++++++++++++++++++++------------ swiftest/tool.py | 104 +++++---- version.txt | 2 +- 9 files changed, 591 insertions(+), 400 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index 9cad3b1d4..0e757b398 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -42,7 +42,6 @@ Setting Simulation Parameters Simulation.set_init_cond_files Simulation.set_output_files Simulation.set_unit_system - Simulation.update_param_units Simulation.set_distance_range Simulation.set_ephemeris_date @@ -88,6 +87,123 @@ File Input and Output Simulation.convert Simulation.clean + +Initial Conditions Generation Functions +======================================= + +.. autosummary:: + :toctree: generated/ + + swiftest.init_cond.solar_system_horizons + swiftest.init_cond.horizons_query + swiftest.init_cond.horizons_get_physical_properties + swiftest.init_cond.vec2xr + + +Input/Output Processing Functions +================================== + +A collection of functions for processing simulation files. + +Reading and writing simulation parameter and initial conditions files +--------------------------------------------------------------------- + +.. autosummary:: + :toctree: generated/ + + swiftest.io.process_netcdf_input + swiftest.io.read_swiftest_param + swiftest.io.read_swifter_param + swiftest.io.read_swift_param + swiftest.io.write_swift_param + swiftest.io.write_labeled_param + swiftest.io.select_active_from_frame + swiftest.io.swiftest_xr2infile + +Tools for fixing differences between NetCDF-Fortran and xarray data structures +------------------------------------------------------------------------------ + +.. autosummary:: + :toctree: generated/ + + swiftest.io.swiftest2xr + swiftest.io.reorder_dims + swiftest.io.fix_types + + +Conversions between legacy integrator formats and Swiftest +---------------------------------------------------------- + +.. autosummary:: + :toctree: generated/ + + swiftest.io.swifter2swiftest + swiftest.io.swifter2xr + swiftest.io.swifter_xr2infile + swiftest.io.swiftest2swifter_param + swiftest.io.swift2swifter + swiftest.io.swift2swiftest + +Tools +===== + +Miscellaneous helper functions + +.. autosummary:: + :toctree: generated/ + + swiftest.tool.magnitude + swiftest.tool.wrap_angle + swiftest.tool.follow_swift + swiftest.tool.danby + swiftest.tool.el2xv_one + swiftest.tool.el2xv_vec + swiftest.tool.xv2el_one + swiftest.tool.xv2el_vec + +Constants +========= + +The `constants` module defines several astronomical and physical constants. Below is a description of each constant: + +.. list-table:: + :widths: 25 75 + :header-rows: 1 + + * - Constant + - Description + * - ``GC`` + - The gravitational constant (G) from Astropy constants, in SI units (m^3 kg^-1 s^-2). + * - ``AU2M`` + - Astronomical Unit in meters, representing the average distance from the Earth to the Sun. + * - ``GMSun`` + - Standard gravitational parameter for the Sun in m^3 s^-2. + * - ``MSun`` + - Mass of the Sun in kilograms. + * - ``RSun`` + - Radius of the Sun in meters. + * - ``MEarth`` + - Mass of the Earth in kilograms. + * - ``REarth`` + - Radius of the Earth in meters. + * - ``GMEarth`` + - Standard gravitational parameter for the Earth in m^3 s^-2. + * - ``JD2S`` + - Number of seconds in a Julian day. + * - ``YR2S`` + - Number of seconds in a Julian year (365.25 days). + * - ``einsteinC`` + - Speed of light in vacuum in meters per second. + * - ``J2Sun`` + - Solar quadrupole moment coefficient (J2), indicating the extent of the Sun's equatorial bulge. + * - ``J4Sun`` + - Higher order coefficient (J4) for the Sun's shape, indicating asymmetry in its mass distribution. + * - ``rotpoleSun`` + - SkyCoord object representing the rotation pole of the Sun in right ascension (ra) and declination (dec), converted to Cartesian coordinates. + * - ``rotSun`` + - Angular velocity of the Sun's rotation in radians per second, considering an average rotational period of 25.05 days. + + Fortran API Documentation ========================= diff --git a/docs/conf.py b/docs/conf.py index ff2b206ec..2aae7ae5b 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -9,9 +9,8 @@ from sphinx.util import logging # Disable import of swiftest._bindings so that we don't have to build the Fortran code when building the docs -import swiftest autodoc_mock_imports = ['swiftest._bindings'] - +import swiftest # -- Project information ----------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information project = 'Swiftest' diff --git a/pyproject.toml b/pyproject.toml index ee18493cf..ffbdad11f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "swiftest" -version = "2023.12.0" +version = "2023.12.1" authors=[ {name = 'David A. Minton', email='daminton@purdue.edu'}, {name = 'Carlisle Wishard'}, @@ -115,12 +115,7 @@ before-all = [ "LIBS=\"\" buildscripts/build_dependencies.sh -p ${PREFIX} -d ${HOME}/Downloads -m ${MACOSX_DEPLOYMENT_TARGET}" ] -[tool.cibuildwheel.linux] -skip = "cp312-* pp* -manylinux_i686* *-musllinux*" -before-all = [ - "yum install doxygen libxml2-devel libcurl-devel -y", - "buildscripts/build_dependencies.sh -p /usr/local" -] + [tool.cibuildwheel.linux.environment] PREFIX="/usr/local" NCDIR="${PREFIX}" @@ -138,22 +133,14 @@ HDF5_DIR="${HDF5_ROOT}/cmake" PATH="${HDF5_ROOT}/bin:${PATH}" CMAKE_INSTALL_LIBDIR="lib" +[tool.cibuildwheel.linux] +skip = "cp312-* pp* -manylinux_i686* *-musllinux*" +before-all = [ + "yum install doxygen libxml2-devel libcurl-devel -y", + "buildscripts/build_dependencies.sh -p /usr/local" +] + [[tool.scikit-build.generate]] path = "version.txt" location = "source" template = '''${version}''' - -[project.optional-dependencies] -docs = [ - "sphinx", - "sphinx-autosummary-accessors", - "sphinx-book-theme >= 0.3.0", - "sphinx-copybutton", - "sphinx-design", - "sphinx-inline-tabs", - "sphinxext-rediraffe", - "sphinxext-opengraph", - "nbsphinx", - "ipython", - "ford" -] \ No newline at end of file diff --git a/src/globals/globals_module.f90 b/src/globals/globals_module.f90 index 8bf98c6cd..03c1ab86a 100644 --- a/src/globals/globals_module.f90 +++ b/src/globals/globals_module.f90 @@ -48,7 +48,7 @@ module globals integer(I4B), parameter :: UPPERCASE_OFFSET = iachar('A') - iachar('a') !! ASCII character set parameter for lower to upper !! conversion - offset between upper and lower - character(*), parameter :: VERSION = "2023.12.0" !! Swiftest version + character(*), parameter :: VERSION = "2023.12.1" !! Swiftest version !> Symbolic name for integrator types character(*), parameter :: UNKNOWN_INTEGRATOR = "UKNOWN INTEGRATOR" diff --git a/swiftest/init_cond.py b/swiftest/init_cond.py index 3190b77ae..cea5912ba 100644 --- a/swiftest/init_cond.py +++ b/swiftest/init_cond.py @@ -37,7 +37,6 @@ def horizons_get_physical_properties(altid,**kwargs): **kwargs: Any Additional keyword arguments to pass to the query method (see https://astroquery.readthedocs.io/en/latest/jplhorizons/jplhorizons.html) - Returns ------- MSun_over_Mpl : float @@ -184,7 +183,6 @@ def horizons_query(id, ephemerides_start_date, exclude_spacecraft=True, verbose= An astroquery.jplhorizons HorizonsClass object. Or None if no match was found. altid: string list | None A list of alternate ids if more than one object matches the list - """ def get_altid(errstr,exclude_spacecraft=True): @@ -416,6 +414,7 @@ def solar_system_horizons(name: str, return id,name,a,e,inc,capom,omega,capm,rh,vh,Gmass,Rpl,rhill,Ip,rot,J2,J4 + def vec2xr(param: Dict, **kwargs: Any): """ Converts and stores the variables of all bodies in an xarray dataset. @@ -458,9 +457,11 @@ def vec2xr(param: Dict, **kwargs: Any): instead of passing Ip1, Ip2, and Ip3 separately time : array of floats Time at start of simulation + Returns ------- ds : xarray dataset + Dataset containing the variables of all bodies passed in kwargs """ scalar_dims = ['id'] vector_dims = ['id','space'] diff --git a/swiftest/io.py b/swiftest/io.py index 99bc3e4b6..33fdb1e8f 100644 --- a/swiftest/io.py +++ b/swiftest/io.py @@ -72,7 +72,7 @@ char_varnames = ["space"] int_varnames = ["id", "ntp", "npl", "nplm", "discard_body_id", "collision_id", "status"] -def bool2yesno(boolval): +def _bool2yesno(boolval): """ Converts a boolean into a string of either "YES" or "NO". @@ -83,8 +83,8 @@ def bool2yesno(boolval): Returns ------- - {"YES","NO"} - + str + {"YES","NO"} """ if boolval: return "YES" @@ -92,7 +92,7 @@ def bool2yesno(boolval): return "NO" -def bool2tf(boolval): +def _bool2tf(boolval): """ Converts a boolean into a string of either "T" or "F". @@ -103,7 +103,8 @@ def bool2tf(boolval): Returns ------- - {"T","F"} + str + {"T","F"} """ if boolval: @@ -112,7 +113,7 @@ def bool2tf(boolval): return "F" -def str2bool(input_str): +def _str2bool(input_str): """ Converts a string into an equivalent boolean. @@ -123,7 +124,7 @@ def str2bool(input_str): Returns ------- - {True, False} + bool """ if type(input_str) is bool: @@ -138,7 +139,7 @@ def str2bool(input_str): raise ValueError(f"{input_str} is not recognized as boolean") -def real2float(realstr): +def _real2float(realstr): """ Converts a Fortran-generated ASCII string of a real value into a numpy float type. Handles cases where double precision numbers in exponential notation use 'd' or 'D' instead of 'e' or 'E' @@ -150,7 +151,7 @@ def real2float(realstr): Returns ------- - : float + float The converted floating point value of the input string """ return float(realstr.replace('d', 'E').replace('D', 'E')) @@ -184,8 +185,8 @@ def read_swiftest_param(param_file_name, param, verbose=True): param[key] = fields[1] # Special case of CHK_QMIN_RANGE requires a second input if fields[0].upper() == 'CHK_QMIN_RANGE': - alo = real2float(fields[1]) - ahi = real2float(fields[2]) + alo = _real2float(fields[1]) + ahi = _real2float(fields[2]) param['CHK_QMIN_RANGE'] = f"{alo} {ahi}" for uc in upper_str_param: @@ -198,19 +199,28 @@ def read_swiftest_param(param_file_name, param, verbose=True): for f in float_param: if f in param and type(param[f]) is str: - param[f] = real2float(param[f]) + param[f] = _real2float(param[f]) for b in bool_param: if b in param: - param[b] = str2bool(param[b]) + param[b] = _str2bool(param[b]) except IOError: print(f"{param_file_name} not found.") return param def reorder_dims(ds): - - # Re-order dimension coordinates so that they are in the same order as the Fortran side + """ + Re-order dimension coordinates so that they are in the same order as the Fortran side + + Parameters + ---------- + ds : xarray dataset + + Returns + ------- + ds : xarray dataset with the dimensions re-ordered + """ idx = ds.indexes if "id" in idx: dim_order = ["time", "id", "space"] @@ -231,10 +241,11 @@ def read_swifter_param(param_file_name, verbose=True): ---------- param_file_name : string File name of the input parameter file - + verbose : bool, default True + Print out information about the file being read Returns ------- - param + param : dict A dictionary containing the entries in the user parameter file """ param = { @@ -278,26 +289,26 @@ def read_swifter_param(param_file_name, verbose=True): if (key == fields[0].upper()): param[key] = fields[1] # Special case of CHK_QMIN_RANGE requires a second input if fields[0].upper() == 'CHK_QMIN_RANGE': - alo = real2float(fields[1]) - ahi = real2float(fields[2]) + alo = _real2float(fields[1]) + ahi = _real2float(fields[2]) param['CHK_QMIN_RANGE'] = f"{alo} {ahi}" param['ISTEP_OUT'] = int(param['ISTEP_OUT']) param['ISTEP_DUMP'] = int(param['ISTEP_DUMP']) - param['T0'] = real2float(param['T0']) - param['TSTOP'] = real2float(param['TSTOP']) - param['DT'] = real2float(param['DT']) - param['J2'] = real2float(param['J2']) - param['J4'] = real2float(param['J4']) - param['CHK_RMIN'] = real2float(param['CHK_RMIN']) - param['CHK_RMAX'] = real2float(param['CHK_RMAX']) - param['CHK_EJECT'] = real2float(param['CHK_EJECT']) - param['CHK_QMIN'] = real2float(param['CHK_QMIN']) + param['T0'] = _real2float(param['T0']) + param['TSTOP'] = _real2float(param['TSTOP']) + param['DT'] = _real2float(param['DT']) + param['J2'] = _real2float(param['J2']) + param['J4'] = _real2float(param['J4']) + param['CHK_RMIN'] = _real2float(param['CHK_RMIN']) + param['CHK_RMAX'] = _real2float(param['CHK_RMAX']) + param['CHK_EJECT'] = _real2float(param['CHK_EJECT']) + param['CHK_QMIN'] = _real2float(param['CHK_QMIN']) for b in bool_param: if b in param: - param[b] = str2bool(param[b]) + param[b] = _str2bool(param[b]) if param['C'] != '-1.0': - param['C'] = real2float(param['C']) + param['C'] = _real2float(param['C']) else: param.pop('C', None) except IOError: @@ -314,7 +325,9 @@ def read_swift_param(param_file_name, startfile="swift.in", verbose=True): ---------- param_file_name : string File name of the input parameter file - + startfile : string, default "swift.in" + File name of the input start file + Returns ------- param : dict @@ -360,14 +373,14 @@ def read_swift_param(param_file_name, startfile="swift.in", verbose=True): line = f.readline().split() for i, l in enumerate(line): line[i] = l - param['T0'] = real2float(line[0]) - param['TSTOP'] = real2float(line[1]) - param['DT'] = real2float(line[2]) + param['T0'] = _real2float(line[0]) + param['TSTOP'] = _real2float(line[1]) + param['DT'] = _real2float(line[2]) line = f.readline().split() for i, l in enumerate(line): line[i] = l - param['DTOUT'] = real2float(line[0]) - param['DTDUMP'] = real2float(line[1]) + param['DTOUT'] = _real2float(line[0]) + param['DTDUMP'] = _real2float(line[1]) line = f.readline().split() param['L1'] = line[0].upper() param['L2'] = line[1].upper() @@ -379,10 +392,10 @@ def read_swift_param(param_file_name, startfile="swift.in", verbose=True): line = f.readline().split() for i, l in enumerate(line): line[i] = l - param['RMIN'] = real2float(line[0]) - param['RMAX'] = real2float(line[1]) - param['RMAXU'] = real2float(line[2]) - param['QMIN'] = real2float(line[3]) + param['RMIN'] = _real2float(line[0]) + param['RMAX'] = _real2float(line[1]) + param['RMAXU'] = _real2float(line[2]) + param['QMIN'] = _real2float(line[3]) param['LCLOSE'] = line[4].upper() param['BINARY_OUTPUTFILE'] = f.readline().strip() param['STATUS_FLAG_FOR_OPEN_STATEMENTS'] = f.readline().strip().upper() @@ -405,7 +418,7 @@ def write_swift_param(param, param_file_name): Returns ------- - Prints a text file containing the parameter information. + None """ outfile = open(param_file_name, 'w') print(param['T0'], param['TSTOP'], param['DT'], file=outfile) @@ -431,7 +444,7 @@ def write_labeled_param(param, param_file_name): Returns ------- - Prints a text file containing the parameter information. + None """ outfile = open(param_file_name, 'w') ptmp = param.copy() @@ -440,28 +453,30 @@ def write_labeled_param(param, param_file_name): val = ptmp.pop(key, None) if val is not None: if type(val) is bool: - print(f"{key:<32} {bool2yesno(val)}", file=outfile) + print(f"{key:<32} {_bool2yesno(val)}", file=outfile) else: print(f"{key:<32} {val}", file=outfile) # Print the remaining key/value pairs in whatever order for key, val in ptmp.items(): if val != "": if type(val) is bool: - print(f"{key:<32} {bool2yesno(val)}", file=outfile) + print(f"{key:<32} {_bool2yesno(val)}", file=outfile) else: print(f"{key:<32} {val}", file=outfile) outfile.close() return -def swifter_stream(f, param): +def _swifter_stream(f, param): """ Reads in a Swifter bin.dat file and returns a single frame of data as a datastream Parameters ---------- f : file object + File object of the Swifter bin.dat file param : dict + Swifter parameters Yields ------- @@ -557,6 +572,7 @@ def swifter_stream(f, param): yield t, npl, [plid, elempvec, GMpl, Rpl], plab, \ ntp, [tpid, elemtvec], tlab + def swifter2xr(param, verbose=True): """ Converts a Swifter binary data file into an xarray DataSet. @@ -565,7 +581,8 @@ def swifter2xr(param, verbose=True): ---------- param : dict Swifter parameters - + verbose : bool, default True + Print out information about the file being read Returns ------- xarray dataset @@ -575,7 +592,7 @@ def swifter2xr(param, verbose=True): tp = [] with FortranFile(param['BIN_OUT'], 'r') as f: for t, npl, pvec, plab, \ - ntp, tvec, tlab in swifter_stream(f, param): + ntp, tvec, tlab in _swifter_stream(f, param): sys.stdout.write('\r' + f"Reading in time {t[0]:.3e}") sys.stdout.flush() @@ -628,6 +645,8 @@ def swiftest2xr(param, verbose=True, dask=False): ---------- param : dict Swiftest parameters + verbose : bool, default True + Print out information about the file being read dask : bool, default False Use Dask to lazily load data (useful for very large datasets) @@ -653,13 +672,14 @@ def swiftest2xr(param, verbose=True, dask=False): return ds -def xstrip_nonstr(a): +def _xstrip_nonstr(a): """ Cleans up the string values in the DataSet to remove extra white space Parameters ---------- - a : xarray dataset + a : xarray + Input dataset Returns ------- @@ -669,7 +689,7 @@ def xstrip_nonstr(a): return xr.apply_ufunc(func, a.str.decode(encoding='utf-8'),dask='parallelized') -def xstrip_str(a): +def _xstrip_str(a): """ Cleans up the string values in the DataSet to remove extra white space @@ -685,7 +705,7 @@ def xstrip_str(a): return xr.apply_ufunc(func, a,dask='parallelized') -def string_converter(da): +def _string_converter(da): """ Converts a string to a unicode string @@ -699,12 +719,12 @@ def string_converter(da): """ da = da.astype(' ") - if GMTINY != '' and real2float(GMTINY.strip()) > 0: - swiftest_param['GMTINY'] = real2float(GMTINY.strip()) + if GMTINY != '' and _real2float(GMTINY.strip()) > 0: + swiftest_param['GMTINY'] = _real2float(GMTINY.strip()) # Remove the unneeded parameters if 'C' in swiftest_param: @@ -1623,6 +1588,7 @@ def swiftest2swifter_param(swiftest_param, J2=0.0, J4=0.0): 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 @@ -1669,6 +1635,7 @@ def swifter2swift_param(swifter_param, J2=0.0, J4=0.0): 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 diff --git a/swiftest/simulation_class.py b/swiftest/simulation_class.py index 09c8b1879..55f4db28b 100644 --- a/swiftest/simulation_class.py +++ b/swiftest/simulation_class.py @@ -11,15 +11,11 @@ from __future__ import annotations -from swiftest import io -from swiftest import init_cond -from swiftest import tool -from swiftest import constants -from swiftest._bindings import driver -from swiftest import __file__ as _pyfile -import json +from . import io +from . import init_cond +from . import tool +from . import constants import os -from glob import glob from pathlib import Path import datetime import xarray as xr @@ -301,13 +297,13 @@ def __init__(self,read_param: bool = False, Turns on the computation of energy, angular momentum, and mass conservation and reports the values every output step of a running simulation. Parameter input file equivalent is `ENERGY` - extra_force: bool, default False + extra_force : bool, default False Turns on user-defined force function. Parameter input file equivalent is `EXTRA_FORCE` - big_discard: bool, default False + big_discard : bool, default False Includes big bodies when performing a discard (Swifter only) Parameter input file equivalent is `BIG_DISCARD` - rhill_present: bool, default False + rhill_present : bool, default False Include the Hill's radius with the input files . Parameter input file equivalent is `RHILL_PRESENT` restart : bool, default False @@ -337,9 +333,8 @@ def __init__(self,read_param: bool = False, Use Dask to lazily load data (useful for very large datasets) coarray : bool, default False If true, will employ Coarrays on test particle structures to run in single program/multiple data parallel mode. - In order to use this capability, Swiftest must be compiled for Coarray support. Only certain integrators - can use Coarrays. RMVS, WHM, Helio are all compatible, but SyMBA is not, due to the way tp-pl close encounters - are handeled. + In order to use this capability, Swiftest must be compiled for Coarray support. Only certain integrators can use + Coarrays. RMVS, WHM, Helio are all compatible, but SyMBA is not, due to the way tp-pl close encounters are handeled. verbose : bool, default True If set to True, then more information is printed by Simulation methods as they are executed. Setting to False suppresses most messages other than errors. @@ -418,25 +413,31 @@ def _run_swiftest_driver(self): """ Internal callable function that executes the swiftest_driver run """ + from ._bindings import driver with _cwd(self.simdir): driver(self.integrator,str(self.param_file), "progress") return - def run(self,dask: bool = False, **kwargs): + def run(self, + dask: bool = False, + **kwargs: Any + ) -> None: """ Runs a Swiftest integration. Uses the parameters set by the `param` dictionary unless overridden by keyword arguments. Accepts any keyword arguments that can be passed to `set_parameter`. Parameters ---------- - **kwargs : Any valid keyword arguments accepted by `set_parameter` + dask : bool, default False + If true, will use Dask to lazily load data (useful for very large datasets) + **kwargs : Any + Any valid keyword arguments accepted by `set_parameter` Returns ------- None - """ if len(kwargs) > 0: @@ -464,7 +465,10 @@ def run(self,dask: bool = False, **kwargs): return - def _get_valid_arg_list(self, arg_list: str | List[str] | None = None, valid_var: Dict | None = None): + def _get_valid_arg_list(self, + arg_list: str | List[str] | None = None, + valid_var: Dict | None = None + ) -> Tuple[List[str], Dict]: """ Internal function for getters that extracts subset of arguments that is contained in the dictionary of valid argument/parameter variable pairs. @@ -485,7 +489,6 @@ def _get_valid_arg_list(self, arg_list: str | List[str] | None = None, 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: @@ -517,9 +520,11 @@ def set_simulation_time(self, dump_cadence: int | None = None, verbose: bool | None = None, **kwargs: Any - ): + ) -> Dict[str, Any]: """ - + Set the parameters that control how a simulation is run, such as start and stop time, step size, and the cadence of output + to both the screen and to file. Returns a dictionary of the parameters that were set. + Parameters ---------- t0 : float, optional @@ -555,7 +560,7 @@ def set_simulation_time(self, The number of output steps (given by `istep_out`) between when the saved data is dumped to a file. Setting it to 0 is equivalent to only dumping data to file at the end of the simulation. Default value is 10. Parameter input file equivalent is `DUMP_CADENCE` - verbose: bool, optional + verbose : bool,optional 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 @@ -565,8 +570,8 @@ def set_simulation_time(self, ------- time_dict : dict A dictionary containing the requested parameters - """ + if t0 is None and tstart is None and tstop is None and dt is None and istep_out is None and \ tstep_out is None and dump_cadence is None: return {} @@ -653,9 +658,12 @@ def set_simulation_time(self, return time_dict - def get_simulation_time(self, arg_list: str | List[str] | None = None, verbose: bool | None = None, **kwargs): + def get_simulation_time(self, + arg_list: str | List[str] | None = None, + verbose: bool | None = None, + **kwargs: Any + ) -> Dict[str, Any]: """ - Returns a subset of the parameter dictionary containing the current simulation time parameters. If the verbose option is set in the Simulation object, then it will also print the values. @@ -665,18 +673,16 @@ def get_simulation_time(self, arg_list: str | List[str] | None = None, verbose: 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", "dump_cadence"] - verbose: bool, optional + verbose : bool,optional If passed, it will override the Simulation object's verbose flag - **kwargs + **kwargs : Any 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 ------- time_dict : dict A dictionary containing the requested parameters - """ valid_var = {"t0": "T0", @@ -721,18 +727,26 @@ def get_simulation_time(self, arg_list: str | List[str] | None = None, verbose: return time_dict - def set_parameter(self, verbose: bool = True, **kwargs): + def set_parameter(self, + verbose: bool = True, + **kwargs: Any + ) -> Dict[str, Any]: """ Setter for all possible parameters. This will call each of the specialized setters using keyword arguments. If no arguments are passed, then default values will be used. + Parameters ---------- - **kwargs : Any argument listed listed in the Simulation class definition. + verbose : bool, default True + If set to True, then more information is printed by Simulation methods as they are executed. Setting to False suppresses + most messages other than errors. + **kwargs : Any + Any argument listed listed in the Simulation class definition. Returns ------- - param : A dictionary of all Simulation parameters that changed - + param : dict + A dictionary of all Simulation parameters that changed """ default_arguments = { @@ -818,17 +832,21 @@ def set_parameter(self, verbose: bool = True, **kwargs): return param_dict - def get_parameter(self, **kwargs): + def get_parameter(self, + **kwargs: Any + ) -> Dict[str, Any]: """ Setter for all possible parameters. Calls each of the specialized setters using keyword arguments + Parameters ---------- - **kwargs : Any of the arguments defined in Simulation. If none provided, it returns all arguments. + **kwargs : Any + Any of the arguments defined in Simulation. If none provided, it returns all arguments. Returns ------- - param : A dictionary of all Simulation parameters requested - + param : dict + A dictionary of all Simulation parameters requested """ # Getters returning parameter dictionary values @@ -852,12 +870,15 @@ def set_integrator(self, gmtiny: float | None = None, verbose: bool | None = None, **kwargs: Any - ): + ) -> Dict[str, Any]: """ - + Sets the integrator to be used when running a simulation. Returns a dictionary of the parameters that were set. + Parameters ---------- codename : {"swiftest", "swifter", "swift"}, optional + The name of the code to use. Case-insensitive valid options are swiftest, swifter, and swift. Currently only swiftest is + is supported for excuting runs with the run() method. integrator : {"symba","rmvs","whm","helio"}, optional Name of the n-body integrator that will be used when executing a run. mtiny : float, optional @@ -866,9 +887,9 @@ def set_integrator(self, gmtiny : float, optional The minimum G*mass of fully interacting bodies. Bodies below this mass interact with the larger bodies, but not each other (SyMBA only). Only mtiny or gmtiny is accepted, not both. - verbose: bool, optional + verbose : bool,optional If passed, it will override the Simulation object's verbose flag - **kwargs + **kwargs : Any A dictionary of additional keyword argument. This allows this method to be called by the more general set_parameter method, which takes all possible Simulation parameters as arguments, so these are ignored. @@ -878,7 +899,6 @@ def set_integrator(self, A dictionary containing the subset of the parameter dictonary that was updated by this setter """ - # TODO: Improve how it finds the executable binary update_list = [] @@ -929,9 +949,11 @@ def set_integrator(self, return integrator_dict - def get_integrator(self,arg_list: str | List[str] | None = None, verbose: bool | None = None, **kwargs: Any): + def get_integrator(self, + arg_list: str | List[str] | None = None, + verbose: bool | None = None, + **kwargs: Any) -> Dict[str, Any]: """ - Returns a subset of the parameter dictionary containing the current values of the distance range parameters. If the verbose option is set in the Simulation object, then it will also print the values. @@ -940,9 +962,9 @@ def get_integrator(self,arg_list: str | List[str] | None = None, verbose: bool | 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: ["integrator"] - verbose: bool, optional + verbose : bool, optional If passed, it will override the Simulation object's verbose flag - **kwargs + **kwargs : Any 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. @@ -1020,7 +1042,7 @@ def set_feature(self, verbose: bool | None = None, simdir: str | os.PathLike = None, **kwargs: Any - ): + ) -> Dict[str, Any]: """ Turns on or off various features of a simulation. @@ -1063,11 +1085,11 @@ def set_feature(self, compute_conservation_values : bool, optional Turns on the computation of energy, angular momentum, and mass conservation and reports the values every output step of a running simulation. - extra_force: bool, optional + extra_force : bool, optional Turns on user-defined force function. - big_discard: bool, optional + big_discard : bool, optional Includes big bodies when performing a discard (Swifter only) - rhill_present: bool, optional + rhill_present : bool, optional Include the Hill's radius with the input files. interaction_loops : {"TRIANGULAR","FLAT"}, default "TRIANGULAR" *Swiftest Experimental feature* @@ -1090,22 +1112,22 @@ def set_feature(self, In order to use this capability, Swiftest must be compiled for Coarray support. Only certain integrators can use Coarrays: RMVS, WHM, Helio are all compatible, but SyMBA is not, due to the way tp-pl close encounters are handeled. - tides: bool, optional + tides : bool, optional Turns on tidal model (IN DEVELOPMENT - IGNORED) - Yarkovsky: bool, optional + Yarkovsky : bool, optional Turns on Yarkovsky model (IN DEVELOPMENT - IGNORED) - YORP: bool, optional + YORP : bool, optional Turns on YORP model (IN DEVELOPMENT - IGNORED) restart : bool, default False If true, will restart an old run. The file given by `output_file_name` must exist before the run can execute. If false, will start a new run. If the file given by `output_file_name` exists, it will be replaced when the run is executed. - simdir: PathLike, optional + simdir : PathLike, optional Directory where simulation data will be stored, including the parameter file, initial conditions file, output file, dump files, and log files. - verbose: bool, optional + verbose : bool,optional If passed, it will override the Simulation object's verbose flag - **kwargs + **kwargs : Any A dictionary of additional keyword argument. This allows this method to be called by the more general set_parameter method, which takes all possible Simulation parameters as arguments, so these are ignored. @@ -1113,7 +1135,6 @@ def set_feature(self, ------- feature_dict : dict A dictionary containing the requested features. - """ update_list = [] @@ -1248,9 +1269,11 @@ def set_feature(self, feature_dict = self.get_feature(update_list, verbose) return feature_dict - def get_feature(self, arg_list: str | List[str] | None = None, verbose: bool | None = None, **kwargs: Any): + def get_feature(self, + arg_list: str | List[str] | None = None, + verbose: bool | None = None, + **kwargs: Any) -> Dict[str, Any]: """ - Returns a subset of the parameter dictionary containing the current value of the feature boolean values. If the verbose option is set in the Simulation object, then it will also print the values. @@ -1260,7 +1283,7 @@ def get_feature(self, arg_list: str | List[str] | None = None, verbose: bool | N A single string or list of strings containing the names of the features to extract. Default is all of: ["close_encounter_check", "general_relativity", "collision_model", "rotation", "compute_conservation_values"] 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 : Any 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. @@ -1269,7 +1292,6 @@ def get_feature(self, arg_list: str | List[str] | None = None, verbose: bool | N ------- feature_dict : dict A dictionary containing the requested features. - """ valid_var = {"close_encounter_check": "CHK_CLOSE", @@ -1316,7 +1338,7 @@ def set_init_cond_files(self, init_cond_format: Literal["EL", "XV"] | None = None, verbose: bool | None = None, **kwargs: Any - ): + ) -> Dict[str, Any]: """ Sets the initial condition file parameters in the parameters dictionary. @@ -1359,7 +1381,6 @@ def set_init_cond_files(self, ------- init_cond_file_dict : dict A dictionary containing the requested parameters - """ update_list = [] @@ -1456,9 +1477,11 @@ def ascii_file_input_error_msg(codename): return init_cond_file_dict - def get_init_cond_files(self, arg_list: str | List[str] | None = None, verbose: bool | None = None, **kwargs): + def get_init_cond_files(self, + arg_list: str | List[str] | None = None, + verbose: bool | None = None, + **kwargs: Any) -> Dict[str, Any]: """ - Returns a subset of the parameter dictionary containing the current initial condition file parameters If the verbose option is set in the Simulation object, then it will also print the values. @@ -1474,12 +1497,10 @@ def get_init_cond_files(self, arg_list: str | List[str] | None = None, verbose: 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 ------- init_cond_file_dict : dict A dictionary containing the requested parameters - """ valid_var = {"init_cond_file_type": "IN_TYPE", @@ -1538,7 +1559,7 @@ def set_output_files(self, restart: bool | None = None, verbose: bool | None = None, **kwargs: Any - ): + ) -> Dict[str, Any]: """ Sets the output file parameters in the parameter dictionary. @@ -1570,7 +1591,6 @@ def set_output_files(self, ------- output_file_dict : dict A dictionary containing the requested parameters - """ update_list = [] if output_file_type is not None: @@ -1634,9 +1654,12 @@ def set_output_files(self, return output_file_dict - def get_output_files(self, arg_list: str | List[str] | None = None, verbose: bool | None = None, **kwargs): + def get_output_files(self, + arg_list: str | List[str] | None = None, + verbose: bool | None = None, + **kwargs: Any + ) -> Dict[str, Any]: """ - Returns a subset of the parameter dictionary containing the current output file parameters If the verbose option is set in the Simulation object, then it will also print the values. @@ -1646,13 +1669,12 @@ def get_output_files(self, arg_list: str | List[str] | None = None, verbose: boo 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 + verbose : bool,optional If passed, it will override the Simulation object's verbose flag - **kwargs + **kwargs : Any 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 ------- output_file_dict : dict @@ -1690,7 +1712,8 @@ def set_unit_system(self, TU_name: str | None = None, recompute_unit_values: bool = True, verbose: bool | None = None, - **kwargs: Any): + **kwargs: Any + ) -> Dict[str, Any]: """ Setter for setting the unit conversion between one of the standard sets. @@ -1873,15 +1896,18 @@ def set_unit_system(self, 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) + self._update_param_units(MU2KG_old, DU2M_old, TU2S_old) unit_dict = self.get_unit_system(update_list, verbose) return unit_dict - def get_unit_system(self, arg_list: str | List[str] | None = None, verbose: bool | None = None, **kwargs): + def get_unit_system(self, + arg_list: str | List[str] | None = None, + verbose: bool | None = None, + **kwargs + ) -> Dict[str, Any]: """ - Returns a subset of the parameter dictionary containing the current simulation unit system. If the verbose option is set in the Simulation object, then it will also print the values. @@ -1890,6 +1916,8 @@ def get_unit_system(self, arg_list: str | List[str] | None = None, verbose: bool arg_list : str | List[str], optional A single string or list of strings containing the names of the simulation unit system Default is all of ["MU", "DU", "TU"] + verbose : bool, optional + If passed, it will override the Simulation object's verbose flag **kwargs : Any 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. @@ -1944,7 +1972,7 @@ def get_unit_system(self, arg_list: str | List[str] | None = None, verbose: bool return unit_dict - def update_param_units(self, MU2KG_old, DU2M_old, TU2S_old): + def _update_param_units(self, MU2KG_old, DU2M_old, TU2S_old): """ Updates the values of parameters that have units when the units have changed. @@ -1960,7 +1988,6 @@ def update_param_units(self, MU2KG_old, DU2M_old, TU2S_old): Returns ------- Updates the set of param dictionary values for the new unit system - """ mass_keys = ['GMTINY', 'MIN_GMFRAG'] @@ -1996,7 +2023,8 @@ def set_distance_range(self, rmax: float | None = None, qmin_coord: Literal["HELIO","BARY"] | None = None, verbose: bool | None = None, - **kwargs: Any): + **kwargs: Any + ) -> Dict[str, Any]: """ Sets the minimum and maximum distances of the simulation. @@ -2008,6 +2036,8 @@ def set_distance_range(self, Maximum distance of the simulation (CHK_RMAX, CHK_QMIN_RANGE[1]) qmin_coord : str, {"HELIO", "BARY"} coordinate frame to use for CHK_QMIN + verbose : bool, optional + If passed, it will override the Simulation object's verbose flag **kwargs : Any A dictionary of additional keyword argument. This allows this method to be called by the more general set_parameter method, which takes all possible Simulation parameters as arguments, so these are ignored. @@ -2016,7 +2046,6 @@ def set_distance_range(self, ------- range_dict : dict A dictionary containing the requested parameters. - """ if rmax is None and rmin is None and qmin_coord is None: return {} @@ -2052,17 +2081,20 @@ def set_distance_range(self, return range_dict - def get_distance_range(self, arg_list: str | List[str] | None = None, verbose: bool | None = None, **kwargs: Any): + def get_distance_range(self, + arg_list: str | List[str] | None = None, + verbose: bool | None = None, + **kwargs: Any + ) -> Dict[str, Any]: """ - Returns a subset of the parameter dictionary containing the current values of the distance range parameters. If the verbose option is set in the Simulation object, then it will also print the values. Parameters ---------- - arg_list: str | List[str], optional + 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 ["rmin", "rmax"] - verbose: bool, optional + verbose : bool,optional If passed, it will override the Simulation object's verbose flag **kwargs : Any A dictionary of additional keyword argument. This allows this method to be called by the more general @@ -2072,7 +2104,6 @@ def get_distance_range(self, arg_list: str | List[str] | None = None, verbose: b ------- range_dict : dict A dictionary containing the requested parameters. - """ valid_var = {"rmin": "CHK_RMIN", @@ -2119,7 +2150,8 @@ def add_solar_system_body(self, ephemeris_id: int | List[int] | None = None, date: str | None = None, source: str = "HORIZONS", - **kwargs: Any): + **kwargs: Any + ) -> None: """ Adds a solar system body to an existing simulation Dataset from the JPL Horizons ephemeris service. The JPL Horizons service will be searched for a body matching the string passed by `name`, or alternatively `ephemeris_id` if passed. Bodies will be @@ -2244,9 +2276,11 @@ def add_solar_system_body(self, def set_ephemeris_date(self, ephemeris_date: str | None = None, verbose: bool | None = None, - **kwargs: Any): + **kwargs: Any + ) -> str: """ - + Sets the date to use when obtaining the ephemerides. + Parameters ---------- ephemeris_date : str, optional @@ -2260,8 +2294,9 @@ def set_ephemeris_date(self, Returns ------- - Sets the `ephemeris_date` instance variable. - + ephemeris_date : str + The ISO-formatted date string for the ephemeris computation. + Also Sets the `ephemeris_date` instance variable. """ if ephemeris_date is None: @@ -2292,27 +2327,29 @@ def set_ephemeris_date(self, return ephemeris_date - def get_ephemeris_date(self, arg_list: str | List[str] | None = None, verbose: bool | None = None, **kwargs: Any): + def get_ephemeris_date(self, + arg_list: str | List[str] | None = None, + verbose: bool | None = None, + **kwargs: Any + ) -> str: """ - - Prints the current value of the ephemeris date + Prints the current value of the ephemeris date. Parameters ---------- - arg_list: str | List[str], optional + 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: ["integrator"] - verbose: bool, optional + verbose : bool, optional If passed, it will override the Simulation object's verbose flag - **kwargs + **kwargs : Any 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 ------- - ephemeris_date: str + ephemeris_date : str The ISO-formatted date string for the ephemeris computation - """ try: @@ -2327,27 +2364,30 @@ def get_ephemeris_date(self, arg_list: str | List[str] | None = None, verbose: b return ephemeris_date - def _get_instance_var(self, arg_list: str | List[str], valid_arg: Dict, verbose: bool | None = None, **kwargs: Any): + def _get_instance_var(self, + arg_list: str | List[str], valid_arg: Dict, + verbose: bool | None = None, + **kwargs: Any + ) -> Tuple[Any, ...]: """ - Prints the current value of an instance variable. Parameters ---------- - arg_list: str | List[str] + arg_list : str | List[str] A single string or list of strings containing the names of the the instance variable to get. - valid_arg: dict + valid_arg : dict A dictionary where the key is the parameter argument and the value is the equivalent instance variable value. - verbose: bool, optional + verbose : bool,optional If passed, it will override the Simulation object's verbose flag - **kwargs + **kwargs : Any 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 ------- - Tuple of instance variable values given by the arg_list - + Tuple[Any, ...] + Instance variable values given by the arg_list """ arg_vals = [] @@ -2381,7 +2421,8 @@ def add_body(self, rot: List[float] | List[npt.NDArray[np.float_]] | npt.NDArray[np.float_] | None=None, Ip: List[float] | npt.NDArray[np.float_] | None=None, J2: float | List[float] | npt.NDArray[np.float_] | None=None, - J4: float | List[float] | npt.NDArray[np.float_] | None=None): + J4: float | List[float] | npt.NDArray[np.float_] | None=None + ) -> None: """ Adds a body (test particle or massive body) to the internal DataSet given a set up 6 vectors (orbital elements or cartesian state vectors, depending on the value of self.param). Input all angles in degress. @@ -2419,16 +2460,15 @@ def add_body(self, Radius values if these are massive bodies rhill : float or array-like of float, optional Hill's radius values if these are massive bodies - rot: (3) or (n,3) array-like of float, optional + rot : (3) or (n,3) array-like of float, optional Rotation rate vectors if these are massive bodies with rotation enabled. - Ip: (3) or (n,3) array-like of float, optional + Ip : (3) or (n,3) array-like of float, optional Principal axes moments of inertia vectors if these are massive bodies with rotation enabled. Returns ------- - data : Xarray Dataset - Dasaset containing the body or bodies that were added - + None + Sets the data and init_cond instance variables each with an Xarray Dataset containing the body or bodies that were added """ #convert all inputs to numpy arrays @@ -2546,9 +2586,12 @@ def input_to_array_3d(val,n=None): return - def _combine_and_fix_dsnew(self,dsnew): + def _combine_and_fix_dsnew(self, + dsnew: xr.Dataset + ) -> xr.Dataset: """ Combines the new Dataset with the old one. Also computes the values of ntp and npl and sets the proper types. + Parameters ---------- dsnew : xarray Dataset @@ -2558,8 +2601,8 @@ def _combine_and_fix_dsnew(self,dsnew): ------- dsnew : xarray Dataset Updated Dataset with ntp, npl values and types fixed. - """ + if "id" not in self.data.dims: if len(np.unique(dsnew['name'])) == len(dsnew['name']): dsnew = dsnew.swap_dims({"id" : "name"}) @@ -2610,7 +2653,8 @@ def read_param(self, read_init_cond : bool | None = None, verbose: bool | None = None, dask: bool = False, - **kwargs : Any): + **kwargs : Any + ) -> bool: """ Reads in an input parameter file and stores the values in the param dictionary. @@ -2630,7 +2674,8 @@ def read_param(self, Returns ------- - True if the parameter file exists and is read correctly. False otherwise. + bool + True if the parameter file exists and is read correctly. False otherwise. """ if param_file is None: param_file = self.simdir / self.param_file @@ -2673,7 +2718,8 @@ def write_param(self, codename: Literal["Swiftest", "Swifter", "Swift"] | None = None, param_file: str | os.PathLike | None = None, param: Dict | None = None, - **kwargs: Any): + **kwargs: Any + ) -> None: """ Writes to a param.in file and determines whether the output format needs to be converted between Swift/Swifter/Swiftest. @@ -2684,9 +2730,9 @@ def write_param(self, variable codename param_file : str or path-like, optional Alternative file name of the input parameter file. Defaults to current instance variable self.param_file - param: Dict, optional + param : Dict, optional An alternative parameter dictionary to write out. Defaults to the current instance variable self.param - **kwargs + **kwargs : Any A dictionary of additional keyword argument. These are ignored. Returns @@ -2727,8 +2773,15 @@ def write_param(self, return - def convert(self, param_file, newcodename="Swiftest", plname="pl.swiftest.in", tpname="tp.swiftest.in", - cbname="cb.swiftest.in", conversion_questions={}, dask = False): + def convert(self, + param_file: str, + newcodename: str="Swiftest", + plname: str="pl.swiftest.in", + tpname: str="tp.swiftest.in", + cbname: str="cb.swiftest.in", + conversion_questions: Dict={}, + dask: bool=False + ) -> xr.Dataset: """ Converts simulation input files from one format to another (Swift, Swifter, or Swiftest). @@ -2752,7 +2805,7 @@ def convert(self, param_file, newcodename="Swiftest", plname="pl.swiftest.in", t Returns ------- oldparam : xarray dataset - The old parameter configuration. + The old parameter configuration. """ oldparam = self.param if self.codename == newcodename: @@ -2783,7 +2836,10 @@ def convert(self, param_file, newcodename="Swiftest", plname="pl.swiftest.in", t warnings.warn(f"Conversion from {self.codename} to {newcodename} is not supported.",stacklevel=2) return oldparam - def read_output_file(self,read_init_cond : bool = True, dask : bool = False): + def read_output_file(self, + read_init_cond : bool = True, + dask : bool = False + ) -> None: """ Reads in simulation data from an output file and stores it as an Xarray Dataset in the `data` instance variable. @@ -2793,9 +2849,11 @@ def read_output_file(self,read_init_cond : bool = True, dask : bool = False): Read in an initial conditions file along with the output file. Default is True dask : bool, default False Use Dask to lazily load data (useful for very large datasets) + Returns ------- - self.data : xarray dataset + None + Sets the data instance variable xarray dataset """ # Make a temporary copy of the parameter dictionary so we can supply the absolute path of the binary file @@ -2833,7 +2891,22 @@ def read_output_file(self,read_init_cond : bool = True, dask : bool = False): warnings.warn('Cannot process unknown code type. Call the read_param method with a valid code name. Valid options are "Swiftest", "Swifter", or "Swift".',stacklevel=2) return - def read_encounter_file(self, dask=False): + def read_encounter_file(self, + dask: bool=False + ) -> None: + """ + Reads in an encounter history file and stores it as an Xarray Dataset in the `encounters` instance variable. + + Parameters + ---------- + dask : bool, default False + Use Dask to lazily load data (useful for very large datasets) + + Returns + ------- + None + Sets the encounters instance variable xarray dataset + """ enc_file = self.simdir / "encounters.nc" if not os.path.exists(enc_file): return @@ -2853,7 +2926,22 @@ def read_encounter_file(self, dask=False): return - def read_collision_file(self, dask=False): + def read_collision_file(self, + dask: bool=False + ) -> None: + """ + Reads in a collision history file and stores it as an Xarray Dataset in the `collisions` instance variable. + + Parameters + ---------- + dask : bool, default False + Use Dask to lazily load data (useful for very large datasets) + + Returns + ------- + None + Sets the collisions instance variable xarray dataset + """ col_file = self.simdir / "collisions.nc" if not os.path.exists(col_file): @@ -2871,18 +2959,22 @@ def read_collision_file(self, dask=False): return - def follow(self, codestyle="Swifter", dask=False): + def follow(self, + codestyle: str="Swifter", + dask: bool=False + ) -> xr.Dataset: """ 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) + codestyle : str, default "Swifter" + Name of the desired format (Swift/Swifter/Swiftest) Returns ------- - fol : xarray dataset + xarray dataset + Dataset containing the variables retrieved from the follow algorithm """ if self.data is None: self.read_output_file(dask=dask) @@ -2914,7 +3006,8 @@ def save(self, param_file: str | os.PathLike | None = None, param: Dict | None = None, framenum: int = -1, - **kwargs: Any): + **kwargs: Any + ) -> None: """ Saves an xarray dataset to a set of input files. @@ -2925,11 +3018,11 @@ def save(self, variable self.codename param_file : str or path-like, optional Alternative file name of the input parameter file. Defaults to current instance variable self.param_file - param: Dict, optional + param : Dict, optional An alternative parameter dictionary to write out. Defaults to the current instance variable self.param framenum : int 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. - **kwargs + **kwargs : Any A dictionary of additional keyword argument. These are ignored. Returns @@ -2967,29 +3060,35 @@ def save(self, 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"): + def initial_conditions_from_bin(self, + framenum: int=-1, + new_param: os.PathLike=None, + new_param_file: os.PathLike="param.new.in", + new_initial_conditions_file: os.PathLike="bin_in.nc", + restart: bool=False, + codename: str="Swiftest" + ) -> xr.Dataset: """ 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) + framenum : int (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 : PathLike, optional + File to copy parameters from. Default is the old parameter file. + new_param_file : PathLike, default "param.new.in" + Name of the new parameter file. + new_initial_conditions_file : PathLike, default "bin_in.nc" + Name of the new NetCDF file containing the new initial conditions. + restart : bool, default False + If True, overwrite the old output file. If False, generate a new output file. + codename : str, default "Swiftest" + Name of the desired format (Swift/Swifter/Swiftest) Returns ------- - frame : NetCDF dataset + xarray dataset A dataset containing the extracted initial condition data. """ @@ -3035,6 +3134,14 @@ def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_fil def clean(self): """ Cleans up simulation directory by deleting old files (dump, logs, and output files). + + Parameters + ---------- + None + + Returns + ------- + None """ old_files = [self.simdir / self.param['BIN_OUT'], self.simdir / "fraggle.log", diff --git a/swiftest/tool.py b/swiftest/tool.py index 62692b667..083945dd3 100644 --- a/swiftest/tool.py +++ b/swiftest/tool.py @@ -11,15 +11,23 @@ import numpy as np import xarray as xr -""" -Functions that recreate the Swift/Swifter tool programs -""" def magnitude(ds,x): + """ + Computes the magnitude of a vector quantity from a Dataset. + + Parameters + ---------- + ds : Xarray Dataset + Dataset containing the vector quantity + x : str + Name of the vector quantity variable in the Dataset, which must have a "space" dimension (x,y,z coordinates) + """ dim = "space" ord = None return xr.apply_ufunc( np.linalg.norm, ds[x].where(~np.isnan(ds[x])), input_core_dims=[[dim]], kwargs={"ord": ord, "axis": -1}, dask="allowed" ) + def wrap_angle(angle): """ @@ -40,22 +48,28 @@ def wrap_angle(angle): angle[angle < 0.0] += 360.0 return angle + def follow_swift(ds, ifol=None, nskp=None): """ Emulates the Swift version of tool_follow.f - - + It will generate a file called follow.out containing 10 columns (angles are all in degrees):: + + 1 2 3 4 5 6 7 8 9 10 + t,a,e,inc,capom,omega,capm,peri,apo,obar + Parameters ---------- - ds : Dataset containing orbital elements + ds : Xarray Dataset + Dataset containing orbital elements + ifol : int, optional + Particle number to follow. The default is None, in which case the user is prompted to enter a particle number. + nskp : int, optional + Print frequency. The default is None, in which case the user is prompted to enter a print frequency. Returns ------- - fol : Dataset containing only the followed body with angles converted to degrees - - Generates a file called follow.out containing 10 columns (angles are all in degrees): - 1 2 3 4 5 6 7 8 9 10 - t,a,e,inc,capom,omega,capm,peri,apo,obar + fol : Xarray Dataset + Dataset containing only the followed body with angles converted to degrees """ fol = None @@ -105,9 +119,10 @@ def follow_swift(ds, ifol=None, nskp=None): return fol + def danby(M, ecc, accuracy=1e-14): """ - Danby's method to solve Kepler's equation. + Danby's method to solve Kepler's equation. See [1]_ and [2]_ for details. Parameters ---------- @@ -125,14 +140,25 @@ def danby(M, ecc, accuracy=1e-14): References __________ - Danby, J.M.A. 1988. Fundamentals of celestial mechanics. Richmond, Va., U.S.A., Willmann-Bell, 1988. 2nd ed. - Murray, C.D., Dermott, S.F., 1999. Solar system dynamics, New York, New York. ed, Cambridge University Press. + .. [1] Danby, J.M.A. 1988. Fundamentals of celestial mechanics. Richmond, Va., U.S.A., Willmann-Bell, 1988. 2nd ed. + .. [2] Murray, C.D., Dermott, S.F., 1999. Solar system dynamics, New York, New York. ed, Cambridge University Press. """ def kepler_root(E, ecc, M, deriv): """ - The Kepler equation root function. + The Kepler equation root function. The returned value depends on the value of `deriv`, where:: + + deriv = 0 : + E - e * np.sin(E) - M + deriv = 1 : + 1 - e * np.cos(E) + deriv = 2 : + e * np.sin(E) + deriv = 3 : + e * np.cos(E) + + The function will return 0 when E is correct for a given e and M Parameters ---------- @@ -147,12 +173,8 @@ def kepler_root(E, ecc, M, deriv): Returns ---------- - deriv = 0: E - e * np.sin(E) - M - deriv = 1: 1 - e * np.cos(E) - deriv = 2: e * np.sin(E) - deriv = 3: e * np.cos(E) - - Note: The function will return 0 when E is correct for a given e and M + float + The value of the Kepler equation root function """ if deriv == 0: @@ -184,8 +206,8 @@ def delta_ij(E, ecc, M, j): Returns ---------- - delta_ij value used in Danby's iterative Kepler equation solver - + float: + delta_ij value used in Danby's iterative Kepler equation solver """ if j == 1: return -kepler_root(E, ecc, M, 0) / kepler_root(E, ecc, M, 1) @@ -217,16 +239,14 @@ def delta_ij(E, ecc, M, j): raise RuntimeError("The danby function did not converge on a solution.") - def el2xv_one(mu, a, ecc, inc, Omega, omega, M): """ - Compute osculating orbital elements from relative Cartesian position and velocity + Compute osculating orbital elements from relative Cartesian position and velocity. All angular measures are returned in radians - If inclination < TINY, longitude of the ascending node is arbitrarily set to 0 + If inclination < TINY, longitude of the ascending node is arbitrarily set to 0 + If eccentricity < sqrt(TINY), argument of pericenter is arbitrarily set to 0 - If eccentricity < sqrt(TINY), argument of pericenter is arbitrarily set to 0 - - ALGORITHM: See Fitzpatrick "Principles of Cel. Mech." + ALGORITHM - See Fitzpatrick "Principles of Cel. Mech." Adapted from Martin Duncan's el2xv.f @@ -249,13 +269,10 @@ def el2xv_one(mu, a, ecc, inc, Omega, omega, M): Returns ---------- - rvec, vvec : tuple of float vectors - rvec : (3) float vector Cartesian position vector vvec : (3) float vector Cartesian velocity vector - """ if ecc < 0.0: @@ -282,12 +299,15 @@ def scget(angle): Parameters ---------- - angle : input angle + angle : float + input angle in radians Returns ------- - sx : sin of angle - cx : cos of angle + sx : float + sin of angle + cx : float + cos of angle """ TWOPI = 2 * np.pi @@ -343,8 +363,8 @@ def scget(angle): def el2xv_vec(mu, a, ecc, inc, Omega, omega, M): """ - Vectorized call to el2xv_one + Parameters ---------- mu : float @@ -364,9 +384,7 @@ def el2xv_vec(mu, a, ecc, inc, Omega, omega, M): Returns ---------- - rvec, vvec : tuple of float vectors - - rvec : (n,3) float rray + rvec : (n,3) float array Cartesian position vector vvec : (n,3) float array Cartesian velocity vector @@ -374,6 +392,7 @@ def el2xv_vec(mu, a, ecc, inc, Omega, omega, M): vecfunc = np.vectorize(el2xv_one, signature='(),(),(),(),(),(),()->(3),(3)') return vecfunc(mu, a, ecc, inc, Omega, omega, M) + def xv2el_one(mu,rvec,vvec): """ Converts from cartesian position and velocity values to orbital elements @@ -389,8 +408,6 @@ def xv2el_one(mu,rvec,vvec): Returns ---------- - a, ecc, inc, Omega, omega, M, varpi, f, lam : tuple of floats - a : float semimajor axis ecc : float @@ -409,7 +426,6 @@ def xv2el_one(mu,rvec,vvec): true anomaly (degrees) lam : float mean longitude (degrees) - """ rmag = np.linalg.norm(rvec) @@ -454,6 +470,7 @@ def xv2el_one(mu,rvec,vvec): return a, ecc, np.rad2deg(inc), np.rad2deg(Omega), np.rad2deg(omega), np.rad2deg(M), np.rad2deg(varpi), np.rad2deg(f), np.rad2deg(lam) + def xv2el_vec(mu, rvec, vvec): """ Vectorized call to xv2el_one. @@ -469,8 +486,6 @@ def xv2el_vec(mu, rvec, vvec): Returns ---------- - a, ecc, inc, Omega, omega, M, varpi, f, lam : tuple of float arrays - a : (n) float array semimajor axis ecc : (n) float array @@ -489,7 +504,6 @@ def xv2el_vec(mu, rvec, vvec): true anomaly (degrees) lam : (n) float array mean longitude (degrees) - """ vecfunc = np.vectorize(xv2el_one, signature='(),(3),(3)->(),(),(),(),(),(),(),(),()') diff --git a/version.txt b/version.txt index aa240b72d..5f73c59e9 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -2023.12.0 \ No newline at end of file +2023.12.1 \ No newline at end of file