From 33cbac85019b9ee5f830b2b1044e8521d016dca4 Mon Sep 17 00:00:00 2001 From: David Minton Date: Fri, 25 Jun 2021 16:07:49 -0400 Subject: [PATCH] Major restructuring of the swiftest Python module. Now with OOP! --- {python/.idea => .idea}/.gitignore | 0 .../swiftestreader-checkpoint.py | 32 - python/swiftest/.idea/.gitignore | 3 + python/{ => swiftest}/README.md | 0 python/{ => swiftest}/setup.py | 0 python/swiftest/swiftest/__init__.py | 2 + python/swiftest/swiftest/constants.py | 15 + python/swiftest/swiftest/simulation_class.py | 47 ++ .../swiftest/swiftestio.py} | 559 +++++++++--------- python/{ => swiftest/tests}/__init__.py | 0 .../swiftest/tests/param.swift.in | 0 .../swiftest/tests}/pl.in | 0 .../swiftest/tests}/swift.in | 0 python/swiftest/tests/swift_reader.py | 6 + .../swiftest/tests}/tp.in | 0 15 files changed, 353 insertions(+), 311 deletions(-) rename {python/.idea => .idea}/.gitignore (100%) delete mode 100644 python/.ipynb_checkpoints/swiftestreader-checkpoint.py create mode 100644 python/swiftest/.idea/.gitignore rename python/{ => swiftest}/README.md (100%) rename python/{ => swiftest}/setup.py (100%) create mode 100644 python/swiftest/swiftest/__init__.py create mode 100644 python/swiftest/swiftest/constants.py create mode 100644 python/swiftest/swiftest/simulation_class.py rename python/{swiftest.py => swiftest/swiftest/swiftestio.py} (69%) rename python/{ => swiftest/tests}/__init__.py (100%) rename examples/swift_conversion/param.in => python/swiftest/tests/param.swift.in (100%) rename {examples/swift_conversion => python/swiftest/tests}/pl.in (100%) rename {examples/swift_conversion => python/swiftest/tests}/swift.in (100%) create mode 100644 python/swiftest/tests/swift_reader.py rename {examples/swift_conversion => python/swiftest/tests}/tp.in (100%) diff --git a/python/.idea/.gitignore b/.idea/.gitignore similarity index 100% rename from python/.idea/.gitignore rename to .idea/.gitignore diff --git a/python/.ipynb_checkpoints/swiftestreader-checkpoint.py b/python/.ipynb_checkpoints/swiftestreader-checkpoint.py deleted file mode 100644 index daa350d8e..000000000 --- a/python/.ipynb_checkpoints/swiftestreader-checkpoint.py +++ /dev/null @@ -1,32 +0,0 @@ -import numpy as np -from matplotlib import pyplot as plt -import os -import sys -from scipy.io import FortranFile - -binfilename = "bin.swiftest.dat" - -with FortranFile(binfilename, 'r') as f: - while True: # Loop until you read the end of file - try: - t = f.read_reals(np.float64) # Try first part of the header - print(t) - except: - break - if t[0] > tstop: break - npl = f.read_ints() - ntp = f.read_ints() - iout_form = f.read_ints() - name = f.read_ints() - px = f.read_reals(np.float64) - py = f.read_reals(np.float64) - pz = f.read_reals(np.float64) - vx = f.read_reals(np.float64) - vy = f.read_reals(np.float64) - vz = f.read_reals(np.float64) - mass = f.read_reals(np.float64) - radius = f.read_reals(np.float64) - a = f.read_reals(np.float64) - e = f.read_reals(np.float64) - inc = f.read_reals(np.float64) - yield t, name, mass, radius, np.c_[a, e] \ No newline at end of file diff --git a/python/swiftest/.idea/.gitignore b/python/swiftest/.idea/.gitignore new file mode 100644 index 000000000..26d33521a --- /dev/null +++ b/python/swiftest/.idea/.gitignore @@ -0,0 +1,3 @@ +# Default ignored files +/shelf/ +/workspace.xml diff --git a/python/README.md b/python/swiftest/README.md similarity index 100% rename from python/README.md rename to python/swiftest/README.md diff --git a/python/setup.py b/python/swiftest/setup.py similarity index 100% rename from python/setup.py rename to python/swiftest/setup.py diff --git a/python/swiftest/swiftest/__init__.py b/python/swiftest/swiftest/__init__.py new file mode 100644 index 000000000..61f0a6c2a --- /dev/null +++ b/python/swiftest/swiftest/__init__.py @@ -0,0 +1,2 @@ +from swiftest.constants import * +from swiftest.simulation_class import Simulation \ No newline at end of file diff --git a/python/swiftest/swiftest/constants.py b/python/swiftest/swiftest/constants.py new file mode 100644 index 000000000..7b031b4b7 --- /dev/null +++ b/python/swiftest/swiftest/constants.py @@ -0,0 +1,15 @@ +import numpy as np +import astropy.constants as const + +# Constants in SI units +GC = np.longdouble(const.G.value) +AU2M = np.longdouble(const.au.value) +GMSunSI = np.longdouble(const.GM_sun.value) +MSun = np.longdouble(const.M_sun.value) +RSun = np.longdouble(const.R_sun.value) +JD2S = 86400 +YR2S = np.longdouble(365.25 * JD2S) +einsteinC = np.longdouble(299792458.0) +# Solar oblatenes values: From Mecheri et al. (2004), using Corbard (b) 2002 values (Table II) +J2Sun = np.longdouble(2.198e-7) +J4Sun = np.longdouble(-4.805e-9) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py new file mode 100644 index 000000000..2adcec70b --- /dev/null +++ b/python/swiftest/swiftest/simulation_class.py @@ -0,0 +1,47 @@ +from swiftest import swiftestio +class Simulation: + """ + This is a class that define the basic Swift/Swifter/Swiftest simulation object + """ + def __init__(self, param_file_name, codename="Swiftest"): + self.read_param(param_file_name, codename) + return + + def read_param(self, param_file_name, codename="Swiftest"): + if codename == "Swiftest": + self.param = swiftestio.read_swiftest_param(param_file_name) + self.codename = "Swiftest" + elif codename == "Swifter": + self.param = swiftestio.read_swifter_param(param_file_name) + self.codename = "Swifter" + elif codename == "Swift": + self.param = swiftestio.read_swift_param(param_file_name) + self.codename = "Swift" + else: + print(f'{codename} is not a recognized code name. Valid options are "Swiftest", "Swifter", or "Swift".') + self.codename = "Unknown" + return + + def write_param(self, param_file_name): + # Check to see if the parameter type matches the output type. If not, we need to convert + codename = self.param['VERSION'].split()[1] + if codename == "Swifter" or codename == "Swiftest": + swiftestio.write_labeled_param(self.param, param_file_name) + elif codename == "Swift": + swiftestio.write_swift_param(self.param, param_file_name) + else: + print('Cannot process unknown code type. Call the read_param method with a valid code name. Valid options are "Swiftest", "Swifter", or "Swift".') + return + + def bin2xr(self): + if self.codename == "Swiftest": + self.ds = swiftestio.swiftest2xr(self.param) + elif self.codename == "Swifter": + self.ds = swiftestio.swifter2xr(self.param) + elif self.codename == "Swift": + print("Reading Swift simulation data is not implemented yet") + else: + print('Cannot process unknown code type. Call the read_param method with a valid code name. Valid options are "Swiftest", "Swifter", or "Swift".') + return + + diff --git a/python/swiftest.py b/python/swiftest/swiftest/swiftestio.py similarity index 69% rename from python/swiftest.py rename to python/swiftest/swiftest/swiftestio.py index db1ce2612..e7505a10f 100644 --- a/python/swiftest.py +++ b/python/swiftest/swiftest/swiftestio.py @@ -2,22 +2,8 @@ from scipy.io import FortranFile import xarray as xr from astroquery.jplhorizons import Horizons -import astropy.constants as const import datetime -# Constants in SI units -GC = np.longdouble(const.G.value) -AU2M = np.longdouble(const.au.value) -GMSunSI = np.longdouble(const.GM_sun.value) -MSun = np.longdouble(const.M_sun.value) -RSun = np.longdouble(const.R_sun.value) -JD2S = 86400 -YR2S = np.longdouble(365.25 * JD2S) -einsteinC = np.longdouble(299792458.0) -# Solar oblatenes values: From Mecheri et al. (2004), using Corbard (b) 2002 values (Table II) -J2Sun = np.longdouble(2.198e-7) -J4Sun = np.longdouble(-4.805e-9) - def read_swiftest_param(param_file_name): """ Reads in a Swiftest param.in file and saves it as a dictionary @@ -34,94 +20,95 @@ def read_swiftest_param(param_file_name): """ param = { 'VERSION': "! Swiftest parameter input file", - 'T0' : "0.0", - 'TSTOP' : "0.0", - 'DT' : "0.0", - 'PL_IN' : "", - 'TP_IN' : "", - 'CB_IN' : "", - 'IN_TYPE' : "ASCII", - 'ISTEP_OUT' : "-1", - 'BIN_OUT' : "bin.dat", - 'OUT_TYPE' : 'REAL8', - 'OUT_FORM' : "XV", - 'OUT_STAT' : "NEW", - 'ISTEP_DUMP' : "-1", - 'J2' : "0.0", - 'J4' : "0.0", - 'CHK_RMIN' : "-1.0", - 'CHK_RMAX' : "-1.0", - 'CHK_EJECT' : "-1.0", - 'CHK_QMIN' : "-1.0", - 'CHK_QMIN_COORD' : "HELIO", - 'CHK_QMIN_RANGE' : "", - 'QMIN_ALO' : "-1.0", - 'QMIN_AHI' : "-1.0", - 'ENC_OUT' : "", - 'MTINY' : "-1.0", - 'MU2KG' : "-1.0", - 'TU2S' : "-1.0", - 'DU2M' : "-1.0", - 'GU' : "-1.0", - 'EXTRA_FORCE' : "NO", - 'BIG_DISCARD' : "NO", - 'CHK_CLOSE' : "NO", - 'FRAGMENTATION' : "NO", - 'MTINY_SET' : "NO", - 'ROTATION' : "NO", - 'TIDES' : "NO", - 'ENERGY' : "NO", - 'GR' : "NO", - 'YARKOVSKY' : "NO", - 'YORP' : "NO", + 'T0': "0.0", + 'TSTOP': "0.0", + 'DT': "0.0", + 'PL_IN': "", + 'TP_IN': "", + 'CB_IN': "", + 'IN_TYPE': "ASCII", + 'ISTEP_OUT': "-1", + 'BIN_OUT': "bin.dat", + 'OUT_TYPE': 'REAL8', + 'OUT_FORM': "XV", + 'OUT_STAT': "NEW", + 'ISTEP_DUMP': "-1", + 'J2': "0.0", + 'J4': "0.0", + 'CHK_RMIN': "-1.0", + 'CHK_RMAX': "-1.0", + 'CHK_EJECT': "-1.0", + 'CHK_QMIN': "-1.0", + 'CHK_QMIN_COORD': "HELIO", + 'CHK_QMIN_RANGE': "", + 'QMIN_ALO': "-1.0", + 'QMIN_AHI': "-1.0", + 'ENC_OUT': "", + 'MTINY': "-1.0", + 'MU2KG': "-1.0", + 'TU2S': "-1.0", + 'DU2M': "-1.0", + 'GU': "-1.0", + 'EXTRA_FORCE': "NO", + 'BIG_DISCARD': "NO", + 'CHK_CLOSE': "NO", + 'FRAGMENTATION': "NO", + 'MTINY_SET': "NO", + 'ROTATION': "NO", + 'TIDES': "NO", + 'ENERGY': "NO", + 'GR': "NO", + 'YARKOVSKY': "NO", + 'YORP': "NO", } - + # Read param.in file - print(f'Reading Swiftest file {param_file_name}' ) + print(f'Reading Swiftest file {param_file_name}') with open(param_file_name, 'r') as f: for line in f.readlines(): fields = line.split() if len(fields) > 0: for key in param: if (key == fields[0].upper()): param[key] = fields[1] - #Special case of CHK_QMIN_RANGE requires a second input + # Special case of CHK_QMIN_RANGE requires a second input if fields[0].upper() == 'CHK_QMIN_RANGE': param['QMIN_ALO'] = fields[1] param['QMIN_AHI'] = fields[2] param['CHK_QMIN_RANGE'] = f"{fields[1]} {fields[2]}" - - param['ISTEP_OUT'] = int(param['ISTEP_OUT']) + + param['ISTEP_OUT'] = int(param['ISTEP_OUT']) param['ISTEP_DUMP'] = int(param['ISTEP_DUMP']) - param['T0'] = float(param['T0'].replace('d','E').replace('D','E')) - param['TSTOP'] = float(param['TSTOP'].replace('d','E').replace('D','E')) - param['DT'] = float(param['DT'].replace('d','E').replace('D','E')) - param['J2'] = float(param['J2'].replace('d','E').replace('D','E')) - param['J4'] = float(param['J4'].replace('d','E').replace('D','E')) - param['CHK_RMIN'] = float(param['CHK_RMIN'].replace('d','E').replace('D','E')) - param['CHK_RMAX'] = float(param['CHK_RMAX'].replace('d','E').replace('D','E')) - param['CHK_EJECT'] = float(param['CHK_EJECT'].replace('d','E').replace('D','E')) - param['CHK_QMIN'] = float(param['CHK_QMIN'].replace('d','E').replace('D','E')) - param['QMIN_ALO'] = float(param['QMIN_ALO'].replace('d','E').replace('D','E')) - param['QMIN_AHI'] = float(param['QMIN_AHI'].replace('d','E').replace('D','E')) - param['MTINY'] = float(param['MTINY'].replace('d','E').replace('D','E')) - param['DU2M'] = float(param['DU2M'].replace('d','E').replace('D','E')) - param['MU2KG'] = float(param['MU2KG'].replace('d','E').replace('D','E')) - param['TU2S'] = float(param['TU2S'].replace('d','E').replace('D','E')) + param['T0'] = float(param['T0'].replace('d', 'E').replace('D', 'E')) + param['TSTOP'] = float(param['TSTOP'].replace('d', 'E').replace('D', 'E')) + param['DT'] = float(param['DT'].replace('d', 'E').replace('D', 'E')) + param['J2'] = float(param['J2'].replace('d', 'E').replace('D', 'E')) + param['J4'] = float(param['J4'].replace('d', 'E').replace('D', 'E')) + param['CHK_RMIN'] = float(param['CHK_RMIN'].replace('d', 'E').replace('D', 'E')) + param['CHK_RMAX'] = float(param['CHK_RMAX'].replace('d', 'E').replace('D', 'E')) + param['CHK_EJECT'] = float(param['CHK_EJECT'].replace('d', 'E').replace('D', 'E')) + param['CHK_QMIN'] = float(param['CHK_QMIN'].replace('d', 'E').replace('D', 'E')) + param['QMIN_ALO'] = float(param['QMIN_ALO'].replace('d', 'E').replace('D', 'E')) + param['QMIN_AHI'] = float(param['QMIN_AHI'].replace('d', 'E').replace('D', 'E')) + param['MTINY'] = float(param['MTINY'].replace('d', 'E').replace('D', 'E')) + param['DU2M'] = float(param['DU2M'].replace('d', 'E').replace('D', 'E')) + param['MU2KG'] = float(param['MU2KG'].replace('d', 'E').replace('D', 'E')) + param['TU2S'] = float(param['TU2S'].replace('d', 'E').replace('D', 'E')) param['EXTRA_FORCE'] = param['EXTRA_FORCE'].upper() param['BIG_DISCARD'] = param['BIG_DISCARD'].upper() - param['CHK_CLOSE'] = param['CHK_CLOSE'].upper() + param['CHK_CLOSE'] = param['CHK_CLOSE'].upper() param['FRAGMENTATION'] = param['FRAGMENTATION'].upper() - param['ROTATION'] = param['ROTATION'].upper() - param['TIDES'] = param['TIDES'].upper() - param['ENERGY'] = param['ENERGY'].upper() - param['GR'] = param['GR'].upper() - param['YORP'] = param['YORP'].upper() - - param['GU'] = GC / (param['DU2M']**3 / (param['MU2KG'] * param['TU2S']**2)) - param['INV_C2'] = einsteinC * param['TU2S'] / param['DU2M'] - param['INV_C2'] = param['INV_C2']**(-2) + param['ROTATION'] = param['ROTATION'].upper() + param['TIDES'] = param['TIDES'].upper() + param['ENERGY'] = param['ENERGY'].upper() + param['GR'] = param['GR'].upper() + param['YORP'] = param['YORP'].upper() + + param['GU'] = GC / (param['DU2M'] ** 3 / (param['MU2KG'] * param['TU2S'] ** 2)) + param['INV_C2'] = einsteinC * param['TU2S'] / param['DU2M'] + param['INV_C2'] = param['INV_C2'] ** (-2) return param + def read_swifter_param(param_file_name): """ Reads in a Swifter param.in file and saves it as a dictionary @@ -138,38 +125,37 @@ def read_swifter_param(param_file_name): """ param = { 'VERSION': "! Swifter parameter input file", - 'T0' : "0.0", - 'TSTOP' : "0.0", - 'DT' : "0.0", - 'PL_IN' : "", - 'TP_IN' : "", - 'IN_TYPE' : "ASCII", - 'ISTEP_OUT' : "-1", - 'BIN_OUT' : "bin.dat", - 'OUT_TYPE' : "REAL8", - 'OUT_FORM' : "XV", - 'OUT_STAT' : "NEW", - 'ISTEP_DUMP' : "-1", - 'J2' : "0.0", - 'J4' : "0.0", - 'CHK_CLOSE' : 'NO', - 'CHK_RMIN' : "-1.0", - 'CHK_RMAX' : "-1.0", - 'CHK_EJECT' : "-1.0", - 'CHK_QMIN' : "-1.0", - 'CHK_QMIN_COORD' : "HELIO", - 'CHK_QMIN_RANGE' : "", - 'QMIN_ALO' : "-1.0", - 'QMIN_AHI' : "-1.0", - 'ENC_OUT' : "", - 'EXTRA_FORCE' : 'NO', - 'BIG_DISCARD' : 'NO', - 'RHILL_PRESENT' : 'NO', - 'GR' : 'NO', - 'C2' : "-1.0", + 'T0': "0.0", + 'TSTOP': "0.0", + 'DT': "0.0", + 'PL_IN': "", + 'TP_IN': "", + 'IN_TYPE': "ASCII", + 'ISTEP_OUT': "-1", + 'BIN_OUT': "bin.dat", + 'OUT_TYPE': "REAL8", + 'OUT_FORM': "XV", + 'OUT_STAT': "NEW", + 'ISTEP_DUMP': "-1", + 'J2': "0.0", + 'J4': "0.0", + 'CHK_CLOSE': 'NO', + 'CHK_RMIN': "-1.0", + 'CHK_RMAX': "-1.0", + 'CHK_EJECT': "-1.0", + 'CHK_QMIN': "-1.0", + 'CHK_QMIN_COORD': "HELIO", + 'CHK_QMIN_RANGE': "", + 'QMIN_ALO': "-1.0", + 'QMIN_AHI': "-1.0", + 'ENC_OUT': "", + 'EXTRA_FORCE': 'NO', + 'BIG_DISCARD': 'NO', + 'RHILL_PRESENT': 'NO', + 'GR': 'NO', + 'C2': "-1.0", } - # Read param.in file print(f'Reading Swifter file {param_file_name}') with open(param_file_name, 'r') as f: @@ -178,35 +164,36 @@ def read_swifter_param(param_file_name): if len(fields) > 0: for key in param: if (key == fields[0].upper()): param[key] = fields[1] - #Special case of CHK_QMIN_RANGE requires a second input + # Special case of CHK_QMIN_RANGE requires a second input if fields[0].upper() == 'CHK_QMIN_RANGE': param['QMIN_ALO'] = fields[1] param['QMIN_AHI'] = fields[2] param['CHK_QMIN_RANGE'] = f"{fields[1]} {fields[2]}" - - param['ISTEP_OUT'] = int(param['ISTEP_OUT']) + + param['ISTEP_OUT'] = int(param['ISTEP_OUT']) param['ISTEP_DUMP'] = int(param['ISTEP_DUMP']) - param['T0'] = float(param['T0'].replace('d','E').replace('D','E')) - param['TSTOP'] = float(param['TSTOP'].replace('d','E').replace('D','E')) - param['DT'] = float(param['DT'].replace('d','E').replace('D','E')) - param['J2'] = float(param['J2'].replace('d','E').replace('D','E')) - param['J4'] = float(param['J4'].replace('d','E').replace('D','E')) - param['CHK_RMIN'] = float(param['CHK_RMIN'].replace('d','E').replace('D','E')) - param['CHK_RMAX'] = float(param['CHK_RMAX'].replace('d','E').replace('D','E')) - param['CHK_EJECT'] = float(param['CHK_EJECT'].replace('d','E').replace('D','E')) - param['CHK_QMIN'] = float(param['CHK_QMIN'].replace('d','E').replace('D','E')) - param['QMIN_ALO'] = float(param['QMIN_ALO'].replace('d','E').replace('D','E')) - param['QMIN_AHI'] = float(param['QMIN_AHI'].replace('d','E').replace('D','E')) - param['C2'] = float(param['C2'].replace('d','E').replace('D','E')) - param['INV_C2'] = param['C2'] + param['T0'] = float(param['T0'].replace('d', 'E').replace('D', 'E')) + param['TSTOP'] = float(param['TSTOP'].replace('d', 'E').replace('D', 'E')) + param['DT'] = float(param['DT'].replace('d', 'E').replace('D', 'E')) + param['J2'] = float(param['J2'].replace('d', 'E').replace('D', 'E')) + param['J4'] = float(param['J4'].replace('d', 'E').replace('D', 'E')) + param['CHK_RMIN'] = float(param['CHK_RMIN'].replace('d', 'E').replace('D', 'E')) + param['CHK_RMAX'] = float(param['CHK_RMAX'].replace('d', 'E').replace('D', 'E')) + param['CHK_EJECT'] = float(param['CHK_EJECT'].replace('d', 'E').replace('D', 'E')) + param['CHK_QMIN'] = float(param['CHK_QMIN'].replace('d', 'E').replace('D', 'E')) + param['QMIN_ALO'] = float(param['QMIN_ALO'].replace('d', 'E').replace('D', 'E')) + param['QMIN_AHI'] = float(param['QMIN_AHI'].replace('d', 'E').replace('D', 'E')) + param['C2'] = float(param['C2'].replace('d', 'E').replace('D', 'E')) + param['INV_C2'] = param['C2'] param['EXTRA_FORCE'] = param['EXTRA_FORCE'].upper() param['BIG_DISCARD'] = param['BIG_DISCARD'].upper() - param['CHK_CLOSE'] = param['CHK_CLOSE'].upper() + param['CHK_CLOSE'] = param['CHK_CLOSE'].upper() param['RHILL_PRESENT'] = param['RHILL_PRESENT'].upper() - param['GR'] = param['GR'].upper() - + param['GR'] = param['GR'].upper() + return param + def read_swift_param(param_file_name): """ Reads in a Swift param.in file and saves it as a dictionary @@ -246,51 +233,56 @@ def read_swift_param(param_file_name): # Read param.in file print(f'Reading Swift file {param_file_name}') - with open(param_file_name, 'r') as f: - line = f.readline().split() - for i, l in enumerate(line): - line[i] = l.replace('d','E').replace('D','E') - param['T0'] = float(line[0]) - param['TSTOP'] = float(line[1]) - param['DT'] = float(line[2]) - line = f.readline().split() - for i, l in enumerate(line): - line[i] = l.replace('d','E').replace('D','E') - param['DTOUT'] = float(line[0]) - param['DTDUMP'] = float(line[1]) - line = f.readline().split() - param['L1'] = line[0].upper() - param['L2'] = line[1].upper() - param['L3'] = line[2].upper() - param['L4'] = line[3].upper() - param['L5'] = line[4].upper() - param['L6'] = line[5].upper() - line = f.readline().split() - for i, l in enumerate(line): - line[i] = l.replace('d','E').replace('D','E') - param['RMIN'] = float(line[0]) - param['RMAX'] = float(line[1]) - param['RMAXU'] = float(line[2]) - param['QMIN'] = float(line[3]) - param['LCLOSE'] = line[4].upper() - param['BINARY_OUTPUTFILE'] = f.readline().strip() - param['STATUS_FLAG_FOR_OPEN_STATEMENTS'] = f.readline().strip().upper() - + try: + with open(param_file_name, 'r') as f: + line = f.readline().split() + for i, l in enumerate(line): + line[i] = l.replace('d', 'E').replace('D', 'E') + param['T0'] = float(line[0]) + param['TSTOP'] = float(line[1]) + param['DT'] = float(line[2]) + line = f.readline().split() + for i, l in enumerate(line): + line[i] = l.replace('d', 'E').replace('D', 'E') + param['DTOUT'] = float(line[0]) + param['DTDUMP'] = float(line[1]) + line = f.readline().split() + param['L1'] = line[0].upper() + param['L2'] = line[1].upper() + param['L3'] = line[2].upper() + param['L4'] = line[3].upper() + param['L5'] = line[4].upper() + param['L6'] = line[5].upper() + line = f.readline().split() + for i, l in enumerate(line): + line[i] = l.replace('d', 'E').replace('D', 'E') + param['RMIN'] = float(line[0]) + param['RMAX'] = float(line[1]) + param['RMAXU'] = float(line[2]) + param['QMIN'] = float(line[3]) + param['LCLOSE'] = line[4].upper() + param['BINARY_OUTPUTFILE'] = f.readline().strip() + param['STATUS_FLAG_FOR_OPEN_STATEMENTS'] = f.readline().strip().upper() + except IOError: + print(f"{param_file_name} not found.") + return param -def write_param(param, param_file_name): +def write_swift_param(param, param_file_name): + outfile = open(param_file_name, 'w') + print(param['T0'], param['TSTOP'], param['DT'], file=outfile) + print(param['DTOUT'], param['DTDUMP'], file=outfile) + print(param['L1'], param['L2'], param['L3'], param['L4'], param['L5'], param['L6'], file=outfile) + print(param['RMIN'], param['RMAX'], param['RMAXU'], param['QMIN'], param['LCLOSE'], file=outfile) + print(param['BINARY_OUTPUTFILE'], file=outfile) + print(param['STATUS_FLAG_FOR_OPEN_STATEMENTS'], file=outfile) + outfile.close() + return + +def write_labeled_param(param, param_file_name): outfile = open(param_file_name, 'w') - codename = param['VERSION'].split()[1] - if codename == "Swifter" or codename == "Swiftest": - for key,val in param.items(): - print(f"{key:<16} {val}",file=outfile) - elif codename == "Swift": - print(param['T0'],param['TSTOP'],param['DT'], file=outfile) - print(param['DTOUT'], param['DTDUMP'], file=outfile) - print(param['L1'],param['L2'],param['L3'],param['L4'],param['L5'],param['L6'], file=outfile) - print(param['RMIN'],param['RMAX'],param['RMAXU'],param['QMIN'],param['LCLOSE'], file=outfile) - print(param['BINARY_OUTPUTFILE'], file=outfile) - print(param['STATUS_FLAG_FOR_OPEN_STATEMENTS'], file=outfile) + for key, val in param.items(): + print(f"{key:<16} {val}", file=outfile) outfile.close() return @@ -327,13 +319,13 @@ def swifter_stream(f, param): while True: # Loop until you read the end of file try: # Read single-line header - record = f.read_record(' 0: for i in range(ntp): - record = f.read_record(' 0: - pvec = np.vstack([p1,p2,p3,p4,p5,p6,Mpl,Rpl]) + pvec = np.vstack([p1, p2, p3, p4, p5, p6, Mpl, Rpl]) else: - pvec = np.empty((8,0)) + pvec = np.empty((8, 0)) plid = np.empty(0) if ntp > 0: - tvec = np.vstack([t1,t2,t3,t4,t5,t6]) + tvec = np.vstack([t1, t2, t3, t4, t5, t6]) else: - tvec = np.empty((6,0)) + tvec = np.empty((6, 0)) tpid = np.empty(0) - cvec = np.array([Mcb,Rcb,J2cb,J4cb]) + cvec = np.array([Mcb, Rcb, J2cb, J4cb]) if param['ROTATION'] == 'YES': cvec = np.vstack([cvec, Ipcbx, Ipcby, Ipcbz, rotcbx, rotcby, rotcbz]) if npl > 0: pvec = np.vstack([pvec, Ipplx, Ipply, Ipplz, rotplx, rotply, rotplz]) if param['TIDES'] == 'YES': - cvec = np.vstack([cvec,k2cb,Qcb]) + cvec = np.vstack([cvec, k2cb, Qcb]) if npl > 0: - pvec = np.vstack([pvec,k2pl,Qpl]) + pvec = np.vstack([pvec, k2pl, Qpl]) yield t, cbid, cvec.T, clab, \ npl, plid, pvec.T, plab, \ ntp, tpid, tvec.T, tlab + def swifter2xr(param): - dims = ['time','id', 'vec'] + """ + Converts a Swifter binary data file into an xarray DataSet. + + Parameters + ---------- + param : dict + Swifter parameters + + Returns + ------- + xarray dataset + """ + dims = ['time', 'id', 'vec'] pl = [] tp = [] with FortranFile(param['BIN_OUT'], 'r') as f: for t, npl, plid, pvec, plab, \ - ntp, tpid, tvec, tlab in swifter_stream(f, param): - - #Prepare frames by adding an extra axis for the time coordinate + ntp, tpid, tvec, tlab in swifter_stream(f, param): + # Prepare frames by adding an extra axis for the time coordinate plframe = np.expand_dims(pvec, axis=0) tpframe = np.expand_dims(tvec, axis=0) - #Create xarray DataArrays out of each body type - plxr = xr.DataArray(plframe, dims = dims, coords = {'time' : t, 'id' : plid, 'vec' : plab}) - tpxr = xr.DataArray(tpframe, dims = dims, coords = {'time' : t, 'id' : tpid, 'vec' : tlab}) - + # Create xarray DataArrays out of each body type + plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab}) + tpxr = xr.DataArray(tpframe, dims=dims, coords={'time': t, 'id': tpid, 'vec': tlab}) + pl.append(plxr) tp.append(tpxr) - + plda = xr.concat(pl, dim='time') tpda = xr.concat(tp, dim='time') - + plds = plda.to_dataset(dim='vec') tpds = tpda.to_dataset(dim='vec') ds = xr.combine_by_coords([plds, tpds]) return ds + def swiftest2xr(param): """ Converts a Swiftest binary data file into an xarray DataSet. @@ -564,46 +570,46 @@ def swiftest2xr(param): Parameters ---------- param : dict - Swiftest paramuration parameters + Swiftest parameters Returns ------- xarray dataset """ - - dims = ['time','id', 'vec'] + + dims = ['time', 'id', 'vec'] cb = [] pl = [] tp = [] with FortranFile(param['BIN_OUT'], 'r') as f: for t, cbid, cvec, clab, \ - npl, plid, pvec, plab, \ - ntp, tpid, tvec, tlab in swiftest_stream(f, param): - - #Prepare frames by adding an extra axis for the time coordinate + npl, plid, pvec, plab, \ + ntp, tpid, tvec, tlab in swiftest_stream(f, param): + # Prepare frames by adding an extra axis for the time coordinate cbframe = np.expand_dims(cvec, axis=0) plframe = np.expand_dims(pvec, axis=0) tpframe = np.expand_dims(tvec, axis=0) - #Create xarray DataArrays out of each body type - cbxr = xr.DataArray(cbframe, dims = dims, coords = {'time' : t, 'id' : cbid, 'vec' : clab}) - plxr = xr.DataArray(plframe, dims = dims, coords = {'time' : t, 'id' : plid, 'vec' : plab}) - tpxr = xr.DataArray(tpframe, dims = dims, coords = {'time' : t, 'id' : tpid, 'vec' : tlab}) - + # Create xarray DataArrays out of each body type + cbxr = xr.DataArray(cbframe, dims=dims, coords={'time': t, 'id': cbid, 'vec': clab}) + plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab}) + tpxr = xr.DataArray(tpframe, dims=dims, coords={'time': t, 'id': tpid, 'vec': tlab}) + cb.append(cbxr) pl.append(plxr) tp.append(tpxr) - - cbda = xr.concat(cb,dim='time') - plda = xr.concat(pl,dim='time') - tpda = xr.concat(tp,dim='time') - - cbds = cbda.to_dataset(dim = 'vec') - plds = plda.to_dataset(dim = 'vec') - tpds = tpda.to_dataset(dim = 'vec') + + cbda = xr.concat(cb, dim='time') + plda = xr.concat(pl, dim='time') + tpda = xr.concat(tp, dim='time') + + cbds = cbda.to_dataset(dim='vec') + plds = plda.to_dataset(dim='vec') + tpds = tpda.to_dataset(dim='vec') ds = xr.combine_by_coords([cbds, plds, tpds]) return ds + def solar_system_pl(param, ephemerides_start_date): """ Initializes a Swiftest dataset containing the major planets of the Solar System at a particular data from JPL/Horizons @@ -631,7 +637,7 @@ def solar_system_pl(param, ephemerides_start_date): 'neptune': '8', 'plutocharon': '9' } - + # Planet MSun/M ratio MSun_over_Mpl = { 'mercury': np.longdouble(6023600.0), @@ -644,7 +650,7 @@ def solar_system_pl(param, ephemerides_start_date): 'neptune': np.longdouble(19412.24), 'plutocharon': np.longdouble(1.35e8) } - + # Planet radii in meters planetradius = { 'mercury': np.longdouble(2439.4e3), @@ -657,27 +663,27 @@ def solar_system_pl(param, ephemerides_start_date): 'neptune': np.longdouble(24622.e3), 'plutocharon': np.longdouble(1188.3e3) } - + # Unit conversion factors DCONV = AU2M / param['DU2M'] VCONV = (AU2M / JD2S) / (param['DU2M'] / param['TU2S']) THIRDLONG = np.longdouble(1.0) / np.longdouble(3.0) - + # Central body value vectors - GMcb = np.array([GMSunSI * param['TU2S']**2 / param['DU2M']**3]) + GMcb = np.array([GMSunSI * param['TU2S'] ** 2 / param['DU2M'] ** 3]) Rcb = np.array([RSun / param['DU2M']]) - J2RP2 = np.array([J2Sun * (RSun / param['DU2M'])**2]) - J4RP4 = np.array([J4Sun * (RSun / param['DU2M'])**4]) + J2RP2 = np.array([J2Sun * (RSun / param['DU2M']) ** 2]) + J4RP4 = np.array([J4Sun * (RSun / param['DU2M']) ** 4]) cbid = np.array([0]) cvec = np.vstack([GMcb, Rcb, J2RP2, J4RP4]) - + # Horizons date time internal variables tstart = datetime.date.fromisoformat(ephemerides_start_date) tstep = datetime.timedelta(days=1) tend = tstart + tstep ephemerides_end_date = tend.isoformat() ephemerides_step = '1d' - + pldata = {} p1 = [] p2 = [] @@ -688,7 +694,7 @@ def solar_system_pl(param, ephemerides_start_date): Rhill = [] Rpl = [] GMpl = [] - + # Fetch solar system ephemerides from Horizons for key, val in planetid.items(): pldata[key] = Horizons(id=val, id_type='majorbody', location='@sun', @@ -714,34 +720,35 @@ def solar_system_pl(param, ephemerides_start_date): # Generate planet value vectors plid = np.fromiter(planetid.values(), dtype=int) pvec = np.vstack([p1, p2, p3, p4, p5, p6, GMpl, Rpl]) - + dims = ['time', 'id', 'vec'] cb = [] pl = [] tp = [] t = np.array([0.0]) - + clab, plab, tlab = make_swiftest_labels(param) - #Prepare frames by adding an extra axis for the time coordinate + # Prepare frames by adding an extra axis for the time coordinate cbframe = np.expand_dims(cvec.T, axis=0) plframe = np.expand_dims(pvec.T, axis=0) - #Create xarray DataArrays out of each body type - cbxr = xr.DataArray(cbframe, dims = dims, coords = {'time' : t, 'id' : cbid, 'vec' : clab}) - plxr = xr.DataArray(plframe, dims = dims, coords = {'time' : t, 'id' : plid, 'vec' : plab}) - + # Create xarray DataArrays out of each body type + cbxr = xr.DataArray(cbframe, dims=dims, coords={'time': t, 'id': cbid, 'vec': clab}) + plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab}) + cb.append(cbxr) pl.append(plxr) - - cbda = xr.concat(cb,dim='time') - plda = xr.concat(pl,dim='time') - - cbds = cbda.to_dataset(dim = 'vec') - plds = plda.to_dataset(dim = 'vec') + + cbda = xr.concat(cb, dim='time') + plda = xr.concat(pl, dim='time') + + cbds = cbda.to_dataset(dim='vec') + plds = plda.to_dataset(dim='vec') ds = xr.combine_by_coords([cbds, plds]) return ds + def swiftest_xr2_infile(ds, param, framenum=-1): """ Writes a set of Swiftest input files from a single frame of a Swiftest xarray dataset @@ -764,12 +771,12 @@ def swiftest_xr2_infile(ds, param, framenum=-1): pl = frame.where(frame.id > 0, drop=True) pl = pl.where(np.invert(np.isnan(pl['Mass'])), drop=True).drop_vars(['J_2', 'J_4']) tp = frame.where(np.isnan(frame['Mass']), drop=True).drop_vars(['Mass', 'Radius', 'J_2', 'J_4']) - + GMSun = np.double(cb['Mass']) RSun = np.double(cb['Radius']) J2 = np.double(cb['J_2']) J4 = np.double(cb['J_4']) - + if param['IN_TYPE'] == 'ASCII': # Swiftest Central body file cbfile = open(param['CB_IN'], 'w') @@ -778,7 +785,7 @@ def swiftest_xr2_infile(ds, param, framenum=-1): print(J2, file=cbfile) print(J4, file=cbfile) cbfile.close() - + plfile = open(param['PL_IN'], 'w') print(pl.id.count().values, file=plfile) for i in pl.id: @@ -788,7 +795,7 @@ def swiftest_xr2_infile(ds, param, framenum=-1): print(pli['px'].values, pli['py'].values, pli['pz'].values, file=plfile) print(pli['vx'].values, pli['vy'].values, pli['vz'].values, file=plfile) plfile.close() - + # TP file tpfile = open(param['TP_IN'], 'w') print(tp.id.count().values, file=tpfile) @@ -807,10 +814,10 @@ def swiftest_xr2_infile(ds, param, framenum=-1): cbfile.write_record(np.double(J2)) cbfile.write_record(np.double(J4)) cbfile.close() - + plfile = FortranFile(swiftest_pl, 'w') plfile.write_record(npl) - + plfile.write_record(plid) plfile.write_record(p_pl[0]) plfile.write_record(p_pl[1]) @@ -834,6 +841,7 @@ def swiftest_xr2_infile(ds, param, framenum=-1): else: print(f"{param['IN_TYPE']} is an unknown file type") + def swifter2swiftest(swifter_param, outparam): swiftest_param = swifter_param.copy() swiftest_param['PL_IN'] = 'pl.swiftest.in' @@ -843,9 +851,9 @@ def swifter2swiftest(swifter_param, outparam): print(f"Swiftest massive body file : {swiftest_param['PL_IN']}") print(f"Swiftest test particle file: {swiftest_param['TP_IN']}") print(f"Swiftest central body file : {swiftest_param['CB_IN']}") - + plnew = open(swiftest_param['PL_IN'], 'w') - + print(f'Writing out new PL file: {swiftest_param["PL_IN"]}') with open(swifter_param['PL_IN'], 'r') as plold: line = plold.readline() @@ -855,7 +863,7 @@ def swifter2swiftest(swifter_param, outparam): print(npl - 1, file=plnew) line = plold.readline() i_list = [i for i in line.split(" ") if i.strip()] - GMcb = float(i_list[1]) # Store central body GM for later + GMcb = float(i_list[1]) # Store central body GM for later line = plold.readline() # Ignore the two zero vector lines line = plold.readline() for n in range(1, npl): # Loop over planets @@ -863,12 +871,12 @@ def swifter2swiftest(swifter_param, outparam): i_list = [i for i in line.split(" ") if i.strip()] name = int(i_list[0]) GMpl = float(i_list[1]) - print(name, GMpl,file=plnew) + print(name, GMpl, file=plnew) if swifter_param['CHK_CLOSE'] == 'YES': line = plold.readline() i_list = [i for i in line.split(" ") if i.strip()] plrad = float(i_list[0]) - print(plrad,file=plnew) + print(plrad, file=plnew) line = plold.readline() i_list = [i for i in line.split(" ") if i.strip()] xh = float(i_list[0]) @@ -881,11 +889,11 @@ def swifter2swiftest(swifter_param, outparam): vy = float(i_list[1]) vz = float(i_list[2]) print(vx, vy, vz, file=plnew) - + plold.close() plnew.close() tpnew = open(swiftest_param['TP_IN'], 'w') - + print(f'Writing out new TP file: {swiftest_param["TP_IN"]}') with open(swifter_param['TP_IN'], 'r') as tpold: line = tpold.readline() @@ -910,10 +918,10 @@ def swifter2swiftest(swifter_param, outparam): vy = float(i_list[1]) vz = float(i_list[2]) print(vx, vy, vz, file=tpnew) - + tpold.close() tpnew.close() - + print(f"\nCentral body G*M = {GMcb}\n") print("Select the unit system to use:") print("1) MSun-AU-year") @@ -936,12 +944,12 @@ def swifter2swiftest(swifter_param, outparam): swiftest_param['MU2KG'] = MSun swiftest_param['DU2M'] = AU2M swiftest_param['TU2S'] = YR2S - elif unit_type == 2: # MSun-AU-day + elif unit_type == 2: # MSun-AU-day print("Unit system is MSun-AU-day") swiftest_param['MU2KG'] = MSun swiftest_param['DU2M'] = AU2M swiftest_param['TU2S'] = JD2S - elif unit_type == 3: # SI: kg-m-s + elif unit_type == 3: # SI: kg-m-s print("Unit system is SI: kg-m-s") swiftest_param['MU2KG'] = 1.0 swiftest_param['DU2M'] = 1.0 @@ -959,7 +967,7 @@ def swifter2swiftest(swifter_param, outparam): swiftest_param['TU2S'] = input("Time unit in seconds: ") GU = GC / (swiftest_param['DU2M'] ** 3 / (swiftest_param['MU2KG'] * swiftest_param['TU2S'] ** 2)) print(f"Central body mass: {GMcb / GU} MU ({(GMcb / GU) * swiftest_param['MU2KG']} kg)") - + print("Set central body radius:") print(f"1) Use Swifter parameter value of CHK_RMIN = {swifter_param['CHK_RMIN']}") print(f"2) Set value manually") @@ -977,18 +985,18 @@ def swifter2swiftest(swifter_param, outparam): cbrad = swifter_param['CHK_RMIN'] elif cbrad_type == 2: cbrad = input("Enter radius of central body in simulation Distance Units: ") - + print(f'Writing out new CB file: {swiftest_param["CB_IN"]}') # Write out new central body file cbnew = open(swiftest_param['CB_IN'], 'w') - + print(GMcb, file=cbnew) print(cbrad, file=cbnew) print(swifter_param['J2'], file=cbnew) print(swifter_param['J4'], file=cbnew) - + cbnew.close() - + # Remove the unneeded parameters swiftest_param.pop('INV_C2', None) swiftest_param.pop('C2', None) @@ -998,10 +1006,3 @@ def swifter2swiftest(swifter_param, outparam): swiftest_param.pop('J4', None) swiftest_param.pop('RHILL_PRESENT', None) return swiftest_param - -if __name__ == '__main__': - - workingdir = '/Users/daminton/git/swiftest/examples/swift_conversion/' - param_file_name = workingdir + 'param.in' - param = read_swift_param(param_file_name) - write_param(param,"test.in") \ No newline at end of file diff --git a/python/__init__.py b/python/swiftest/tests/__init__.py similarity index 100% rename from python/__init__.py rename to python/swiftest/tests/__init__.py diff --git a/examples/swift_conversion/param.in b/python/swiftest/tests/param.swift.in similarity index 100% rename from examples/swift_conversion/param.in rename to python/swiftest/tests/param.swift.in diff --git a/examples/swift_conversion/pl.in b/python/swiftest/tests/pl.in similarity index 100% rename from examples/swift_conversion/pl.in rename to python/swiftest/tests/pl.in diff --git a/examples/swift_conversion/swift.in b/python/swiftest/tests/swift.in similarity index 100% rename from examples/swift_conversion/swift.in rename to python/swiftest/tests/swift.in diff --git a/python/swiftest/tests/swift_reader.py b/python/swiftest/tests/swift_reader.py new file mode 100644 index 000000000..de62d59c1 --- /dev/null +++ b/python/swiftest/tests/swift_reader.py @@ -0,0 +1,6 @@ +import swiftest +inparam = "param.swift.in" +outparam = "param.swift.new" +print(f"Reading Swift parameter {inparam} and saving it to {outparam}") +sim = swiftest.Simulation(inparam, codename="Swift") +sim.write_param(outparam) \ No newline at end of file diff --git a/examples/swift_conversion/tp.in b/python/swiftest/tests/tp.in similarity index 100% rename from examples/swift_conversion/tp.in rename to python/swiftest/tests/tp.in