#!/usr/bin/env python
##############################################################################
# IMPORT ALL NECESSARY PYTHON PACKAGES
##############################################################################
import argparse
import itertools
import os
import random
import shutil
import stat
import string
import sys
import time
from collections import ChainMap
import numpy as np
import pandas
from posydon.active_learning.psy_cris.utils import parse_inifile
from posydon.grids.psygrid import PSyGrid
from posydon.utils import gridutils as utils
from posydon.utils.posydonwarning import Pwarn
try:
from mpi4py import MPI
MPI_IMPORTED = True
except ImportError:
MPI_IMPORTED = False
###############################################################################
# DEFINE COMMANDLINE ARGUMENTS
###############################################################################
[docs]
def parse_commandline():
"""Parse the arguments given on the command-line.
"""
parser = argparse.ArgumentParser(description=__doc__)
grid_arguments = parser.add_argument_group('Arguments related to the running of the grid, including '
'specifying the grid values and type of grid to be run, '
'where the MESA simulation will be run, '
'and where the output of the simulation should go')
grid_arguments.add_argument("--mesa-grid",
help="Grid of mesa points with values already known,"
"or new grid values to run MESA on",
required=True)
grid_arguments.add_argument("--grid-type",
help="Either you are supplying a grid "
"of points to run MESA on (fixed) or you are supplying a pre computed MESA "
"grid and want to sampling new points to run MESA on (dynamic).",
required=True)
grid_arguments.add_argument("--output-directory",
help="This is where the result of the MESA simulatio will go",
required=True)
grid_arguments.add_argument("--temporary-directory",
help="This is where your MESA simulations will run", default=os.getcwd())
fixed_grid_arguments = parser.add_argument_group('Arguments specific to running a fixed grid')
fixed_grid_arguments.add_argument("--grid-point-index",
help="Override grid point slection by "
"explictly saying which grid point you want."
"This must be run with mpirun -np 1!!",
type=int)
dynamic_grid_arguments = parser.add_argument_group('Arguments specific to running a dynamic grid')
dynamic_grid_arguments.add_argument("--psycris-inifile",
help="Path to the psycris inifile. ")
dynamic_grid_arguments.add_argument("--Niter",
help="Number of iterations of binaries "
"to try, will check ever Nstep for convergence",
type=int, default=200)
mesa_executables = parser.add_argument_group('absolute path to MESA executables to be used in grid')
mesa_executables.add_argument("--mesa-binary-executable",
help="MESA binary executable to be used for this grid",
required=True)
mesa_executables.add_argument("--mesa-star1-executable",
help="optionally use star executable to form star1",
)
mesa_executables.add_argument("--mesa-star2-executable",
help="optionally use star executable to form star2",
)
inlists_for_running_binary = parser.add_argument_group('absolute paths to static inlists related to the call of the binary executable')
inlists_for_running_binary.add_argument("--mesa-binary-inlist-project",
help="This is the shared inlist_project that will be used by "
"every MESA run for this grid. inlist_project has both binary_job and binary_controls in it",
required=True)
inlists_for_running_binary.add_argument("--mesa-binary-inlist1",
help="This is the shared inlist1 that will be used by "
"every MESA run for this grid. inlist1 has both star_job and controls in it that determine "
"the properties of the first object",
required=True)
inlists_for_running_binary.add_argument("--mesa-binary-inlist2",
help="This is the shared inlist2 that will be used by "
"every MESA run for this grid. inlist2 has both star_job and controls in it that determine "
"the properties of the second object",
required=True)
inlists_for_running_star1 = parser.add_argument_group('absolute paths to static inlists related to the formation of star1')
inlists_for_running_star1.add_argument("--mesa-star1-inlist-project",
help="optionally use specific inlist(s) to form star1",
nargs='+')
inlists_for_running_star2 = parser.add_argument_group('absolute paths to static inlists related to the formation of star2')
inlists_for_running_star2.add_argument("--mesa-star2-inlist-project",
help="optionally use specific inlist(s) to form star2",
nargs='+')
mesa_columns = parser.add_argument_group('absolute paths to list of columns desired from each MESA simulation on this grid')
mesa_columns.add_argument("--mesa-star-history-columns",
help="This is the path to the star_history_columns.list file you want to use for this grid",
required=True)
mesa_columns.add_argument("--mesa-binary-history-columns",
help="This is the path to the binary_history_columns.list file you want to use for this grid",
required=True)
mesa_columns.add_argument("--mesa-profile-columns",
help="This is the path to the profile_columns.list file you want to use for this grid",
required=True)
parser.add_argument("--verbose", action="store_true", default=False,
help="Run in Verbose Mode")
parser.add_argument("--keep_profiles", action="store_true", default=False,
help="Do not delete profiles at run completion")
parser.add_argument("--keep_photos", action="store_true", default=False,
help="Do not delete photos at run completion")
now = int(time.time())
parser.add_argument("--job_end", type=int, default=now+60*60*24*365,
help="End time of the job (in seconds)")
parser.add_argument("--job_before_end", type=int, default=300,
help="Time span till end of job (in seconds)")
args = parser.parse_args()
if args.grid_type not in ['fixed', 'dynamic']:
raise parser.error("--run-type must be either fixed or dynamic")
# if we are running a dynamic grid then these arguments are required otherwise
# they are not
if args.grid_type == 'dynamic':
if (args.psycris_inifile is None):
raise parser.error("You must supply a psycris inifile when running a dynamic grid")
return args
[docs]
def id_generator(size=10,
chars=(string.ascii_uppercase +
string.digits +
string.ascii_lowercase)):
"""Obtain omicron triggers run gravityspy on
Parameters:
x (str): the item you would like a random id to be generated for
Returns:
"""
return ''.join(random.SystemRandom().choice(chars) for _ in range(size))
[docs]
def create_working_directory(grid_param_dict, args):
"""Function to build the directory name and make inlists for a given grid point
Params:
grid_param_dict:
Dictionary with keys being the parameter names of the grid
and values the actual values for those parameters for this grid point
args:
The arguments passed to posydon-run-grid
"""
############# set directories
variable_names_to_log = ['initial_z', 'initial_period_in_days']
# first name directories which we do not need to log
directory_name_from_dict_no_log = '_'.join(['{0}_{1:.4f}'.format(k, v) for k, v in grid_param_dict.items() if k not in variable_names_to_log])
directory_name_from_dict_log = '_'.join(['{0}_{1:.4e}'.format(k, v) for k, v in grid_param_dict.items() if k in variable_names_to_log])
# drop parenthesis
directory_name_from_dict_no_log = directory_name_from_dict_no_log.replace('(','')
directory_name_from_dict_no_log = directory_name_from_dict_no_log.replace(')','')
# combine and add on what grid number we are running if exists
directory_name_from_dict = '{0}_{1}'.format(directory_name_from_dict_no_log, directory_name_from_dict_log)
if args.grid_point_index is not None:
directory_name_from_dict = '{0}_{1}_{2}'.format(directory_name_from_dict, 'grid_index', args.grid_point_index)
if args.grid_type == 'dynamic':
directory_name_from_dict = '{0}_{1}'.format(directory_name_from_dict, id_generator())
temp_dir = args.temporary_directory #all sims will be created inside this directory (e.g. /scratch/pablo_nu/sim1/)
work_dir = os.path.join(temp_dir, directory_name_from_dict)
final_dir = os.path.join(args.output_directory, directory_name_from_dict)
############# check for existance of precomputed model
if os.path.isdir(work_dir):
Pwarn("Directory {} exists, removing and rerunning".format(work_dir),
"OverwriteWarning")
shutil.rmtree(work_dir)
############# create work directory and change into it
os.makedirs(work_dir)
shutil.copy(args.mesa_star_history_columns, work_dir)
shutil.copy(args.mesa_binary_history_columns, work_dir)
shutil.copy(args.mesa_profile_columns, work_dir)
return work_dir, final_dir
[docs]
def create_binary_inlists(binary_controls, binary_job,
star1_binary_controls, star1_binary_job,
star2_binary_controls, star2_binary_job,
work_dir, binary_inlist_project):
"""Create the inlist and inlist grid points needed to run this part of grid
"""
binary_controls_str = """
&binary_controls
read_extra_binary_controls_inlist1 = .true.
extra_binary_controls_inlist1_name = '{0}'
read_extra_binary_controls_inlist2 = .true.
extra_binary_controls_inlist2_name = 'inlist_grid_points'
/ ! end of controls namelist
""".format(binary_inlist_project)
binary_job_str = """
&binary_job
read_extra_binary_job_inlist1 = .true.
extra_binary_job_inlist1_name = '{0}'
/ ! end of job namelist
""".format(binary_inlist_project)
filename = os.path.join(work_dir, 'inlist')
if os.path.exists(filename):
Pwarn('Replace '+filename, "OverwriteWarning")
with open(filename, 'w') as f:
f.write(binary_controls_str)
f.write('\n')
f.write(binary_job_str)
filename = os.path.join(work_dir, 'inlist_grid_points')
if os.path.exists(filename):
Pwarn('Replace '+filename, "OverwriteWarning")
with open(filename, 'w') as inlist_grid_points:
if binary_controls:
# Write the binary controls for this grid point
inlist_grid_points.write("&binary_controls\n")
for k, v in binary_controls.items():
if (type(v) == float) and (not v.is_integer()):
inlist_grid_points.write("\t{0} = ".format(k)+"{0:.10e}\n"\
.format(v).replace('e','d'))
elif (type(v) == float) and (v.is_integer()):
inlist_grid_points.write("\t{0} = {1}\n".format(k, int(v)))
elif type(v) == bool:
if v:
inlist_grid_points.write("\t{0} = .true.\n".format(k))
else:
inlist_grid_points.write("\t{0} = .false.\n".format(k))
else:
inlist_grid_points.write("{0} = {1}\n".format(k, v))
inlist_grid_points.write("/ ! end of binary_controls namelist\n")
if binary_job:
# Write the binary job for this grid point
inlist_grid_points.write("&binary_job\n")
for k, v in binary_job.items():
if (type(v) == float) and (not v.is_integer()):
inlist_grid_points.write("\t{0} = ".format(k)+"{0:.10e}\n"\
.format(v).replace('e','d'))
elif (type(v) == float) and (v.is_integer()):
inlist_grid_points.write("\t{0} = {1}\n".format(k, int(v)))
elif type(v) == bool:
if v:
inlist_grid_points.write("\t{0} = .true.\n".format(k))
else:
inlist_grid_points.write("\t{0} = .false.\n".format(k))
else:
inlist_grid_points.write("{0} = {1}\n".format(k, v))
inlist_grid_points.write("/ ! end of binary_job namelist\n")
if star1_binary_controls:
filename = os.path.join(work_dir, 'inlist_grid_star1_binary_controls')
if os.path.exists(filename):
Pwarn('Replace '+filename, "OverwriteWarning")
with open(filename, 'w') as inlist_grid_points:
# Write the star1 controls for this grid point
inlist_grid_points.write("&controls\n")
for k, v in star1_binary_controls.items():
if (type(v) == float) and (not v.is_integer()):
inlist_grid_points.write("\t{0} = ".format(k)+"{0:.10e}\n"\
.format(v).replace('e','d'))
elif (type(v) == float) and (v.is_integer()):
inlist_grid_points.write("\t{0} = {1}\n".format(k, int(v)))
elif type(v) == bool:
if v:
inlist_grid_points.write("\t{0} = .true.\n".format(k))
else:
inlist_grid_points.write("\t{0} = .false.\n".format(k))
else:
inlist_grid_points.write("{0} = {1}\n".format(k, v))
inlist_grid_points.write("/ ! end of controls namelist\n")
if star2_binary_controls:
filename = os.path.join(work_dir, 'inlist_grid_star2_binary_controls')
if os.path.exists(filename):
Pwarn('Replace '+filename, "OverwriteWarning")
with open(filename, 'w') as inlist_grid_points:
# Write the star2 controls for this grid point
inlist_grid_points.write("&controls\n")
for k, v in star2_binary_controls.items():
if (type(v) == float) and (not v.is_integer()):
inlist_grid_points.write("\t{0} = ".format(k)+"{0:.10e}\n"\
.format(v).replace('e','d'))
elif (type(v) == float) and (v.is_integer()):
inlist_grid_points.write("\t{0} = {1}\n".format(k, int(v)))
elif type(v) == bool:
if v:
inlist_grid_points.write("\t{0} = .true.\n".format(k))
else:
inlist_grid_points.write("\t{0} = .false.\n".format(k))
else:
inlist_grid_points.write("{0} = {1}\n".format(k, v))
inlist_grid_points.write("/ ! end of controls namelist\n")
if star1_binary_job:
filename = os.path.join(work_dir, 'inlist_grid_star1_binary_job')
if os.path.exists(filename):
Pwarn('Replace '+filename, "OverwriteWarning")
with open(filename, 'w') as inlist_grid_points:
# Write the star1 job for this grid point
inlist_grid_points.write("&star_job\n")
for k, v in star1_binary_job.items():
if (type(v) == float) and (not v.is_integer()):
inlist_grid_points.write("\t{0} = ".format(k)+"{0:.10e}\n"\
.format(v).replace('e','d'))
elif (type(v) == float) and (v.is_integer()):
inlist_grid_points.write("\t{0} = {1}\n".format(k, int(v)))
elif type(v) == bool:
if v:
inlist_grid_points.write("\t{0} = .true.\n".format(k))
else:
inlist_grid_points.write("\t{0} = .false.\n".format(k))
else:
inlist_grid_points.write("{0} = {1}\n".format(k, v))
inlist_grid_points.write("/ ! end of star_job namelist\n")
if star2_binary_job:
filename = os.path.join(work_dir, 'inlist_grid_star2_binary_job')
if os.path.exists(filename):
Pwarn('Replace '+filename, "OverwriteWarning")
with open(filename, 'w') as inlist_grid_points:
# Write the star2 job for this grid point
inlist_grid_points.write("&star_job\n")
for k, v in star2_binary_job.items():
if (type(v) == float) and (not v.is_integer()):
inlist_grid_points.write("\t{0} = ".format(k)+"{0:.10e}\n"\
.format(v).replace('e','d'))
elif (type(v) == float) and (v.is_integer()):
inlist_grid_points.write("\t{0} = {1}\n".format(k, int(v)))
elif type(v) == bool:
if v:
inlist_grid_points.write("\t{0} = .true.\n".format(k))
else:
inlist_grid_points.write("\t{0} = .false.\n".format(k))
else:
inlist_grid_points.write("{0} = {1}\n".format(k, v))
inlist_grid_points.write("/ ! end of star_job namelist\n")
return
[docs]
def move_mesa_output(work_dir, final_dir, copyinstead=False):
"""Moves the data from the working directory to the final directory.
Parameters:
work_dir:
Working directory (i.e. where the mesa exectuable should run)
final_dir:
When completed where would you like the output to go
copyinstead:
Indicates that the files should be copied instead of moved
"""
if not work_dir == final_dir:
if os.path.isdir(final_dir):
Pwarn("Directory {} exists,".format(work_dir) +\
" replacing with new data", "OverwriteWarning")
shutil.rmtree(final_dir)
if copyinstead:
print ("copy data from " + work_dir+ " to "+ final_dir)
shutil.copytree(work_dir, final_dir)
else:
print ("move data from " + work_dir+ " to "+ final_dir)
shutil.move(work_dir, final_dir)
sys.stdout.flush()
return
[docs]
def run_mesa(mesa_executable, work_dir, final_dir, **kwargs):
"""Provides minor utility for actually spawning the mesa job for each grid
mesa_executable:
Path to the executable you are running
work_dir:
Working directory (i.e. where the mesa exectuable should run)
final_dir:
When completed where would you like the output to go
"""
outfile_name = kwargs.pop('outfile_name', 'out.txt')
keep_profiles = kwargs.pop('keep_profiles', False)
keep_photos = kwargs.pop('keep_photos', False)
keep_work_dir = kwargs.pop('keep_work_dir', False)
os.chdir(work_dir)
if not work_dir == final_dir:
child_ID = os.fork()
else:
child_ID = 0.1
if child_ID>0:
#parent process: run MESA and do the clean up in the working directory
os.system("{0} &> {1}".format(mesa_executable, os.path.join(work_dir, outfile_name)))
if not keep_profiles:
if os.path.exists(os.path.join(work_dir, 'LOGS')):
os.chmod(os.path.join(work_dir, 'LOGS'), stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH )
os.system("rm LOGS/profile*")
if os.path.exists(os.path.join(work_dir, 'LOGS1')):
os.chmod(os.path.join(work_dir, 'LOGS1'), stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH )
os.system("rm LOGS1/profile*")
if os.path.exists(os.path.join(work_dir, 'LOGS2')):
os.chmod(os.path.join(work_dir, 'LOGS2'), stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH )
os.system("rm LOGS2/profile*")
if os.path.exists(os.path.join(work_dir, '.mesa_temp_cache')): os.system("rm -rf .mesa_temp_cache")
if not keep_photos:
os.system("rm -rf photos*")
if child_ID>=1:
#kill child process, because it is no longer needed
os.kill(child_ID, 9)
move_mesa_output(work_dir, final_dir, keep_work_dir)
elif child_ID==0:
now = int(time.time())
child_end = kwargs.pop('job_end', now+60*60*24*365) - kwargs.pop('job_before_end', 300)
#child process: waits till short before the job will get killed by slurm to copy the data beforehand and exit this process after the copy finished
if child_end>now:
#if the child job still needs to wait, get time to wait
wait_time = child_end - now
else:
#otherwise wait 1 minute
wait_time = 60
time.sleep(wait_time)
move_mesa_output(work_dir, final_dir, copyinstead=True)
sys.exit(0)
os.chdir(final_dir)
print("Done")
return
[docs]
def get_next_grid_point(psycris_inifile_path, grid, n_new_points, **kwargs):
"""Sample a new grid point
Parameters
----------
psycris_inifile_path : str
Path to psycris inifile
grid : pandas DataFrame
DataFrame of initial and final values of training data
n_new_points : int
Number of new points requested from the algorithm
Returns
-------
query_points : array
Initial conditions of new MESA run to label.
pred_classess : array
Predicted class of query point.
"""
from posydon.active_learning.psy_cris.utils import (
get_new_query_points,
parse_inifile,
)
from posydon.grids.psygrid import PSyGrid
from posydon.interpolation.data_scaling import DataScaler
normed_grid = grid.copy()
mesa_kwargs = parse_inifile(psycris_inifile_path)
input_cols = mesa_kwargs["TableData_kwargs"]["input_cols"]
output_cols = mesa_kwargs["TableData_kwargs"]["output_cols"]
in_scaling = mesa_kwargs["posydon_dynamic_sampling_kwargs"]["in_scaling"]
out_scaling = mesa_kwargs["posydon_dynamic_sampling_kwargs"]["out_scaling"]
# Check if sampling input and output columns match data from psygrid
if not all([name in grid.columns for name in input_cols]):
raise ValueError("Bad input columns given: {0},\nMust be in: {1}".format(input_cols, grid.columns) )
if not all([name in grid.columns for name in output_cols]):
raise ValueError("Bad output columns given: {0}\nMust be in: {1}".format(output_cols, grid.columns) )
# default scaling is from [-1,1]
# scale input data
input_scalers = []
for i, input_col_name in enumerate(input_cols):
ds = DataScaler()
col_data = ds.fit_and_transform( grid[input_col_name].values, method=in_scaling[i] )
normed_grid[input_col_name] = col_data
input_scalers.append(ds)
# scale output data
output_scalers = []
for i, output_col_name in enumerate(output_cols):
if output_col_name==mesa_kwargs["TableData_kwargs"]["class_col_name"]:
continue
ds = DataScaler()
col_data = ds.fit_and_transform( grid[output_col_name].values, method=out_scaling[i] )
normed_grid[output_col_name] = col_data
output_scalers.append(ds)
# protect overwriting nested dicts
holder = mesa_kwargs['TableData_kwargs'].copy()
holder["my_DataFrame"] = normed_grid
mesa_kwargs['TableData_kwargs'] = holder
normed_query_points, pred_classess = get_new_query_points(N_new_points=n_new_points, **mesa_kwargs)
# inv transform the query points from sapling space for MESA variables
query_points = np.empty(normed_query_points.shape, dtype=float)
for n in range( len(input_scalers) ):
query_points[:,n] = input_scalers[n].inv_transform(normed_query_points[:,n])
return query_points, pred_classess
[docs]
def do_root_process_logic(comm, grid, args, grid_params_binary_controls, grid_params_binary_job,
grid_params_star1_binary_controls, grid_params_star1_binary_job,
grid_params_star2_binary_controls, grid_params_star2_binary_job):
"""Execute the logic that the root process does for the dynamic grid
"""
from posydon.active_learning.psy_cris.utils import parse_inifile
nstep = 0
print('Beginning dynamic grid...')
# As first step find the new grid points for all the child process to run on and send that data to them
proposed_points, predicted_classes_of_points = get_next_grid_point(args.psycris_inifile, grid=grid, n_new_points=size)
for child in range(1,size):
print('Sampling first grid point for process {0}'.format(child))
sys.stdout.flush()
# make a dataframe from the proposed points with the names of the columns
dynamic_grid_params = parse_inifile(args.psycris_inifile)
grid_df = pandas.DataFrame(np.atleast_2d(proposed_points[child-1,:]),columns=dynamic_grid_params["TableData_kwargs"]["input_cols"],)
grid_df = convert_input_cols_to_mesa_cols(grid_df)
# make a dict so we can send the information around
grid_dict = grid_df.to_dict('records')[0]
print('Sending process number {0} a new grid point...'.format(child))
sys.stdout.flush()
req = comm.isend(grid_dict, dest=child,)
req.wait()
print('Process number {0} received their new grid point...'.format(child))
sys.stdout.flush()
while nstep < args.Niter:
# Now we wait for one of the child processes to finish their grid point and send the result back to the root process so we can
# sample a new point using that extra info
state = MPI.Status()
mesa_result = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=state)
child = state.Get_source()
# calculate some columns for convenience
if ('initial_star_2_mass' in mesa_result.columns) and ('initial_star_1_mass' in mesa_result.columns):
mesa_result['q'] = mesa_result['initial_star_2_mass']/mesa_result['initial_star_1_mass']
if 'initial_period_days' in mesa_result.columns:
mesa_result['log_p'] = np.log10(mesa_result['initial_period_days'])
try:
grid = grid.append(mesa_result[grid.columns], sort=False)
except:
raise ValueError("The version of POSYDON used to make the fixed grid psygrid object is likely different then the version psygrid object"
" being used for this dynamic grid run")
print('Process number {0} finished their run'.format(child))
print('Saving grid data and sampling a new grid point...')
sys.stdout.flush()
filename = os.path.join(args.output_directory, 'grid_results.csv')
if os.path.exists(filename):
Pwarn('Replace '+filename, "OverwriteWarning")
with open(filename, 'w') as f:
grid.to_csv(f, index=False,)
print('Sampling next grid point for process {0}'.format(child))
sys.stdout.flush()
proposed_points, predicted_classes_of_points = get_next_grid_point(args.psycris_inifile, grid=grid, n_new_points=1)
# make a dataframe from the proposed points with the names of the columns
dynamic_grid_params = parse_inifile(args.psycris_inifile)
grid_df = pandas.DataFrame(proposed_points, columns=dynamic_grid_params["TableData_kwargs"]["input_cols"],)
grid_df = convert_input_cols_to_mesa_cols(grid_df)
# make a dict so we can send the information around
grid_dict = grid_df.to_dict('records')[0]
print('Sending process number {0} a new grid point...'.format(child))
sys.stdout.flush()
req = comm.isend(grid_dict, dest=child,)
req.wait()
print('Process number {0} received their new grid point...'.format(child))
sys.stdout.flush()
nstep += 1
[docs]
def do_child_process_logic(comm, grid, args, star1_formation=False, star2_formation=False):
"""Execute the logic that the child process does for the dynamic grid
"""
req = comm.irecv(source=0,)
grid_dict = req.wait()
# per each type of grid value (i.e. binary control or job or star1 control or job etc)
# create a dictionary of the column name and value for this grid point
binary_controls = {k: grid_dict[k] for k in grid_params_binary_controls}
binary_job = {k: grid_dict[k] for k in grid_params_binary_job}
star1_binary_controls = {k: grid_dict[k] for k in grid_params_star1_binary_controls}
star1_binary_job = {k: grid_dict[k] for k in grid_params_star1_binary_job}
star2_binary_controls = {k: grid_dict[k] for k in grid_params_star2_binary_controls}
star2_binary_job = {k: grid_dict[k] for k in grid_params_star2_binary_job}
grid_param_dict = dict(ChainMap({}, binary_controls, binary_job, star1_binary_controls, star1_binary_job, star2_binary_controls,star2_binary_job))
grid_params_string = ', '.join(["{0} : {1:.4f}".format(k,v) for k, v in grid_param_dict.items()])
print("Process {0} is running with {1}".format(rank, grid_params_string))
sys.stdout.flush()
mesa_result = run_grid_point(grid, star1_formation, star2_formation,
grid_param_dict, binary_controls, binary_job,
star1_binary_controls, star1_binary_job,
star2_binary_controls, star2_binary_job, args)
# we now send the result data to the root node
comm.send(mesa_result, dest=0)
print("Job completed: From process {0}".format(rank))
sys.stdout.flush()
[docs]
def run_grid_point(grid, star1_formation, star2_formation,
grid_param_dict, binary_controls, binary_job,
star1_binary_controls, star1_binary_job,
star2_binary_controls, star2_binary_job, args):
"""Execute the logic to run this MESA grid point
"""
# based on the values create directory
work_dir, final_dir = create_working_directory(grid_param_dict, args)
# are we first generating a model before running binary for the first star?
if star1_formation:
# if yes then create inlists specific to this grid points
# moreover we should loop over all the star1 formation steps
for index, inlist_step in enumerate(args.mesa_star1_inlist_project):
# this is only called to create single HeMS stars, only second step1
# requires metallicity
if index == 0:
create_star_formation(grid_param_dict['m1'], work_dir, inlist_step)
elif index == 1:
create_star_formation(grid_param_dict['m1'], work_dir, inlist_step, grid_param_dict['initial_z'], True)
run_mesa(args.mesa_star1_executable, work_dir, final_dir, keep_work_dir=True,
outfile_name='out_star1_formation_step{0}.txt'.format(index))
# are we first generating a model before running binary for the second star?
if star2_formation:
# if yes then create inlists specific to this grid points
# moreover we should loop over all the star1 formation steps
for index, inlist_step in enumerate(args.mesa_star2_inlist_project):
# this is only called to create single HeMS stars, only second step1
# requires metallicity
if index == 0:
create_star_formation(grid_param_dict['m2'], work_dir, inlist_step)
elif index == 1:
create_star_formation(grid_param_dict['m2'], work_dir, inlist_step, grid_param_dict['initial_z'], True)
run_mesa(args.mesa_star2_executable, work_dir, final_dir, keep_work_dir=True,
outfile_name='out_star2_formation_step{0}.txt'.format(index))
# now we can create the grid specific binary inlists and run the binary executable
create_binary_inlists(binary_controls, binary_job,
star1_binary_controls, star1_binary_job,
star2_binary_controls, star2_binary_job,
work_dir, args.mesa_binary_inlist_project)
# run the binary exectuable
run_mesa(args.mesa_binary_executable, work_dir, final_dir,
keep_profiles=args.keep_profiles, keep_photos=args.keep_photos,
job_end=args.job_end, job_before_end=args.job_before_end)
# find out what happened
mesa_result = extract_mesa_results(final_dir)
return mesa_result
###############################################################################
# BEGIN MAIN FUNCTION
###############################################################################
if __name__ == '__main__':
# check MPI was imported
if not MPI_IMPORTED:
raise ImportError('MPI module mpi4py not installed! Plese check your POSYDON installation!')
# READ COMMANDLINE ARGUMENTS
###########################################################################
args = parse_commandline()
# do to the fact that the arguments args.mesa_star2_inlist_project and args.mesa_star1_inlist_project
# are interpreted as lists when you supply None it stringifies it and makes it a list
if 'None' in args.mesa_star1_inlist_project:
args.mesa_star1_inlist_project = None
if 'None' in args.mesa_star2_inlist_project:
args.mesa_star2_inlist_project = None
if (args.mesa_star1_executable is not None) and (args.mesa_star1_inlist_project is not None):
# Then we must first make star1 model before running binary executable
star1_formation = True
else:
star1_formation = False
if (args.mesa_star2_executable is not None) and (args.mesa_star2_inlist_project is not None):
# Then we must first make star2 model before running binary executable
star2_formation = True
else:
star2_formation = False
# read in grid
if '.csv' in args.mesa_grid:
grid = pandas.read_csv(args.mesa_grid)
elif '.h5' in args.mesa_grid:
psy_grid = PSyGrid()
psy_grid.load(args.mesa_grid)
grid = psy_grid.get_pandas_initial_final()
psy_grid.close()
else:
raise ValueError('Grid format not recognized, please feed in an acceptable format: csv')
if args.grid_type == 'dynamic':
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
name = MPI.Get_processor_name()
if (args.grid_point_index is not None) and (size > 1):
raise ValueError("You can not have multiple MPI tasks and "
"request to run a specific grid point."
"Please make sure you cal with mpirun -np 1.")
# identify all columns in the CSV file which effect MESA controls, i.e. what all the parameter
# values that change simulation to simulation are
binary_inlist_params = utils.clean_inlist_file(args.mesa_binary_inlist_project)
star1_binary_params = utils.clean_inlist_file(args.mesa_binary_inlist1)
star2_binary_params = utils.clean_inlist_file(args.mesa_binary_inlist2)
final_binary_controls = binary_inlist_params['&binary_controls']
final_binary_job = binary_inlist_params['&binary_job']
final_star1_binary_controls = star1_binary_params['&controls']
final_star1_binary_job = star1_binary_params['&star_job']
final_star2_binary_controls = star2_binary_params['&controls']
final_star2_binary_job = star2_binary_params['&star_job']
# detemine which is any of the parameters are binary_controls or binary_job params
if args.grid_type == 'dynamic':
dynamic_grid_params = parse_inifile(args.psycris_inifile)
mesa_params_to_run_grid_over = dynamic_grid_params["posydon_dynamic_sampling_kwargs"]["mesa_column_names"]
grid_params_binary_controls = [param for param in mesa_params_to_run_grid_over if param in final_binary_controls.keys()]
grid_params_binary_job = [param for param in mesa_params_to_run_grid_over if param in final_binary_job.keys()]
grid_params_star1_binary_controls = [param for param in mesa_params_to_run_grid_over if param in final_star1_binary_controls.keys()]
grid_params_star1_binary_job = [param for param in mesa_params_to_run_grid_over if param in final_star1_binary_job.keys()]
grid_params_star2_binary_controls = [param for param in mesa_params_to_run_grid_over if param in final_star2_binary_controls.keys()]
grid_params_star2_binary_job = [param for param in mesa_params_to_run_grid_over if param in final_star2_binary_job.keys()]
else:
grid_params_binary_controls = [param for param in grid.columns if param in final_binary_controls.keys()]
grid_params_binary_job = [param for param in grid.columns if param in final_binary_job.keys()]
grid_params_star1_binary_controls = [param for param in grid.columns if param in final_star1_binary_controls.keys()]
grid_params_star1_binary_job = [param for param in grid.columns if param in final_star1_binary_job.keys()]
grid_params_star2_binary_controls = [param for param in grid.columns if param in final_star2_binary_controls.keys()]
grid_params_star2_binary_job = [param for param in grid.columns if param in final_star2_binary_job.keys()]
# check if we are running a single star grid!
if star1_formation and not list(set(grid_params_binary_controls).union(grid_params_binary_job)):
# we need to check for something other than m1, we want to be looking for initial_mass
# identify all columns in the CSV file which effect MESA controls, i.e. what all the parameter
# values that change simulation to simulation are
star1_formation_params = utils.clean_inlist_file(args.mesa_star1_inlist_project[0])
final_star1_formation_controls = star1_formation_params['&controls']
final_star1_formation_job = star1_formation_params['&star_job']
# detemine which is any of the parameters are formation_controls or formation_job params
grid_params_star1_formation_controls = [param for param in grid.columns if param in final_star1_formation_controls.keys()]
grid_params_star1_formation_job = [param for param in grid.columns if param in final_star1_formation_job.keys()]
star1_formation_controls = grid.iloc[args.grid_point_index][grid_params_star1_formation_controls].to_dict()
star1_formation_job = grid.iloc[args.grid_point_index][grid_params_star1_formation_job].to_dict()
grid_param_dict = dict(ChainMap({}, star1_formation_controls, star1_formation_job,))
print("You are running a single star grid, we are checking for what parameters you are varying in your grid")
print("Grid parameters that effect star1_formation_controls: {0}".format(','.join(grid_params_star1_formation_controls)))
print("Grid parameters that effect star1_formation_job: {0}".format(','.join(grid_params_star1_formation_job)))
# based on the values create directory
work_dir, final_dir = create_working_directory(grid_param_dict, args)
# run the single star grid
last_step = len(args.mesa_star1_inlist_project) - 1
for index, inlist_step in enumerate(args.mesa_star1_inlist_project):
# single_HMS star has only step0, metallicity is required but not new_Z
# single_HeMS star has step0, step1, step2, last two steps require metallicity
# but only step1 requires new_Z
if index == last_step:
create_star_formation(grid_param_dict['initial_mass'], work_dir, inlist_step, grid_param_dict['initial_z'], False)
elif index == 1:
create_star_formation(grid_param_dict['initial_mass'], work_dir, inlist_step, grid_param_dict['initial_z'], True)
else :
create_star_formation(grid_param_dict['initial_mass'], work_dir, inlist_step)
run_mesa(args.mesa_star1_executable, work_dir, final_dir,
outfile_name='out_star1_formation_step{0}.txt'.format(index),
keep_profiles=args.keep_profiles, keep_photos=args.keep_photos,
job_end=args.job_end, job_before_end=args.job_before_end,
keep_work_dir=index<last_step)
sys.exit(0)
# only print this information is this is the root process
if args.grid_type == 'dynamic':
if rank == 0:
print("Grid parameters that effect binary_controls: {0}".format(','.join(grid_params_binary_controls)))
print("Grid parameters that effect binary_job: {0}".format(','.join(grid_params_binary_job)))
print("Grid parameters that effect star1_binary_controls: {0}".format(','.join(grid_params_star1_binary_controls)))
print("Grid parameters that effect star1_binary_job: {0}".format(','.join(grid_params_star1_binary_job)))
print("Grid parameters that effect star2_binary_controls: {0}".format(','.join(grid_params_star2_binary_controls)))
print("Grid parameters that effect star2_binary_job: {0}".format(','.join(grid_params_star2_binary_job)))
else:
print("Grid parameters that effect binary_controls: {0}".format(','.join(grid_params_binary_controls)))
print("Grid parameters that effect binary_job: {0}".format(','.join(grid_params_binary_job)))
print("Grid parameters that effect star1_binary_controls: {0}".format(','.join(grid_params_star1_binary_controls)))
print("Grid parameters that effect star1_binary_job: {0}".format(','.join(grid_params_star1_binary_job)))
print("Grid parameters that effect star2_binary_controls: {0}".format(','.join(grid_params_star2_binary_controls)))
print("Grid parameters that effect star2_binary_job: {0}".format(','.join(grid_params_star2_binary_job)))
# calculate some columns for convenience
if ('m2' in grid.columns) and ('m1' in grid.columns):
grid['q'] = grid['m2']/grid['m1']
if 'initial_period_in_days' in grid.columns:
grid['log_p'] = np.log10(grid['initial_period_in_days'])
# calculate some columns for convenience
if ('initial_star_2_mass' in grid.columns) and ('initial_star_1_mass' in grid.columns):
grid['q'] = grid['initial_star_2_mass']/grid['initial_star_1_mass']
if 'initial_period_days' in grid.columns:
grid['log_p'] = np.log10(grid['initial_period_days'])
# if we are using a job array or just running a single point of grid
if args.grid_point_index is not None:
# per each type of grid value (i.e. binary control or job or star1 control or job etc)
# create a dictionary of the column name and value for this grid point
binary_controls = grid.iloc[args.grid_point_index][grid_params_binary_controls].to_dict()
binary_job = grid.iloc[args.grid_point_index][grid_params_binary_job].to_dict()
star1_binary_controls = grid.iloc[args.grid_point_index][grid_params_star1_binary_controls].to_dict()
star1_binary_job = grid.iloc[args.grid_point_index][grid_params_star1_binary_job].to_dict()
star2_binary_controls = grid.iloc[args.grid_point_index][grid_params_star2_binary_controls].to_dict()
star2_binary_job = grid.iloc[args.grid_point_index][grid_params_star2_binary_job].to_dict()
grid_param_dict = dict(ChainMap({}, binary_controls, binary_job, star1_binary_controls, star1_binary_job, star2_binary_controls,star2_binary_job))
mesa_result = run_grid_point(grid, star1_formation, star2_formation,
grid_param_dict, binary_controls, binary_job,
star1_binary_controls, star1_binary_job,
star2_binary_controls, star2_binary_job, args)
if not mesa_result.empty:
mesa_results_file = os.path.join(args.output_directory, 'grid_results.csv')
if not os.path.isfile(mesa_results_file):
with open(mesa_results_file, 'w') as f:
mesa_result.to_csv(f, index=False)
else:
with open(mesa_results_file, 'a') as f:
mesa_result.to_csv(f, index=False, header=False)
# If we are running a dynamic grid and/or using a single MPI task to manage all grid points
else:
if args.grid_type == 'fixed':
pass
elif args.grid_type == 'dynamic':
if rank == 0:
do_root_process_logic(comm, grid, args, grid_params_binary_controls, grid_params_binary_job,
grid_params_star1_binary_controls, grid_params_star1_binary_job,
grid_params_star2_binary_controls, grid_params_star2_binary_job)
else:
while True:
do_child_process_logic(comm, grid, args, star1_formation=star1_formation, star2_formation=star2_formation)