m_sim_setup.py

"""
Setup for single simulations
This file,
1. Gets multiple input params from specs file
2. Generates full parameter dict from specs file 
    and default parameters for all simulations.
5. Saves params file and matrix in cache/
5. Calls m_run.py to run simulation
6. Clears cache files
7. Saves specs file and params dict in 
    data directory for future reference
"""
import subprocess
import os
os.environ["NEURON_MODULE_OPTIONS"] = "-nogui" #Stops no gui warnings in output file
import sim_utils as s_utils
import time
import importlib
from param import mParam
import logging
import shutil
import sys

tstart = time.perf_counter()
args,unk = s_utils.sim_setup_arg_parser().parse_known_args()
#Process specs file
fname = args.specs_file
mod_name = s_utils.get_module_from_path(fname)

#load input parameters
param_file = importlib.import_module(mod_name)
mult_input_params = param_file.generate_mult_input_params()
first_key = next(iter(mult_input_params))
#check if sim_id declared
if not mult_input_params[first_key].get("sim_id",None):
    raise ValueError("sim_id not set in specs file")

#Initialize params dictionary
mult_params = mParam()
mult_params.load_update_mult_params(mult_input_params)
first_key = next(iter(mult_params))
sim_id = mult_params[first_key]["sim_id"]

n_cpus = mult_params[first_key]["n_cpus"] if mult_params[first_key]["n_cpus"] else os.cpu_count()//2

#start logger
logging.basicConfig(handlers=[logging.FileHandler(f"logs/setup_{sim_id}.log",mode="w"),
        logging.StreamHandler()], encoding='utf-8', level=args.verbose,
        format=f'%(asctime)s:%(levelname)s:{sim_id}:%(message)s')

# save parameters to cache for running the simulation
s_utils.json_save(mult_params,f"cache/params_{sim_id}.json")

#create data directory in data_root/
data_root = s_utils.process_data_root(mult_params[first_key]["data_root"])
data_loc = data_root+f"{sim_id}/"

try:
    if args.overwrite_data:
        logging.info(f"Overwriting data at {data_loc}")
        if os.path.exists(data_loc):
            shutil.rmtree(data_loc)
        os.makedirs(data_loc)
    else:
        os.makedirs(data_loc)
except FileExistsError as err:
    logging.error(f"Simulation data already exists at {data_loc}")
    err.add_note(f"Simulation data already exists at {data_loc}. Use -o to overwrite data")
    raise err

logging.debug(f"Created data location at {data_loc}")

# copy specs file to data_loc for future reference
os.system(f"cp {fname} {data_loc}{sim_id}.py")

#run simulation
logging.debug(f"Launching m_run")
verbosity = "-v" if args.verbose == logging.DEBUG else ""
proc=subprocess.run(["mpiexec","-n", f"{n_cpus}",f"{sys.executable}", "m_run.py", 
                     "--sim_id",f"{sim_id}",f"{verbosity}"],check=True)

# save simulation time
t_simulation=round(time.perf_counter() - tstart, 2)
mult_params[first_key]["t_simulation"]= t_simulation
mult_params[first_key]["git_commit_hash"]= s_utils.get_git_commit_hash()

#remove cache file
try:
    os.remove(f"cache/params_{sim_id}.json")
except FileNotFoundError as err:
    logging.debug(f"Cannot remove cached params file: {err}")

# copy params file to data for future reference
param_fname = data_loc + f"{sim_id}.json"
s_utils.json_save(mult_params, param_fname)
logging.info(f"Total time (s): {round(time.perf_counter() - tstart, 2)}")