diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..b7a4aaf Binary files /dev/null and b/.DS_Store differ diff --git a/.coveragerc b/.coveragerc.bak similarity index 100% rename from .coveragerc rename to .coveragerc.bak diff --git a/multiplanet/multiplanet_original.py.bak b/multiplanet/multiplanet_original.py.bak new file mode 100644 index 0000000..4ab5035 --- /dev/null +++ b/multiplanet/multiplanet_original.py.bak @@ -0,0 +1,370 @@ +import argparse +import multiprocessing as mp +import os +import subprocess as sub +import sys + +import h5py +import numpy as np +from bigplanet.read import GetVplanetHelp +from bigplanet.process import DictToBP, GatherData + +# -------------------------------------------------------------------- + + +def GetSNames(in_files, sims): + # get system and the body names + body_names = [] + + for file in in_files: + # gets path to infile + full_path = os.path.join(sims[0], file) + # if the infile is the vpl.in, then get the system name + if "vpl.in" in file: + with open(full_path, "r") as vpl: + content = [line.strip().split() for line in vpl.readlines()] + for line in content: + if line: + if line[0] == "sSystemName": + system_name = line[1] + else: + with open(full_path, "r") as infile: + content = [line.strip().split() for line in infile.readlines()] + for line in content: + if line: + if line[0] == "sName": + body_names.append(line[1]) + + return system_name, body_names + + +def GetSims(folder_name): + """Pass it folder name where simulations are and returns list of simulation folders.""" + # gets the list of sims + sims = sorted( + [ + f.path + for f in os.scandir(os.path.abspath(folder_name)) + if f.is_dir() + ] + ) + return sims + + +def GetDir(vspace_file): + """Give it input file and returns name of folder where simulations are located.""" + + infiles = [] + # gets the folder name with all the sims + with open(vspace_file, "r") as vpl: + content = [line.strip().split() for line in vpl.readlines()] + for line in content: + if line: + if line[0] == "sDestFolder" or line[0] == "destfolder": + folder_name = line[1] + + if ( + line[0] == "sBodyFile" + or line[0] == "sPrimaryFile" + or line[0] == "file" + ): + infiles.append(line[1]) + if folder_name is None: + raise IOError( + "Name of destination folder not provided in file '%s'." + "Use syntax 'destfolder '" % vspace_file + ) + + if os.path.isdir(folder_name) == False: + print( + "ERROR: Folder", + folder_name, + "does not exist in the current directory.", + ) + exit() + + return folder_name, infiles + + +## parallel implementation of running vplanet over a directory ## +def parallel_run_planet(input_file, cores, quiet, verbose, bigplanet, force): + # gets the folder name with all the sims + folder_name, in_files = GetDir(input_file) + # gets the list of sims + sims = GetSims(folder_name) + # Get the SNames (sName and sSystemName) for the simuations + # Save the name of the log file + system_name, body_list = GetSNames(in_files, sims) + logfile = system_name + ".log" + # initalizes the checkpoint file + checkpoint_file = os.getcwd() + "/" + "." + folder_name + # checks if the files doesn't exist and if so then it creates it + if os.path.isfile(checkpoint_file) == False: + CreateCP(checkpoint_file, input_file, sims) + + # if it does exist, it checks for any 0's (sims that didn't complete) and + # changes them to -1 to be re-ran + else: + ReCreateCP( + checkpoint_file, input_file, verbose, sims, folder_name, force + ) + + lock = mp.Lock() + workers = [] + + master_hdf5_file = os.getcwd() + "/" + folder_name + ".bpa" + # with h5py.File(master_hdf5_file, "w") as Master: + for i in range(cores): + workers.append( + mp.Process( + target=par_worker, + args=( + checkpoint_file, + system_name, + body_list, + logfile, + in_files, + verbose, + lock, + bigplanet, + master_hdf5_file, + ), + ) + ) + for w in workers: + print("Starting worker") + w.start() + + for w in workers: + w.join() + + if bigplanet == False: + if os.path.isfile(master_hdf5_file) == True: + sub.run(["rm", master_hdf5_file]) + + +def CreateCP(checkpoint_file, input_file, sims): + with open(checkpoint_file, "w") as cp: + cp.write("Vspace File: " + os.getcwd() + "/" + input_file + "\n") + cp.write("Total Number of Simulations: " + str(len(sims)) + "\n") + for f in range(len(sims)): + cp.write(sims[f] + " " + "-1 \n") + cp.write("THE END \n") + + +def ReCreateCP(checkpoint_file, input_file, verbose, sims, folder_name, force): + if verbose: + print("WARNING: multi-planet checkpoint file already exists!") + + datalist = [] + with open(checkpoint_file, "r") as re: + for newline in re: + datalist.append(newline.strip().split()) + + for l in datalist: + if l[1] == "0": + l[1] = "-1" + if datalist[-1] != ["THE", "END"]: + lest = datalist[-2][0] + idx = sims.index(lest) + for f in range(idx + 2, len(sims)): + datalist.append([sims[f], "-1"]) + datalist.append(["THE", "END"]) + + with open(checkpoint_file, "w") as wr: + for newline in datalist: + wr.writelines(" ".join(newline) + "\n") + + if all(l[1] == "1" for l in datalist[2:-2]) == True: + print("All simulations have been ran") + + if force: + if verbose: + print("Deleting folder...") + os.remove(folder_name) + if verbose: + print("Deleting Checkpoint File...") + os.remove(checkpoint_file) + if verbose: + print("Recreating Checkpoint File...") + CreateCP(checkpoint_file, input_file, sims) + else: + exit() + + +## parallel worker to run vplanet ## +def par_worker( + checkpoint_file, + system_name, + body_list, + log_file, + in_files, + verbose, + lock, + bigplanet, + h5_file, +): + + while True: + + lock.acquire() + datalist = [] + if bigplanet == True: + data = {} + vplanet_help = GetVplanetHelp() + + with open(checkpoint_file, "r") as f: + for newline in f: + datalist.append(newline.strip().split()) + + folder = "" + + for l in datalist: + if l[1] == "-1": + folder = l[0] + l[1] = "0" + break + if not folder: + lock.release() + return + + with open(checkpoint_file, "w") as f: + for newline in datalist: + f.writelines(" ".join(newline) + "\n") + + lock.release() + + if verbose: + print(folder) + os.chdir(folder) + + # runs vplanet on folder and writes the output to the log file + with open("vplanet_log", "a+") as vplf: + vplanet = sub.Popen( + "vplanet vpl.in", + shell=True, + stdout=sub.PIPE, + stderr=sub.PIPE, + universal_newlines=True, + ) + return_code = vplanet.poll() + for line in vplanet.stderr: + vplf.write(line) + + for line in vplanet.stdout: + vplf.write(line) + + lock.acquire() + datalist = [] + + with open(checkpoint_file, "r") as f: + for newline in f: + datalist.append(newline.strip().split()) + + if return_code is None: + for l in datalist: + if l[0] == folder: + l[1] = "1" + break + if verbose: + print(folder, "completed") + if bigplanet == True: + with h5py.File(h5_file, "a") as Master: + group_name = folder.split("/")[-1] + if group_name not in Master: + data = GatherData( + data, + system_name, + body_list, + log_file, + in_files, + vplanet_help, + folder, + verbose, + ) + DictToBP( + data, + vplanet_help, + Master, + verbose, + group_name, + archive=True, + ) + else: + for l in datalist: + if l[0] == folder: + l[1] = "-1" + break + if verbose: + print(folder, "failed") + + with open(checkpoint_file, "w") as f: + for newline in datalist: + f.writelines(" ".join(newline) + "\n") + + lock.release() + + os.chdir("../../") + + +def Arguments(): + max_cores = mp.cpu_count() + parser = argparse.ArgumentParser( + description="Using multi-processing to run a large number of simulations" + ) + parser.add_argument( + "-c", + "--cores", + type=int, + default=max_cores, + help="The total number of processors used", + ) + parser.add_argument( + "-bp", + "--bigplanet", + action="store_true", + help="Runs bigplanet and creates the Bigplanet Archive file alongside running multiplanet", + ) + parser.add_argument( + "-f", + "--force", + action="store_true", + help="forces rerun of multi-planet if completed", + ) + + parser.add_argument("InputFile", help="name of the vspace file") + + # adds the quiet and verbose as mutually exclusive groups + group = parser.add_mutually_exclusive_group() + group.add_argument( + "-q", "--quiet", action="store_true", help="no output for multiplanet" + ) + group.add_argument( + "-v", + "--verbose", + action="store_true", + help="Prints out excess output for multiplanet", + ) + + args = parser.parse_args() + + try: + if sys.version_info >= (3, 0): + help = sub.getoutput("vplanet -h") + else: + help = sub.check_output(["vplanet", "-h"]) + except OSError: + raise Exception("Unable to call VPLANET. Is it in your PATH?") + + parallel_run_planet( + args.InputFile, + args.cores, + args.quiet, + args.verbose, + args.bigplanet, + args.force, + ) + + +if __name__ == "__main__": + Arguments() diff --git a/multiplanet/multiplanet_refactored.py b/multiplanet/multiplanet_refactored.py new file mode 100644 index 0000000..cc19266 --- /dev/null +++ b/multiplanet/multiplanet_refactored.py @@ -0,0 +1,556 @@ +import argparse +import multiprocessing as mp +import os +import subprocess as sub +import sys + +import h5py +import numpy as np +from bigplanet.read import GetVplanetHelp +from bigplanet.process import DictToBP, GatherData + +# -------------------------------------------------------------------- +# HELPER FUNCTIONS (extracted from BigPlanet architecture) +# -------------------------------------------------------------------- + + +def fnGetNextSimulation(sCheckpointFile, lockFile): + """ + Find and mark the next simulation to process from checkpoint file. + + Thread-safe with file locking. Reads checkpoint, finds first simulation + with status -1, marks it as 0 (in-progress), and returns the folder path. + + Parameters + ---------- + sCheckpointFile : str + Path to checkpoint file + lockFile : multiprocessing.Lock + Lock for thread-safe file access + + Returns + ------- + str or None + Absolute path to simulation folder, or None if all done + """ + lockFile.acquire() + listData = [] + + with open(sCheckpointFile, "r") as f: + for sLine in f: + listData.append(sLine.strip().split()) + + sFolder = "" + for listLine in listData: + if len(listLine) > 1 and listLine[1] == "-1": + sFolder = listLine[0] + listLine[1] = "0" + break + + if not sFolder: + lockFile.release() + return None + + with open(sCheckpointFile, "w") as f: + for listLine in listData: + f.writelines(" ".join(listLine) + "\n") + + lockFile.release() + return os.path.abspath(sFolder) + + +def fnMarkSimulationComplete(sCheckpointFile, sFolder, lockFile): + """ + Mark simulation as complete in checkpoint file. + + Thread-safe with file locking. Updates status from 0 or -1 to 1. + + Parameters + ---------- + sCheckpointFile : str + Path to checkpoint file + sFolder : str + Folder path to mark as complete + lockFile : multiprocessing.Lock + Lock for thread-safe file access + + Returns + ------- + None + """ + lockFile.acquire() + listData = [] + + with open(sCheckpointFile, "r") as f: + for sLine in f: + listData.append(sLine.strip().split()) + + for listLine in listData: + if len(listLine) > 1 and listLine[0] == sFolder: + listLine[1] = "1" + break + + with open(sCheckpointFile, "w") as f: + for listLine in listData: + f.writelines(" ".join(listLine) + "\n") + + lockFile.release() + + +def fnMarkSimulationFailed(sCheckpointFile, sFolder, lockFile): + """ + Mark simulation as failed in checkpoint file. + + Thread-safe with file locking. Updates status back to -1 for retry. + + Parameters + ---------- + sCheckpointFile : str + Path to checkpoint file + sFolder : str + Folder path to mark as failed + lockFile : multiprocessing.Lock + Lock for thread-safe file access + + Returns + ------- + None + """ + lockFile.acquire() + listData = [] + + with open(sCheckpointFile, "r") as f: + for sLine in f: + listData.append(sLine.strip().split()) + + for listLine in listData: + if len(listLine) > 1 and listLine[0] == sFolder: + listLine[1] = "-1" + break + + with open(sCheckpointFile, "w") as f: + for listLine in listData: + f.writelines(" ".join(listLine) + "\n") + + lockFile.release() + + +# -------------------------------------------------------------------- +# ORIGINAL FUNCTIONS (unchanged) +# -------------------------------------------------------------------- + + +def GetSNames(in_files, sims): + # get system and the body names + body_names = [] + + for file in in_files: + # gets path to infile + full_path = os.path.join(sims[0], file) + # if the infile is the vpl.in, then get the system name + if "vpl.in" in file: + with open(full_path, "r") as vpl: + content = [line.strip().split() for line in vpl.readlines()] + for line in content: + if line: + if line[0] == "sSystemName": + system_name = line[1] + else: + with open(full_path, "r") as infile: + content = [line.strip().split() for line in infile.readlines()] + for line in content: + if line: + if line[0] == "sName": + body_names.append(line[1]) + + return system_name, body_names + + +def GetSims(folder_name): + """Pass it folder name where simulations are and returns list of simulation folders.""" + # gets the list of sims + sims = sorted( + [ + f.path + for f in os.scandir(os.path.abspath(folder_name)) + if f.is_dir() + ] + ) + return sims + + +def GetDir(vspace_file): + """Give it input file and returns name of folder where simulations are located.""" + + infiles = [] + folder_name = None + # gets the folder name with all the sims + with open(vspace_file, "r") as vpl: + content = [line.strip().split() for line in vpl.readlines()] + for line in content: + if line: + if line[0] == "sDestFolder" or line[0] == "destfolder": + folder_name = line[1] + + if ( + line[0] == "sBodyFile" + or line[0] == "sPrimaryFile" + or line[0] == "file" + ): + infiles.append(line[1]) + if folder_name is None: + raise IOError( + "Name of destination folder not provided in file '%s'." + "Use syntax 'destfolder '" % vspace_file + ) + + if os.path.isdir(folder_name) == False: + print( + "ERROR: Folder", + folder_name, + "does not exist in the current directory.", + ) + exit() + + return folder_name, infiles + + +def CreateCP(checkpoint_file, input_file, sims): + with open(checkpoint_file, "w") as cp: + cp.write("Vspace File: " + os.getcwd() + "/" + input_file + "\n") + cp.write("Total Number of Simulations: " + str(len(sims)) + "\n") + for f in range(len(sims)): + cp.write(sims[f] + " " + "-1 \n") + cp.write("THE END \n") + + +def ReCreateCP(checkpoint_file, input_file, verbose, sims, folder_name, force): + if verbose: + print("WARNING: multi-planet checkpoint file already exists!") + + datalist = [] + with open(checkpoint_file, "r") as re: + for newline in re: + datalist.append(newline.strip().split()) + + for l in datalist: + if len(l) > 1 and l[1] == "0": + l[1] = "-1" + if datalist[-1] != ["THE", "END"]: + lest = datalist[-2][0] + idx = sims.index(lest) + for f in range(idx + 2, len(sims)): + datalist.append([sims[f], "-1"]) + datalist.append(["THE", "END"]) + + with open(checkpoint_file, "w") as wr: + for newline in datalist: + wr.writelines(" ".join(newline) + "\n") + + if all(len(l) > 1 and l[1] == "1" for l in datalist[2:-2]) == True: + print("All simulations have been ran") + + if force: + if verbose: + print("Deleting folder...") + os.remove(folder_name) + if verbose: + print("Deleting Checkpoint File...") + os.remove(checkpoint_file) + if verbose: + print("Recreating Checkpoint File...") + CreateCP(checkpoint_file, input_file, sims) + else: + exit() + + +# -------------------------------------------------------------------- +# REFACTORED WORKER (adopting BigPlanet architecture) +# -------------------------------------------------------------------- + + +def par_worker( + checkpoint_file, + system_name, + body_list, + log_file, + in_files, + verbose, + lock, + bigplanet, + h5_file, + vplanet_help, +): + """ + Worker process for running vplanet simulations. + + REFACTORED to adopt BigPlanet's architecture: + - Uses fnGetNextSimulation() for thread-safe checkpoint access + - GetVplanetHelp() passed as parameter (called once in main) + - Minimal critical sections (lock held only during file I/O) + - Proper subprocess return code handling (wait() not poll()) + - No os.chdir() calls (uses cwd parameter instead) + + Parameters + ---------- + checkpoint_file : str + Path to checkpoint file + system_name : str + Name of system + body_list : list + List of body names + log_file : str + Name of log file + in_files : list + List of input files + verbose : bool + Verbose output flag + lock : multiprocessing.Lock + Lock for thread-safe operations + bigplanet : bool + Create BigPlanet archive + h5_file : str + Path to HDF5 archive file + vplanet_help : dict or None + Vplanet help data (pre-fetched in main process) + + Returns + ------- + None + """ + while True: + # STEP 1: Get next simulation (with lock - minimal critical section) + sFolder = fnGetNextSimulation(checkpoint_file, lock) + if sFolder is None: + return # No more work + + if verbose: + print(f"Processing: {sFolder}") + + # STEP 2: Run vplanet simulation (NO LOCK - independent work) + vplanet_log_path = os.path.join(sFolder, "vplanet_log") + + with open(vplanet_log_path, "a+") as vplf: + vplanet = sub.Popen( + ["vplanet", "vpl.in"], + cwd=sFolder, + stdout=sub.PIPE, + stderr=sub.PIPE, + universal_newlines=True, + ) + # FIXED: Use communicate() to wait for completion and get output + stdout, stderr = vplanet.communicate() + + # Write output to log + vplf.write(stderr) + vplf.write(stdout) + + # FIXED: Check actual return code (not poll()) + return_code = vplanet.returncode + + # STEP 3: Process BigPlanet data if needed (NO LOCK - CPU-bound work) + if return_code == 0 and bigplanet and vplanet_help is not None: + # Gather simulation data + data = {} + data = GatherData( + data, + system_name, + body_list, + log_file, + in_files, + vplanet_help, + sFolder, + verbose, + ) + + # STEP 4: Write to HDF5 (WITH LOCK - minimal critical section) + lock.acquire() + try: + with h5py.File(h5_file, "a") as Master: + group_name = os.path.basename(sFolder) + if group_name not in Master: + DictToBP( + data, + vplanet_help, + Master, + verbose, + group_name, + archive=True, + ) + finally: + lock.release() + + # STEP 5: Update checkpoint (with lock) + if return_code == 0: + fnMarkSimulationComplete(checkpoint_file, sFolder, lock) + if verbose: + print(f"{sFolder} completed") + else: + fnMarkSimulationFailed(checkpoint_file, sFolder, lock) + if verbose: + print(f"{sFolder} failed with return code {return_code}") + + +# -------------------------------------------------------------------- +# REFACTORED MAIN FUNCTION +# -------------------------------------------------------------------- + + +def parallel_run_planet(input_file, cores, quiet, verbose, bigplanet, force): + """ + Run vplanet simulations in parallel. + + REFACTORED to fix BigPlanet deadlock: + - GetVplanetHelp() called ONCE in main process (not in workers) + - Passed to workers as immutable parameter + - No subprocess calls within multiprocessing context + + Parameters + ---------- + input_file : str + Vspace input file + cores : int + Number of CPU cores to use + quiet : bool + Suppress output + verbose : bool + Verbose output + bigplanet : bool + Create BigPlanet HDF5 archive + force : bool + Force rerun if already completed + + Returns + ------- + None + """ + # gets the folder name with all the sims + folder_name, in_files = GetDir(input_file) + # gets the list of sims + sims = GetSims(folder_name) + # Get the SNames (sName and sSystemName) for the simuations + # Save the name of the log file + system_name, body_list = GetSNames(in_files, sims) + logfile = system_name + ".log" + # initalizes the checkpoint file + checkpoint_file = os.getcwd() + "/" + "." + folder_name + # checks if the files doesn't exist and if so then it creates it + if os.path.isfile(checkpoint_file) == False: + CreateCP(checkpoint_file, input_file, sims) + + # if it does exist, it checks for any 0's (sims that didn't complete) and + # changes them to -1 to be re-ran + else: + ReCreateCP( + checkpoint_file, input_file, verbose, sims, folder_name, force + ) + + lock = mp.Lock() + workers = [] + + master_hdf5_file = os.getcwd() + "/" + folder_name + ".bpa" + + # CRITICAL FIX: Call GetVplanetHelp() ONCE in main process + # This is passed to workers instead of being called inside worker loop + if bigplanet: + vplanet_help = GetVplanetHelp() + else: + vplanet_help = None + + # Spawn worker processes + for i in range(cores): + workers.append( + mp.Process( + target=par_worker, + args=( + checkpoint_file, + system_name, + body_list, + logfile, + in_files, + verbose, + lock, + bigplanet, + master_hdf5_file, + vplanet_help, # PASSED as parameter + ), + ) + ) + + # Start all workers + for w in workers: + if verbose: + print("Starting worker") + w.start() + + # Wait for all workers to complete + for w in workers: + w.join() + + # Clean up HDF5 file if not using bigplanet + if bigplanet == False: + if os.path.isfile(master_hdf5_file) == True: + sub.run(["rm", master_hdf5_file]) + + +def Arguments(): + max_cores = mp.cpu_count() + parser = argparse.ArgumentParser( + description="Using multi-processing to run a large number of simulations" + ) + parser.add_argument( + "-c", + "--cores", + type=int, + default=max_cores, + help="The total number of processors used", + ) + parser.add_argument( + "-bp", + "--bigplanet", + action="store_true", + help="Runs bigplanet and creates the Bigplanet Archive file alongside running multiplanet", + ) + parser.add_argument( + "-f", + "--force", + action="store_true", + help="forces rerun of multi-planet if completed", + ) + + parser.add_argument("InputFile", help="name of the vspace file") + + # adds the quiet and verbose as mutually exclusive groups + group = parser.add_mutually_exclusive_group() + group.add_argument( + "-q", "--quiet", action="store_true", help="no output for multiplanet" + ) + group.add_argument( + "-v", + "--verbose", + action="store_true", + help="Prints out excess output for multiplanet", + ) + + args = parser.parse_args() + + try: + if sys.version_info >= (3, 0): + help = sub.getoutput("vplanet -h") + else: + help = sub.check_output(["vplanet", "-h"]) + except OSError: + raise Exception("Unable to call VPLANET. Is it in your PATH?") + + parallel_run_planet( + args.InputFile, + args.cores, + args.quiet, + args.verbose, + args.bigplanet, + args.force, + ) + + +if __name__ == "__main__": + Arguments() diff --git a/pyproject.toml b/pyproject.toml index b7b8828..579e142 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,3 +22,21 @@ exclude = ''' | dist )/ ''' + +[tool.coverage.run] +source = ["multiplanet"] +omit = [ + "multiplanet/multiplanet_refactored.py", + "*/tests/*", +] + +[tool.coverage.report] +precision = 2 +show_missing = true +exclude_lines = [ + "pragma: no cover", + "if __name__ == .__main__.:", +] + +[tool.coverage.html] +directory = "htmlcov" diff --git a/tests/fixtures/conftest.py b/tests/fixtures/conftest.py new file mode 100644 index 0000000..5805809 --- /dev/null +++ b/tests/fixtures/conftest.py @@ -0,0 +1,57 @@ +import pytest +import tempfile +import os + + +@pytest.fixture +def temp_vspace_file(): + """Create temporary vspace.in file with standard content.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f: + f.write("destfolder TestSims\n") + f.write("file vpl.in\n") + f.write("file earth.in\n") + yield f.name + os.unlink(f.name) + + +@pytest.fixture +def temp_checkpoint_file(): + """Create temporary checkpoint file with test data.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: + f.write("Vspace File: /test/path/vspace.in\n") + f.write("Total Number of Simulations: 3\n") + f.write("sim_001 -1\n") + f.write("sim_002 -1\n") + f.write("sim_003 -1\n") + f.write("THE END\n") + yield f.name + os.unlink(f.name) + + +@pytest.fixture +def temp_sim_directory(tmp_path): + """Create temporary directory with mock simulation folders.""" + sim_dir = tmp_path / "TestSims" + sim_dir.mkdir() + (sim_dir / "sim_001").mkdir() + (sim_dir / "sim_002").mkdir() + (sim_dir / "sim_003").mkdir() + return sim_dir + + +@pytest.fixture +def temp_vpl_in_file(tmp_path): + """Create temporary vpl.in file with system name.""" + vpl_file = tmp_path / "vpl.in" + vpl_file.write_text("sSystemName TestSystem\n") + return vpl_file + + +@pytest.fixture +def temp_body_in_file(tmp_path): + """Create temporary body.in file with body name.""" + def _create_body_file(body_name): + body_file = tmp_path / f"{body_name}.in" + body_file.write_text(f"sName {body_name}\n") + return body_file + return _create_body_file diff --git a/tests/unit/test_helpers.py b/tests/unit/test_helpers.py new file mode 100644 index 0000000..591c23c --- /dev/null +++ b/tests/unit/test_helpers.py @@ -0,0 +1,215 @@ +""" +Unit tests for multiplanet helper functions. + +Tests GetSNames, GetSims, and GetDir functions for parsing vspace files +and simulation directories. +""" + +import os +import tempfile +import pytest + +from multiplanet import multiplanet + + +class TestGetSNames: + """Tests for GetSNames function.""" + + def test_get_snames_single_body(self, tmp_path): + """Parse system name and single body name from input files.""" + # Create vpl.in with system name + vpl_file = tmp_path / "vpl.in" + vpl_file.write_text("sSystemName SolarSystem\n") + + # Create earth.in with body name + earth_file = tmp_path / "earth.in" + earth_file.write_text("sName Earth\n") + + # Call GetSNames + system_name, body_names = multiplanet.GetSNames( + ["vpl.in", "earth.in"], + [str(tmp_path)] + ) + + assert system_name == "SolarSystem" + assert body_names == ["Earth"] + + def test_get_snames_multiple_bodies(self, tmp_path): + """Parse system name and multiple body names.""" + # Create vpl.in + vpl_file = tmp_path / "vpl.in" + vpl_file.write_text("sSystemName MultiBody\n") + + # Create three body files + venus = tmp_path / "venus.in" + venus.write_text("sName Venus\n") + + earth = tmp_path / "earth.in" + earth.write_text("sName Earth\n") + + mars = tmp_path / "mars.in" + mars.write_text("sName Mars\n") + + # Call GetSNames + system_name, body_names = multiplanet.GetSNames( + ["vpl.in", "venus.in", "earth.in", "mars.in"], + [str(tmp_path)] + ) + + assert system_name == "MultiBody" + assert body_names == ["Venus", "Earth", "Mars"] + + def test_get_snames_order_preservation(self, tmp_path): + """Ensure body names maintain order from file list.""" + # Create vpl.in + vpl_file = tmp_path / "vpl.in" + vpl_file.write_text("sSystemName Test\n") + + # Create body files + mars = tmp_path / "mars.in" + mars.write_text("sName Mars\n") + + earth = tmp_path / "earth.in" + earth.write_text("sName Earth\n") + + venus = tmp_path / "venus.in" + venus.write_text("sName Venus\n") + + # Call with specific order: mars, earth, venus + system_name, body_names = multiplanet.GetSNames( + ["vpl.in", "mars.in", "earth.in", "venus.in"], + [str(tmp_path)] + ) + + # Verify order matches input order + assert body_names == ["Mars", "Earth", "Venus"] + + +class TestGetSims: + """Tests for GetSims function.""" + + def test_get_sims_returns_sorted_directories(self, tmp_path): + """Returns sorted list of simulation directories.""" + # Create directories in non-sorted order + (tmp_path / "sim_003").mkdir() + (tmp_path / "sim_001").mkdir() + (tmp_path / "sim_002").mkdir() + + # Call GetSims + sims = multiplanet.GetSims(str(tmp_path)) + + # Verify sorted order and absolute paths + assert len(sims) == 3 + assert sims[0].endswith("sim_001") + assert sims[1].endswith("sim_002") + assert sims[2].endswith("sim_003") + assert all(os.path.isabs(s) for s in sims) + + def test_get_sims_ignores_files(self, tmp_path): + """Ignores files, only returns directories.""" + # Create directories + (tmp_path / "sim_001").mkdir() + (tmp_path / "sim_002").mkdir() + + # Create files + (tmp_path / "vpl.in").write_text("test") + (tmp_path / "README.md").write_text("test") + + # Call GetSims + sims = multiplanet.GetSims(str(tmp_path)) + + # Verify only directories returned + assert len(sims) == 2 + assert all("sim_" in s for s in sims) + + def test_get_sims_empty_folder(self, tmp_path): + """Returns empty list for folder with no subdirs.""" + # Create some files but no directories + (tmp_path / "file.txt").write_text("test") + + # Call GetSims + sims = multiplanet.GetSims(str(tmp_path)) + + # Verify empty list + assert sims == [] + + +class TestGetDir: + """Tests for GetDir function.""" + + def test_get_dir_valid_vspace_file(self, tmp_path): + """Parse folder name and input files from vspace file.""" + # Create the destination folder first + dest_folder = tmp_path / "TestSims" + dest_folder.mkdir() + + # Create vspace.in file + vspace_file = tmp_path / "vspace.in" + vspace_file.write_text( + "destfolder TestSims\n" + "file vpl.in\n" + "sBodyFile earth.in\n" + "sPrimaryFile sun.in\n" + ) + + # Change to tmp_path so relative path works + original_cwd = os.getcwd() + os.chdir(tmp_path) + try: + # Call GetDir + folder_name, in_files = multiplanet.GetDir(str(vspace_file)) + + assert folder_name == "TestSims" + assert "vpl.in" in in_files + assert "earth.in" in in_files + assert "sun.in" in in_files + finally: + os.chdir(original_cwd) + + def test_get_dir_missing_destfolder(self, tmp_path): + """Raise IOError if destfolder not specified.""" + # Create vspace.in without destfolder + vspace_file = tmp_path / "vspace.in" + vspace_file.write_text( + "file vpl.in\n" + "sBodyFile earth.in\n" + ) + + # Call GetDir and expect IOError + with pytest.raises(IOError, match="destination folder not provided"): + multiplanet.GetDir(str(vspace_file)) + + def test_get_dir_folder_not_exists(self, tmp_path): + """Exit gracefully if destination folder missing.""" + # Create vspace.in with non-existent folder + vspace_file = tmp_path / "vspace.in" + vspace_file.write_text("destfolder NonExistentFolder\n") + + # Call GetDir and expect SystemExit + with pytest.raises(SystemExit): + multiplanet.GetDir(str(vspace_file)) + + def test_get_dir_alternate_syntax(self, tmp_path): + """Support both 'destfolder' and 'sDestFolder' syntax.""" + # Create the destination folder first + dest_folder = tmp_path / "AltSims" + dest_folder.mkdir() + + # Create vspace.in with sDestFolder syntax + vspace_file = tmp_path / "vspace.in" + vspace_file.write_text( + "sDestFolder AltSims\n" + "file vpl.in\n" + ) + + # Change to tmp_path so relative path works + original_cwd = os.getcwd() + os.chdir(tmp_path) + try: + # Call GetDir + folder_name, in_files = multiplanet.GetDir(str(vspace_file)) + + assert folder_name == "AltSims" + assert "vpl.in" in in_files + finally: + os.chdir(original_cwd) diff --git a/tests/unit/test_mpstatus.py b/tests/unit/test_mpstatus.py new file mode 100644 index 0000000..1292bad --- /dev/null +++ b/tests/unit/test_mpstatus.py @@ -0,0 +1,153 @@ +""" +Unit tests for mpstatus module. + +Tests the mpstatus function and CLI argument parsing. +""" + +import os +import tempfile +import pytest +from unittest import mock + +from multiplanet import mpstatus + + +class TestMpstatus: + """Tests for mpstatus function.""" + + def test_mpstatus_all_pending(self, tmp_path, capsys): + """Count all simulations as pending (status -1).""" + # Create vspace.in + vspace_file = tmp_path / "vspace.in" + vspace_file.write_text( + "Vspace File: /test/vspace.in\n" + "destfolder TestSims\n" + ) + + # Create checkpoint file with all pending + checkpoint_file = tmp_path / ".TestSims" + checkpoint_file.write_text( + "Vspace File: /test/vspace.in\n" + "Total Number of Simulations: 3\n" + "sim_001 -1\n" + "sim_002 -1\n" + "sim_003 -1\n" + "THE END\n" + ) + + # Change to tmp_path directory + original_cwd = os.getcwd() + os.chdir(tmp_path) + try: + mpstatus.mpstatus(str(vspace_file)) + captured = capsys.readouterr() + + assert "Number of Simulations completed: 0" in captured.out + assert "Number of Simulations in progress: 0" in captured.out + assert "Number of Simulations remaining: 3" in captured.out + finally: + os.chdir(original_cwd) + + def test_mpstatus_mixed_status(self, tmp_path, capsys): + """Count simulations with mixed statuses.""" + # Create vspace.in (mpstatus reads line [1], so need header line) + vspace_file = tmp_path / "vspace.in" + vspace_file.write_text("header line\ndestfolder MixedSims\n") + + # Create checkpoint with mixed statuses + checkpoint_file = tmp_path / ".MixedSims" + checkpoint_file.write_text( + "Vspace File: /test/vspace.in\n" + "Total Number of Simulations: 5\n" + "sim_001 1\n" + "sim_002 0\n" + "sim_003 -1\n" + "sim_004 -1\n" + "sim_005 1\n" + "THE END\n" + ) + + # Change to tmp_path directory + original_cwd = os.getcwd() + os.chdir(tmp_path) + try: + mpstatus.mpstatus(str(vspace_file)) + captured = capsys.readouterr() + + assert "Number of Simulations completed: 2" in captured.out + assert "Number of Simulations in progress: 1" in captured.out + assert "Number of Simulations remaining: 2" in captured.out + finally: + os.chdir(original_cwd) + + def test_mpstatus_no_checkpoint(self, tmp_path): + """Raise exception if checkpoint file missing.""" + # Create vspace.in but no checkpoint (need header line) + vspace_file = tmp_path / "vspace.in" + vspace_file.write_text("header line\ndestfolder NoCheckpoint\n") + + # Change to tmp_path directory + original_cwd = os.getcwd() + os.chdir(tmp_path) + try: + with pytest.raises(Exception, match="Multi-Planet must be running"): + mpstatus.mpstatus(str(vspace_file)) + finally: + os.chdir(original_cwd) + + def test_mpstatus_folder_parsing(self, tmp_path, capsys): + """Parse destination folder from vspace file.""" + # Create vspace.in with custom folder name (need header line) + vspace_file = tmp_path / "vspace.in" + vspace_file.write_text("header line\ndestfolder CustomFolderName\n") + + # Create checkpoint for custom folder + checkpoint_file = tmp_path / ".CustomFolderName" + checkpoint_file.write_text( + "Vspace File: /test/vspace.in\n" + "Total Number of Simulations: 1\n" + "sim_001 1\n" + "THE END\n" + ) + + # Change to tmp_path directory + original_cwd = os.getcwd() + os.chdir(tmp_path) + try: + mpstatus.mpstatus(str(vspace_file)) + captured = capsys.readouterr() + + # Verify it successfully parsed and read the checkpoint + assert "Number of Simulations completed: 1" in captured.out + finally: + os.chdir(original_cwd) + + +class TestMpstatusArguments: + """Tests for mpstatus Arguments function.""" + + def test_mpstatus_arguments_parsing(self, tmp_path): + """Parse input file argument correctly.""" + # Create vspace.in and checkpoint + vspace_file = tmp_path / "test_vspace.in" + vspace_file.write_text("destfolder TestArgs\n") + + checkpoint_file = tmp_path / ".TestArgs" + checkpoint_file.write_text( + "Vspace File: /test/vspace.in\n" + "Total Number of Simulations: 1\n" + "sim_001 -1\n" + "THE END\n" + ) + + # Mock sys.argv + original_cwd = os.getcwd() + os.chdir(tmp_path) + try: + with mock.patch('sys.argv', ['mpstatus', str(vspace_file)]): + # Mock mpstatus function to verify it's called + with mock.patch('multiplanet.mpstatus.mpstatus') as mock_mpstatus: + mpstatus.Arguments() + mock_mpstatus.assert_called_once_with(str(vspace_file)) + finally: + os.chdir(original_cwd) diff --git a/tests/unit/test_multiplanet.py b/tests/unit/test_multiplanet.py index b012864..e6583fe 100644 --- a/tests/unit/test_multiplanet.py +++ b/tests/unit/test_multiplanet.py @@ -9,6 +9,7 @@ import tempfile import multiprocessing as mp import pytest +from unittest import mock from multiplanet import multiplanet @@ -120,3 +121,342 @@ def test_get_snames_exists(self): """ assert hasattr(multiplanet, 'GetSNames') assert callable(multiplanet.GetSNames) + + +class TestCreateCP: + """Tests for CreateCP function.""" + + def test_create_cp_basic(self): + """Create checkpoint file with correct format.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: + checkpoint_file = f.name + + try: + input_file = "test_vspace.in" + sims = ["sim_001", "sim_002", "sim_003"] + + # Call CreateCP + multiplanet.CreateCP(checkpoint_file, input_file, sims) + + # Read checkpoint file + with open(checkpoint_file, 'r') as f: + lines = f.readlines() + + # Verify header + assert "Vspace File:" in lines[0] + assert input_file in lines[0] + assert "Total Number of Simulations: 3" in lines[1] + + # Verify each sim has status -1 + assert "sim_001 -1" in lines[2] + assert "sim_002 -1" in lines[3] + assert "sim_003 -1" in lines[4] + + # Verify ends with "THE END" + assert "THE END" in lines[5] + + finally: + os.unlink(checkpoint_file) + + def test_create_cp_header_format(self): + """Verify header contains vspace file path and count.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: + checkpoint_file = f.name + + try: + input_file = "my_vspace.in" + sims = ["sim_a", "sim_b"] + + multiplanet.CreateCP(checkpoint_file, input_file, sims) + + with open(checkpoint_file, 'r') as f: + lines = f.readlines() + + # Verify line 1 format + assert lines[0].startswith("Vspace File:") + assert os.getcwd() in lines[0] + assert input_file in lines[0] + + # Verify line 2 format + assert lines[1].strip() == "Total Number of Simulations: 2" + + finally: + os.unlink(checkpoint_file) + + +class TestReCreateCP: + """Tests for ReCreateCP function.""" + + def test_recreate_cp_reset_in_progress(self): + """Reset in-progress simulations (0 → -1).""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: + f.write("Vspace File: /test/vspace.in\n") + f.write("Total Number of Simulations: 3\n") + f.write("sim_001 0\n") # in-progress + f.write("sim_002 -1\n") # pending + f.write("sim_003 0\n") # in-progress + f.write("THE END\n") + checkpoint_file = f.name + + try: + sims = ["sim_001", "sim_002", "sim_003"] + multiplanet.ReCreateCP(checkpoint_file, "vspace.in", False, sims, "folder", False) + + with open(checkpoint_file, 'r') as f: + lines = f.readlines() + + # Verify all status=0 changed to -1 + assert "sim_001 -1" in lines[2] + assert "sim_002 -1" in lines[3] + assert "sim_003 -1" in lines[4] + + finally: + os.unlink(checkpoint_file) + + def test_recreate_cp_append_missing_sims(self): + """Append simulations added after initial checkpoint.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: + f.write("Vspace File: /test/vspace.in\n") + f.write("Total Number of Simulations: 2\n") + f.write("sim_001 -1\n") + f.write("sim_002 -1\n") + # Missing "THE END" - this triggers append behavior + checkpoint_file = f.name + + try: + # Provide 4 sims, but checkpoint only has 2 + sims = ["sim_001", "sim_002", "sim_003", "sim_004"] + multiplanet.ReCreateCP(checkpoint_file, "vspace.in", False, sims, "folder", False) + + with open(checkpoint_file, 'r') as f: + content = f.read() + + # Verify new simulations were appended + assert "sim_003 -1" in content + assert "sim_004 -1" in content + assert "THE END" in content + + finally: + os.unlink(checkpoint_file) + + def test_recreate_cp_missing_end_marker(self): + """Handle 'THE END' marker missing.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: + f.write("Vspace File: /test/vspace.in\n") + f.write("Total Number of Simulations: 2\n") + f.write("sim_001 -1\n") + f.write("sim_002 -1\n") + # No "THE END" marker + checkpoint_file = f.name + + try: + sims = ["sim_001", "sim_002", "sim_003"] + multiplanet.ReCreateCP(checkpoint_file, "vspace.in", False, sims, "folder", False) + + with open(checkpoint_file, 'r') as f: + lines = f.readlines() + + # Verify "THE END" was appended + assert any("THE END" in line for line in lines) + # Verify sim_003 was added + assert any("sim_003" in line for line in lines) + + finally: + os.unlink(checkpoint_file) + + def test_recreate_cp_all_complete(self, capsys): + """Detect when all simulations finished.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: + f.write("Vspace File: /test/vspace.in\n") + f.write("Total Number of Simulations: 2\n") + f.write("sim_001 1\n") + f.write("sim_002 1\n") + f.write("THE END\n") + checkpoint_file = f.name + + try: + sims = ["sim_001", "sim_002"] + + # Should exit when all complete and force=False + with pytest.raises(SystemExit): + multiplanet.ReCreateCP(checkpoint_file, "vspace.in", False, sims, "folder", False) + + captured = capsys.readouterr() + assert "All simulations have been ran" in captured.out + + finally: + if os.path.exists(checkpoint_file): + os.unlink(checkpoint_file) + + def test_recreate_cp_force_flag_deletes(self): + """Force recreates checkpoint when all complete.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f: + f.write("Vspace File: /test/vspace.in\n") + f.write("Total Number of Simulations: 2\n") + f.write("sim_001 1\n") + f.write("sim_002 1\n") + f.write("THE END\n") + checkpoint_file = f.name + + # Create a dummy folder file to be removed + folder_file = tempfile.mktemp() + with open(folder_file, 'w') as f: + f.write("dummy") + + try: + sims = ["sim_001", "sim_002"] + + # Call with force=True + multiplanet.ReCreateCP(checkpoint_file, "vspace.in", False, sims, folder_file, True) + + # Verify checkpoint was recreated (all sims reset to -1) + with open(checkpoint_file, 'r') as f: + content = f.read() + + assert "sim_001 -1" in content + assert "sim_002 -1" in content + + finally: + if os.path.exists(checkpoint_file): + os.unlink(checkpoint_file) + if os.path.exists(folder_file): + os.unlink(folder_file) + + +class TestArguments: + """Tests for Arguments CLI function.""" + + def test_arguments_default_cores(self): + """Use all CPU cores by default.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f: + f.write("destfolder TestSims\n") + vspace_file = f.name + + try: + # Mock sys.argv and subprocess calls + with mock.patch('sys.argv', ['multiplanet', vspace_file]): + with mock.patch('subprocess.getoutput', return_value="vplanet help"): + with mock.patch('multiplanet.multiplanet.parallel_run_planet') as mock_run: + multiplanet.Arguments() + + # Verify parallel_run_planet was called + assert mock_run.called + # Verify cores argument equals cpu_count + call_args = mock_run.call_args[0] + assert call_args[1] == mp.cpu_count() + + finally: + os.unlink(vspace_file) + + def test_arguments_custom_cores(self): + """Accept custom core count.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f: + vspace_file = f.name + + try: + with mock.patch('sys.argv', ['multiplanet', vspace_file, '-c', '4']): + with mock.patch('subprocess.getoutput', return_value="vplanet help"): + with mock.patch('multiplanet.multiplanet.parallel_run_planet') as mock_run: + multiplanet.Arguments() + + call_args = mock_run.call_args[0] + assert call_args[1] == 4 + + finally: + os.unlink(vspace_file) + + def test_arguments_bigplanet_flag(self): + """Enable bigplanet mode with -bp flag.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f: + vspace_file = f.name + + try: + with mock.patch('sys.argv', ['multiplanet', vspace_file, '-bp']): + with mock.patch('subprocess.getoutput', return_value="vplanet help"): + with mock.patch('multiplanet.multiplanet.parallel_run_planet') as mock_run: + multiplanet.Arguments() + + # bigplanet is the 5th argument (index 4) + call_args = mock_run.call_args[0] + assert call_args[4] == True + + finally: + os.unlink(vspace_file) + + def test_arguments_force_flag(self): + """Enable force mode with -f flag.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f: + vspace_file = f.name + + try: + with mock.patch('sys.argv', ['multiplanet', vspace_file, '-f']): + with mock.patch('subprocess.getoutput', return_value="vplanet help"): + with mock.patch('multiplanet.multiplanet.parallel_run_planet') as mock_run: + multiplanet.Arguments() + + # force is the 6th argument (index 5) + call_args = mock_run.call_args[0] + assert call_args[5] == True + + finally: + os.unlink(vspace_file) + + def test_arguments_quiet_mode(self): + """Enable quiet mode with -q flag.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f: + vspace_file = f.name + + try: + with mock.patch('sys.argv', ['multiplanet', vspace_file, '-q']): + with mock.patch('subprocess.getoutput', return_value="vplanet help"): + with mock.patch('multiplanet.multiplanet.parallel_run_planet') as mock_run: + multiplanet.Arguments() + + call_args = mock_run.call_args[0] + # quiet is 3rd arg, verbose is 4th arg + assert call_args[2] == True # quiet=True + assert call_args[3] == False # verbose=False + + finally: + os.unlink(vspace_file) + + def test_arguments_verbose_mode(self): + """Enable verbose mode with -v flag.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f: + vspace_file = f.name + + try: + with mock.patch('sys.argv', ['multiplanet', vspace_file, '-v']): + with mock.patch('subprocess.getoutput', return_value="vplanet help"): + with mock.patch('multiplanet.multiplanet.parallel_run_planet') as mock_run: + multiplanet.Arguments() + + call_args = mock_run.call_args[0] + # quiet is 3rd arg, verbose is 4th arg + assert call_args[2] == False # quiet=False + assert call_args[3] == True # verbose=True + + finally: + os.unlink(vspace_file) + + def test_arguments_vplanet_check(self): + """Verify vplanet executable check.""" + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.in') as f: + vspace_file = f.name + + try: + # Test successful vplanet check + with mock.patch('sys.argv', ['multiplanet', vspace_file]): + with mock.patch('subprocess.getoutput', return_value="vplanet help"): + with mock.patch('multiplanet.multiplanet.parallel_run_planet'): + # Should not raise exception + multiplanet.Arguments() + + # Test failed vplanet check + with mock.patch('sys.argv', ['multiplanet', vspace_file]): + with mock.patch('subprocess.getoutput', side_effect=OSError("command not found")): + with pytest.raises(Exception, match="Unable to call VPLANET"): + multiplanet.Arguments() + + finally: + os.unlink(vspace_file)