Source code for swiftest.tool
"""
Copyright 2025 - David Minton.
This file is part of Swiftest.
Swiftest is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License
as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
Swiftest is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License along with Swiftest.
If not, see: https://www.gnu.org/licenses.
"""
import numpy as np
[docs]
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):
angle[angle < 0.0] += 360.0
return angle
[docs]
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 : 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 : Xarray Dataset
Dataset containing only the followed body with angles converted to degrees
"""
fol = None
if ifol is None:
intxt = input(" Input the particle number to follow \n")
ifol = int(intxt)
print(f"Following particle {ifol}")
if ifol < 0: # Negative numbers are planets
fol = ds.where(np.invert(np.isnan(ds["Gmass"])), drop=True)
fol = fol.where(
np.invert(np.isnan(fol["a"])), drop=True
) # Remove times where this body doesn't exist (but this also gets rid of the central body)
fol = fol.isel(
id=-ifol - 2
) # Take 1 off for 0-indexed arrays in Python, and take 1 more off because the central body is gone
elif ifol > 0: # Positive numbers are test particles
fol = ds.where(np.isnan(ds["Gmass"]), drop=True).drop_vars(["Gmass", "radius"], errors="ignore")
fol = fol.where(np.invert(np.isnan(fol["a"])), drop=True)
fol = fol.isel(id=ifol - 1) # Take 1 off for 0-indexed arrays in Python
if nskp is None:
intxt = input("Input the print frequency\n")
nskp = int(intxt)
fol["obar"] = fol["capom"] + fol["omega"]
fol["obar"] = fol["obar"].fillna(0)
fol["obar"] = wrap_angle(fol["obar"])
fol["peri"] = fol["a"] * (1.0 - fol["e"])
fol["apo"] = fol["a"] * (1.0 + fol["e"])
tslice = slice(None, None, nskp)
try:
with open("follow.out", "w") as f:
print("# 1 2 3 4 5 6 7 8 9 10", file=f)
print("# t,a,e,inc,capom,omega,capm,peri,apo,obar", file=f)
for t in fol.isel(time=tslice).time:
a = fol["a"].sel(time=t).values
e = fol["e"].sel(time=t).values
inc = fol["inc"].sel(time=t).values
capom = fol["capom"].sel(time=t).values
omega = fol["omega"].sel(time=t).values
capm = fol["capm"].sel(time=t).values
peri = fol["peri"].sel(time=t).values
apo = fol["apo"].sel(time=t).values
obar = fol["obar"].sel(time=t).values
print(
f"{t.values:15.7e} {a:22.16f} {e:22.16f} {inc:22.16f} {capom:22.16f} {omega:22.16f} {capm:22.16f} {peri:22.16f} {apo:22.16f} {obar:22.16f}",
file=f,
)
except OSError:
print("Error writing to follow.out")
return fol
