mesaport.ProjectOps.ops_helper

  1import subprocess
  2import shlex
  3import sys, os, time
  4from pathlib import Path
  5import shutil
  6from rich import print
  7import traceback
  8
  9from ..Access.support import *
 10from ..Access.access_helper import toFortranType, toPythonType
 11from ..Access import MesaAccess, GyreAccess
 12
 13def check_exists(exists, projName):
 14        """Checks if the project exists."""
 15        if not exists:
 16            raise FileNotFoundError(f"Project '{projName}' does not exist. Please create it first.")
 17
 18
 19def run_subprocess(commands, wdir, silent=True, runlog='', status=None, 
 20                    filename="", data_format="FGONG", parallel=False, 
 21                    gyre_in=None, gyre_input_params=None, trace=None, env=None):
 22    
 23    """
 24    Runs a subprocess with the given commands.
 25
 26    Args:
 27    commands (str): The commands to be run.
 28    wdir (str): The working directory.
 29    silent (bool, optional): If True, the output of the subprocess is not printed. Defaults to True.
 30    runlog (str, optional): The path to the runlog file. Defaults to ''.
 31    status (rich.status.Status, optional): The status bar. Defaults to None.
 32    gyre (bool, optional): If True, the subprocess is a GYRE run. Defaults to False.
 33    filename (str, optional): The name of the file to be used in the GYRE run. Defaults to "".
 34    data_format (str, optional): The data format of the file to be used in the GYRE run. Defaults to "FGONG".
 35    parallel (bool, optional): If True, the subprocess is a parallel GYRE run. Defaults to False.
 36    gyre_in (str, optional): The path to the GYRE input file. Defaults to "gyre.in".
 37    gyre_input_params (dict, optional): A dictionary with the parameters to be changed in the GYRE input file. Defaults to None.
 38    trace (str or list, optional): The trace to be followed in the GYRE run. Defaults to None.
 39    env (dict, optional): The environment variables. Defaults to None. 
 40                Pass os.environ.copy() to use the current environment. Or pass a dictionary with the environment variables to be used.
 41    """   
 42    if gyre_in is not None:
 43        try:
 44            gyre_obj = GyreAccess()
 45            if parallel:
 46                # profile_num = filename.split('/')[-1].split(f".{data_format}")[0]
 47                profile_stem = Path(filename).stem.split(".data")[0]
 48                new_gyre_in = os.path.join(wdir, f"gyre{profile_stem}.in")
 49                gyre_obj.modify_gyre_params(wdir, filename, data_format, gyre_in=new_gyre_in)
 50                gyre_obj.set(arg=gyre_input_params, wdir=wdir, gyre_in=new_gyre_in)
 51                time.sleep(1)
 52
 53                # Update gyre_in to the new file
 54                gyre_in = new_gyre_in
 55                commands = commands.replace("gyre.in", f"gyre{profile_stem}.in")
 56                runlog = os.path.join(wdir, f"gyre{profile_stem}.log")
 57            else:
 58                new_gyre_in = os.path.join(wdir, "gyre.in")
 59                shutil.copyfile(gyre_in, new_gyre_in)
 60                gyre_in = new_gyre_in
 61                gyre_obj.modify_gyre_params(wdir, filename, data_format, gyre_in=gyre_in)
 62                gyre_obj.set(arg=gyre_input_params, wdir=wdir, gyre_in=gyre_in)
 63        except Exception as e:
 64            print(traceback.format_exc())
 65            print(f"Error: {e}")
 66            return False
 67        
 68
 69    evo_terminated = None
 70    termination_code = None
 71    if trace is not None:
 72        trace_values = [None for i in range(len(trace))]
 73    with subprocess.Popen(shlex.split(commands), bufsize=0, cwd=wdir, env=env,
 74        stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True) as proc:
 75        with open(runlog, "a+") as logfile:
 76            for outline in proc.stdout:
 77                logfile.write(outline)
 78                logfile.flush()
 79                if silent is False:
 80                    sys.stdout.write(outline)
 81                elif gyre_in is None:
 82                    if "terminated evolution:" in outline or "ERROR" in outline:
 83                        evo_terminated = 1
 84                        termination_code = outline.split()[-1]
 85                        if "cannot find acceptable model" in outline:
 86                            termination_code = 'cannot find acceptable model'
 87                    if "termination code:" in outline:
 88                        evo_terminated = 0
 89                        termination_code = outline.split()[-1]
 90                    if "photo" in outline and "does not exist" in outline:
 91                        evo_terminated = 1
 92                        termination_code = "photo does not exist"
 93                    if not parallel:
 94                        age = process_outline(outline)
 95                        if age is not None:
 96                            if age < 1/365:
 97                                age_str = f"[b]Age: [cyan]{age*365*24:.4f}[/cyan] hours"
 98                            elif 1/365 < age < 1:
 99                                age_str = f"[b]Age: [cyan]{age*365:.4f}[/cyan] days"
100                            elif 1 < age < 1000:
101                                age_str = f"[b]Age: [cyan]{age:.3f}[/cyan] years"
102                            else:
103                                age_str = f"[b]Age: [cyan]{age:.3e}[/cyan] years"
104                        if trace is not None:
105                            trace_values = process_trace(trace, outline, trace_values)
106                            trace_ = [trace[i] for i in range(len(trace)) if trace_values[i] is not None]
107                            values = [val for val in trace_values if val is not None]
108                            if len(values) > 0 and parallel is False:
109                                trace_str = ""
110                                for i in range(len(trace_)):
111                                    trace_str += f"[b]{trace_[i]}[/b]: [cyan]{values[i]:.5f}[/cyan]\n"
112                                status.update(status=f"[b i cyan3]Running....[/b i cyan3]\n"+age_str+"\n"+trace_str, spinner="moon")
113                            elif age is not None:
114                                status.update(status=f"[b i cyan3]Running....[/b i cyan3]\n"+age_str, spinner="moon")
115                        elif age is not None:
116                            status.update(status=f"[b i cyan3]Running....[/b i cyan3]\n"+age_str, spinner="moon")
117            for errline in proc.stderr:
118                logfile.write(errline)
119                sys.stdout.write(errline)
120            logfile.write( "\n\n"+("*"*100)+"\n\n" )
121        _data, error = proc.communicate()
122    if proc.returncode or error:
123        print('The process raised an error:', proc.returncode, error)
124        return False
125    elif evo_terminated == 1:
126        return termination_code, None
127    else:
128        if gyre_in is None:
129            age = 0
130            with open(runlog, "r") as logfile:
131                for line in logfile.readlines():
132                    age_ = process_outline(line)
133                    if age_ is not None:
134                        age = age_
135            return termination_code, age
136        else:
137            working_dir = wdir.replace("LOGS", "")
138            with open(f'{working_dir}/gyre.log', 'a+') as f:
139                f.write(f"Done with {filename}.\n")
140            if parallel:
141                os.remove(gyre_in)
142            return True
143
144def process_outline(outline):
145    try:
146        keyword1 = outline.split()[-1]
147        keyword2 = outline.split()[-2] + " " + outline.split()[-1]
148        keyword3 = outline.split()[-3] + " " + outline.split()[-2] + " " + outline.split()[-1]
149        if keyword1 in dt_limit_values or keyword2 in dt_limit_values or keyword3 in dt_limit_values:
150            return float(outline.split()[0])
151        else:
152            return None
153    except:
154        return None
155
156def setup_trace(trace, work_dir):
157    if isinstance(trace, str):
158        trace = [trace]
159    elif isinstance(trace, list):
160        pass
161    else:
162        raise TypeError("Trace must be a string or a list of strings.")
163    star = MesaAccess(work_dir)
164    num_trace_history_values = star.get("num_trace_history_values")
165    if num_trace_history_values is None:
166        num_trace_history_values = star.getDefault("num_trace_history_values")
167    for tr in trace:
168        exists = False
169        for i in range(num_trace_history_values+1):
170            if tr == star.get(f'trace_history_value_name({i+1})'):
171                exists = True
172                break
173        if not exists:
174            num_trace_history_values += 1
175            star.set({f'trace_history_value_name({num_trace_history_values})': f'{tr}'})
176    star.set({'num_trace_history_values': num_trace_history_values})
177
178
179def process_trace(trace, outline, values):
180    if isinstance(trace, str):
181        trace = [trace]
182    splitline = outline.split()
183    for i in range(len(trace)):
184        if trace[i] in splitline:
185            try:
186                values[i] = float(toPythonType(splitline[1]))
187            except:
188                pass
189    return values
190
191dt_limit_values = ['burn steps', 'Lnuc', 'Lnuc_cat', 'Lnuc_H', 'Lnuc_He', 'lgL_power_phot', 'Lnuc_z', 'bad_X_sum',
192                  'dH', 'dH/H', 'dHe', 'dHe/He', 'dHe3', 'dHe3/He3', 'dL/L', 'dX', 'dX/X', 'dX_nuc_drop', 'delta mdot',
193                  'delta total J', 'delta_HR', 'delta_mstar', 'diff iters', 'diff steps', 'min_dr_div_cs', 'dt_collapse',
194                  'eps_nuc_cntr', 'error rate', 'highT del Ye', 'hold', 'lgL', 'lgP', 'lgP_cntr', 'lgR', 'lgRho', 'lgRho_cntr',
195                  'lgT', 'lgT_cntr', 'lgT_max', 'lgT_max_hi_T', 'lgTeff', 'dX_div_X_cntr', 'lg_XC_cntr', 'lg_XH_cntr', 
196                  'lg_XHe_cntr', 'lg_XNe_cntr', 'lg_XO_cntr', 'lg_XSi_cntr', 'XC_cntr', 'XH_cntr', 'XHe_cntr', 'XNe_cntr',
197                  'XO_cntr', 'XSi_cntr', 'log_eps_nuc', 'max_dt', 'neg_mass_frac', 'adjust_J_q', 'solver iters', 'rel_E_err',
198                  'varcontrol', 'max increase', 'max decrease', 'retry', 'b_****']