From 86575292ddc32247bbc7eb16a851b70739643926 Mon Sep 17 00:00:00 2001 From: Steven Rieder Date: Fri, 29 Nov 2024 19:14:26 +0100 Subject: [PATCH] Updates --- src/sunbather/__init__.py | 102 ++++++++++++------ src/sunbather/construct_parker.py | 29 +++-- .../data/{ => workingdir}/planets.txt | 0 src/sunbather/install_cloudy.py | 16 ++- 4 files changed, 96 insertions(+), 51 deletions(-) rename src/sunbather/data/{ => workingdir}/planets.txt (100%) diff --git a/src/sunbather/__init__.py b/src/sunbather/__init__.py index 4542efc..dae299e 100644 --- a/src/sunbather/__init__.py +++ b/src/sunbather/__init__.py @@ -9,14 +9,17 @@ from sunbather.install_cloudy import GetCloudy -def check_cloudy(): +def check_cloudy(quiet=False, cloudy_version="23.01"): """ Checks if Cloudy executable exists, and if not, prompts to download and build it. + :quiet: bool, if True, does not ask for input + :cloudy_version: str, Cloudy version (default: "23.01", environment variable + CLOUDY_VERSION overrides this) """ try: - cloudyversion = os.environ["CLOUDY_VERSION"] + cloudy_version = os.environ["CLOUDY_VERSION"] except KeyError: - cloudyversion = "23.01" + pass sunbatherpath = os.path.dirname( os.path.abspath(__file__) ) # the absolute path where this code lives @@ -24,23 +27,23 @@ def check_cloudy(): # the path where Cloudy is installed cloudypath = os.environ["cloudy_path"] except KeyError: - cloudypath = f"{sunbatherpath}/cloudy/c{cloudyversion}" + cloudypath = f"{sunbatherpath}/cloudy/c{cloudy_version}" if not os.path.exists(f"{cloudypath}/source/cloudy.exe"): - q = input( - f"Cloudy not found and CLOUDY_PATH is not set. " - f"Do you want to install Cloudy {cloudyversion} now in the sunbather path? " - f"(y/n) " - ) - while q.lower() not in ["y", "n"]: - q = input("Please enter 'y' or 'n'") - if q == "n": - raise KeyError( - "Cloudy not found, and the environment variable 'CLOUDY_PATH' is not " - "set. Please set this variable in your .bashrc/.zshrc file " - "to the path where the Cloudy installation is located. " - "Do not point it to the /source/ subfolder, but to the main folder." + if not quiet: + q = input( + f"Cloudy not found and CLOUDY_PATH is not set. Do you want to install " + f"Cloudy {cloudy_version} now in the sunbather path? (y/n) " ) - installer = GetCloudy(version=cloudyversion) + while q.lower() not in ["y", "n"]: + q = input("Please enter 'y' or 'n'") + if q == "n": + raise KeyError( + "Cloudy not found, and the environment variable 'CLOUDY_PATH' is " + "not set. Please set this variable in your .bashrc/.zshrc file " + "to the path where the Cloudy installation is located. " + "Do not point it to the /source/ subfolder, but to the main folder." + ) + installer = GetCloudy(version=cloudy_version) installer.download() installer.extract() installer.compile() @@ -48,33 +51,62 @@ def check_cloudy(): installer.copy_data() -def make_workingdir(): +def make_workingdir(workingdir=None, quiet=False): """ Checks if the SUNBATHER_PROJECT_PATH environment variable has been set and asks for input if not. Also asks to copy the default files to the working dir. + + :workingdir: str, path to the working dir. If None, checks the + SUNBATHER_PROJECT_PATH environment variable, and asks for input if this is + not set. (default: None) + :quiet: bool, if True, does not ask for input (default: False) """ - try: - workingdir = os.environ["SUNBATHER_PROJECT_PATH"] - except KeyError: - workingdir = input("Enter the working dir for Sunbather: ") - q = input(f"Copy default files to the working dir ({workingdir})? (y/n) ") - while q.lower() not in ["y", "n"]: - q = input("Please enter 'y' or 'n': ") - if q == "n": - return + if workingdir is None: + try: + workingdir = os.environ["SUNBATHER_PROJECT_PATH"] + except KeyError: + if not quiet: + workingdir = input("Enter the working dir for Sunbather: ") + else: + # if quiet, use the current dir + workingdir = "./" + if not quiet: + q = input(f"Copy default files to the working dir ({workingdir})? (y/n) ") + while q.lower() not in ["y", "n"]: + q = input("Please enter 'y' or 'n': ") + if q == "n": + return sunbatherpath = f"{pathlib.Path(__file__).parent.resolve()}" - shutil.copytree( - f"{sunbatherpath}/data/workingdir", - workingdir, - ) + for file in os.listdir(f"{sunbatherpath}/data/workingdir"): + if not os.path.exists(f"{workingdir}/{file}"): + shutil.copyfile( + f"{sunbatherpath}/data/workingdir/{file}", + f"{workingdir}/{file}", + ) + else: + if not quiet: + print("File already exists! Overwrite?") + q = input("(y/n) ") + while q.lower() not in ["y", "n"]: + q = input("Please enter 'y' or 'n': ") + if q == "n": + continue + else: + continue + shutil.copyfile( + f"{sunbatherpath}/data/workingdir/{file}", + f"{workingdir}/{file}", + ) + + return -def firstrun(): +def firstrun(quiet=False, workingdir=None, cloudy_version="23.01"): """ Runs 'check_cloudy()' and 'make_workingdir()'. """ - check_cloudy() - make_workingdir() + check_cloudy(quiet=quiet, cloudy_version=cloudy_version) + make_workingdir(quiet=quiet, workingdir=workingdir) print("Sunbather is ready to go!") diff --git a/src/sunbather/construct_parker.py b/src/sunbather/construct_parker.py index a95d7c2..2c60c72 100644 --- a/src/sunbather/construct_parker.py +++ b/src/sunbather/construct_parker.py @@ -990,7 +990,7 @@ def run_g( p.join() -def new_argument_parser(args, **kwargs): +def new_argument_parser(): parser = argparse.ArgumentParser( description="Creates 1D Parker profile(s) using the p_winds code and Cloudy.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, @@ -1141,15 +1141,22 @@ def __call__(self, parser, namespace, values, option_string=None): action="store_true", help="neglect the stellar tidal gravity term", ) - args = parser.parse_args(args, **kwargs) + return parser - return args - -def main(args, **kwargs): +def main(**kwargs): + """ + Main function to construct a Parker profile. + """ t0 = time.time() - - args = new_argument_parser(args, **kwargs) + parser = new_argument_parser() + if not kwargs: + args = parser.parse_args(sys.argv[1:]) + else: + # This is a bit ugly, but it allows us to either call main directly or + # to call the script with command line arguments + print(f"-{key}={value}" for key, value in kwargs.items()) + args = parser.parse_args([f'-{key}={value}' for key, value in kwargs.items()]) if args.z is not None: zdict = tools.get_zdict(z=args.z, zelem=args.zelem) @@ -1174,18 +1181,18 @@ def main(args, **kwargs): if not os.path.isdir(projectpath + "/parker_profiles/" + args.plname + "/"): os.mkdir(projectpath + "/parker_profiles/" + args.plname) if not os.path.isdir( - projectpath + "/parker_profiles/" + args.plname + "/" + args.pdir + "/" + f"{projectpath}/parker_profiles/{args.plname}/{args.pdir}/" ): os.mkdir( - projectpath + "/parker_profiles/" + args.plname + "/" + args.pdir + "/" + f"{projectpath}/parker_profiles/{args.plname}/{args.pdir}/" ) if (args.fH is None) and ( not os.path.isdir( - projectpath + "/parker_profiles/" + args.plname + "/" + args.pdir + "/temp/" + f"{projectpath}/parker_profiles/{args.plname}/{args.pdir}/temp/" ) ): os.mkdir( - projectpath + "/parker_profiles/" + args.plname + "/" + args.pdir + "/temp" + f"{projectpath}/parker_profiles/{args.plname}/{args.pdir}/temp" ) if len(args.T) == 1 and len(args.Mdot) == 1: # then we run a single model diff --git a/src/sunbather/data/planets.txt b/src/sunbather/data/workingdir/planets.txt similarity index 100% rename from src/sunbather/data/planets.txt rename to src/sunbather/data/workingdir/planets.txt diff --git a/src/sunbather/install_cloudy.py b/src/sunbather/install_cloudy.py index d80e951..18e974e 100644 --- a/src/sunbather/install_cloudy.py +++ b/src/sunbather/install_cloudy.py @@ -73,16 +73,22 @@ def test(self): exited OK" at the end. """ os.chdir(f"{self.cloudypath}/c{self.version}/source/") - print( - 'Type "test" and hit return twice. ' - 'It should print "Cloudy exited OK" at the end.' - ) with subprocess.Popen( [ "./cloudy.exe", - ] + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, ) as p: + cloudy_output = p.communicate(input=b"test\n\n")[0] p.wait() + try: + assert b"Cloudy exited OK" in cloudy_output + except AssertionError: + print("Cloudy did not test OK...") + else: + print("Cloudy tested OK.") def copy_data(self): """