mesaport.ProjectOps.project_ops

  1import os
  2import shutil
  3import psutil
  4import subprocess
  5import glob
  6from itertools import repeat
  7import traceback
  8
  9from rich import print, progress, prompt, status
 10progress_columns = (progress.SpinnerColumn(spinner_name="moon"),
 11                    progress.MofNCompleteColumn(),
 12                    *progress.Progress.get_default_columns(),
 13                    progress.TimeElapsedColumn())
 14import multiprocessing as mp
 15
 16from ..Access import MesaAccess, GyreAccess, MesaEnvironmentHandler
 17from . import ops_helper
 18from . import istarmap
 19
 20class ProjectOps:
 21    """This class handles MESA project operations.
 22    """    
 23    def __init__(self, name='work', astero=False, binary=False):
 24        """Constructor for ProjectOps class.
 25
 26        Args:
 27            name (str, optional): Name of the project. Defaults to 'work'.
 28            binary (bool, optional): True for a binary star system. Defaults to False.
 29        """        
 30        self.projName = name
 31        self.binary = binary
 32        self.astero = astero
 33        self.envObject = MesaEnvironmentHandler()
 34        if self.binary:
 35            self.defaultWork = os.path.join(self.envObject.mesaDir, 'binary/work')
 36        elif self.astero:
 37            self.defaultWork = os.path.join(self.envObject.mesaDir, 'astero/work')
 38        else:
 39            self.defaultWork = os.path.join(self.envObject.mesaDir, 'star/work')
 40
 41        if os.path.exists(self.projName):
 42            self.exists = True               ## Proj already present flag
 43            if os.path.isabs(self.projName):
 44                self.work_dir = self.projName
 45            else:
 46                self.work_dir = os.path.abspath(os.path.join(os.getcwd(), self.projName))
 47        else:
 48            self.exists = False
 49        
 50    
 51
 52    def create(self, overwrite=None, clean=None): 
 53        """Creates a new MESA project.
 54
 55        Args:
 56            overwrite (bool, optional): Overwrite the existing project. Defaults to None.
 57            clean (bool, optional): Clean the existing project. Defaults to None.
 58        """   
 59
 60        def useExisting():
 61            """A helper function to use the existing project."""         
 62            if not prompt.Confirm.ask(f"Use the already existing '{self.projName}' project as it is?", default=False):
 63                raise ValueError("Aborting!!! No project specified.")
 64
 65        def cleanCheck():
 66            """A helper function to check if the user wants to clean the existing project."""
 67            if clean is None:
 68                if prompt.Confirm.ask(f"Clean the existing '{self.projName}' project for re-use?", default=False):
 69                    self.clean()
 70                else:
 71                    useExisting()
 72            elif clean is True:
 73                self.clean()
 74            elif clean is False:
 75                print(f"Using the already existing '{self.projName}' project as it is.")
 76            else:
 77                raise ValueError("Invalid input for argument 'clean'.")
 78        
 79        def writeover():
 80            """A helper function to overwrite the existing project."""
 81            try:
 82                shutil.rmtree(self.projName)
 83                shutil.copytree(self.defaultWork, self.projName)
 84
 85            except shutil.Error:
 86                raise Exception(f"Could not overwrite the existing '{self.projName}' project!")
 87
 88        if self.exists is True:
 89            if os.path.isabs(self.projName):
 90                self.work_dir = self.projName
 91            else:
 92                self.work_dir = os.path.abspath(os.path.join(os.getcwd(), self.projName))
 93            if overwrite is True:
 94                writeover()
 95            elif overwrite is False:
 96                cleanCheck()
 97            elif overwrite is None:
 98                print(f"Mesa project named '{self.projName}' already exists!")
 99                if not prompt.Confirm.ask(f"Use the already existing '{self.projName}' project as it is?", default=False):
100                    if prompt.Confirm.ask("Do you wish to overwrite?", default=False):
101                        writeover()
102                    else:
103                        cleanCheck()
104            else:
105                raise ValueError("Invalid input for argument 'overwrite'.") 
106        else:
107            try:
108                shutil.copytree(self.defaultWork, self.projName)
109                if os.path.isabs(self.projName):
110                    self.work_dir = self.projName
111                else:
112                    self.work_dir = os.path.abspath(os.path.join(os.getcwd(), self.projName))
113                self.exists = True
114            except shutil.Error:
115                raise Exception(f"Could not create the project '{self.projName}'!")
116    
117    def delete(self):
118        """Deletes the project.
119        """        
120        if self.exists is True:
121            shutil.rmtree(self.work_dir)
122            print(f"Deleted project '{self.projName}'.")
123        else:
124            print(f"Project '{self.projName}' does not exist.")
125
126    def clean(self):
127        """Cleans the project.
128
129        Raises:
130            Exception: If the clean fails.
131        """        
132        ops_helper.check_exists(self.exists, self.projName)
133        ## clean files are missing a shebang (#!/bin/bash) and hence need to be run with bash
134        res = subprocess.call('/bin/bash ./clean', cwd=self.work_dir, shell=True, stderr=subprocess.STDOUT)
135        runlog = os.path.join(self.work_dir, "runlog")
136        if os.path.exists(runlog):
137            os.remove(runlog)
138        if res!=0:
139            raise Exception(f"Clean failed! Returned non-zero exit code ({res})")
140        else:
141            print("Clean successful.\n")
142            
143
144    def make(self, silent=False):
145        """Makes the project.
146
147        Args:
148            silent (bool, optional): Run the command silently. Defaults to False.
149
150        Raises:
151            Exception: If the make fails.
152        """        
153        ops_helper.check_exists(self.exists, self.projName)
154        if silent:
155            res = subprocess.call('./mk', cwd=self.work_dir, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
156        else:
157            with status.Status("[b i cyan3]Making...", spinner="moon"):
158                res = subprocess.call('./mk', cwd=self.work_dir, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
159        if res!=0:
160            raise Exception(f"Make failed! Returned non-zero exit code ({res})")
161        else:    
162            print("Make successful.\n")
163
164
165    
166    def run(self, silent=True, logging=True, parallel=False, trace=None, env=os.environ.copy()):
167        """
168        Runs the project.
169        Args:
170            silent (bool, optional): Run the command silently. Defaults to True.
171            logging (bool, optional): Log the run. Defaults to True.
172            parallel (bool, optional): Run in parallel. Defaults to False.
173            trace (list of str, optional): Trace specific history variables. Defaults to None.
174            env (dict, optional): Environment variables. Defaults to os.environ.copy().
175
176        Raises:
177            Exception: If the project is not made yet.
178            ValueError: If the input for argument 'silent' is invalid.
179            Exception: If the run fails.
180        
181        Returns: (If run is successful)
182            termination_code (str): Termination code.
183            age (float): Age of the star in years.
184        """        
185        if trace is not None:
186            ops_helper.setup_trace(trace, self.work_dir)
187        ops_helper.check_exists(self.exists, self.projName)
188        if logging:
189            runlog = os.path.join(self.work_dir, "run.log")
190        else:
191            runlog = os.devnull
192        if not os.path.exists(os.path.join(self.work_dir, "star")) and \
193            not os.path.exists(os.path.join(self.work_dir, "binary")):
194            raise Exception("Aborting! Run 'make()' first.")
195        else:
196            if silent not in [True, False]:
197                raise ValueError("Invalid input for argument 'silent'")
198            else:
199                if parallel:
200                    res = ops_helper.run_subprocess(commands='./rn', wdir=self.work_dir, 
201                                silent=silent, runlog=runlog, parallel=True, trace=trace, env=env)
202                else:
203                    with status.Status("[b i cyan3]Running...", spinner="moon") as status_:
204                        res = ops_helper.run_subprocess(commands='./rn', wdir=self.work_dir, 
205                                    silent=silent, runlog=runlog, status=status_, trace=trace, env=env)
206            if res == False:
207                raise Exception("Run failed! Check runlog.")
208            else:
209                termination_code, age = res
210                if age is not None:
211                    print("Run successful.\n")
212                    return termination_code, age
213                else:
214                    print("Run unsuccessful.\n")
215                    return termination_code, None    
216
217        
218    
219    def resume(self, photo=None, silent=True, target=None, logging=True, parallel=False, trace=None, env=os.environ.copy()):
220        """Resumes the run from a given photo.
221
222        Args:
223            photo (str, optional): Photo name from which the run is to be resumed. 
224                                If None, the last photo is used. Defaults to None.
225            silent (bool, optional): Run the command silently. Defaults to True.
226            target (str, optional): Target photo name. Defaults to None.
227            logging (bool, optional): Log the run. Defaults to True.
228            parallel (bool, optional): Run in parallel. Defaults to False.
229            trace (list of str, optional): Trace specific history variables. Defaults to None.
230            env (dict, optional): Environment variables. Defaults to os.environ.copy().
231
232        Raises:
233            FileNotFoundError: If the photo does not exist.
234            ValueError: If the input for argument 'silent' is invalid.
235
236        Returns: (If run is successful)
237            termination_code (str): Termination code.
238            age (float): Age of the star in years.
239        """
240        if trace is not None:
241            ops_helper.setup_trace(trace, self.work_dir)
242        ops_helper.check_exists(self.exists, self.projName)
243        if logging:
244            runlog = os.path.join(self.work_dir, "run.log")
245        else:
246            runlog = os.devnull
247        if photo == None:
248            if parallel:
249                res = ops_helper.run_subprocess(commands='./re', wdir=self.work_dir, 
250                        silent=silent, runlog=runlog, parallel=True, trace=trace)
251            else:
252                # print(f"[b i  cyan3]Resuming run from the most recent photo.")
253                with status.Status("[b i  cyan3]Resuming run from the most recent photo.\nRunning...", spinner="moon") as status_:
254                    res = ops_helper.run_subprocess(commands=f'./re', wdir=self.work_dir, 
255                            silent=silent, runlog=runlog, status=status_, trace=trace, env=env)
256        else:
257            if self.binary:
258                if target == 'primary':
259                    photo_path = os.path.join(self.work_dir, "photos1", photo)
260                elif target == 'secondary':
261                    photo_path = os.path.join(self.work_dir, "photos2", photo)
262                else:
263                    raise ValueError('''Invalid input for argument 'target'.  
264                                    Please use 'primary' or 'secondary' ''')
265            else:
266                photo_path = os.path.join(self.work_dir, "photos", photo)
267                
268            if not os.path.isfile(photo_path):
269                raise FileNotFoundError(f"Photo '{photo}' could not be exists.")
270            else:
271                if silent not in [True, False]:
272                    raise ValueError("Invalid input for argument 'silent'.")
273                else:
274                    if parallel:
275                        res = ops_helper.run_subprocess(commands=f'./re {photo}', wdir=self.work_dir, 
276                                silent=silent, runlog=runlog, parallel=True, trace=trace, env=env)
277                    else:
278                        with status.Status(f"[b i  cyan3]Resuming run from photo {photo}.\nRunning...", spinner="moon") as status_:
279                            res = ops_helper.run_subprocess(commands=f'./re {photo}', wdir=self.work_dir, 
280                                    silent=silent, runlog=runlog, status=status_, trace=trace, env=env)
281        if res is False:
282            raise Exception("Resume from photo failed! Check runlog.")
283        else:
284            termination_code, age = res
285            if age is not None:
286                print("Run successful.\n")
287                return termination_code, age
288            else:
289                print("Run unsuccessful.\n")
290                return termination_code, None
291
292        
293    def runGyre(self, gyre_in, files='all', wdir=None, data_format="GYRE", silent=True, target=None, logging=True, logfile="gyre.log", 
294                    parallel=False, n_cores=None, gyre_input_params=None, env=os.environ.copy()):
295        """
296        Runs GYRE.
297
298        Arguments:
299            gyre_in (str): GYRE input file.
300            files (str or list of strings, optional): Profile files in the LOGS directory 
301                                            to be processed by GYRE. Defaults to 'all'.
302            wdir (str, optional): Working directory. Defaults to None and uses the project directory.
303            silent (bool, optional): Run the command silently. Defaults to True.
304            target (str, optional): Target star. Defaults to None.
305            logging (bool, optional): Log the output. Defaults to True.
306            logdir (str, optional): Log file name. Defaults to "run.log".
307            parallel (bool, optional): Run GYRE in parallel. Defaults to False.
308            n_cores (int, optional): Number of cores to use. Defaults to None.
309            gyre_input_params (dict or list of dicts, optional): Dictionary of GYRE input parameters.
310                                                                {parameter: value}. Parameter must be a string. Value
311                                                                can be a string, int, float, or bool.
312                                                                List of dictionaries for multiple profiles. 
313                                                                list of dicts must be in the same order as the
314                                                                list of profile files. len(list of dicts) must be
315                                                                equal to len(list of profile files). Defaults to None.
316            env (dict, optional): Environment variables. Defaults to os.environ.copy().
317        Raises:
318            FileNotFoundError: If the GYRE input file does not exist.
319            ValueError: If the input for argument 'silent' is invalid.
320            ValueError: If the input for argument 'files' is invalid.
321        """
322        if wdir is not None:
323            wdir = os.path.abspath(wdir)
324        gyre_in = os.path.abspath(gyre_in)
325
326        if 'GYRE_DIR' in os.environ:
327            gyre_ex = os.path.join(os.environ['GYRE_DIR'], "bin", "gyre")
328        else:
329            raise FileNotFoundError("GYRE_DIR is not set in your enviroment. Be sure to set it properly!!")
330        if self.binary:
331            if target == 'primary':
332                LOGS_dir = os.path.join(self.work_dir, "LOGS1") if wdir is None else wdir
333            elif target == 'secondary':
334                LOGS_dir = os.path.join(self.work_dir, "LOGS2") if wdir is None else wdir
335            else:
336                raise ValueError("""Invalid input for argument 'star'.  
337                                Please use primary or secondary""")
338        else:
339            LOGS_dir = os.path.join(self.work_dir, "LOGS") if wdir is None else wdir
340
341        if wdir is None:
342            wdir = self.work_dir
343
344        if logging:
345            runlog = os.path.join(wdir, logfile)
346        else:
347            runlog = os.devnull
348        
349        
350        if not silent in [True, False]:
351            raise ValueError("Invalid input for argument 'silent'")
352
353        if files == 'all' or isinstance(files, list) or isinstance(files, str):
354            ## ALL FILES
355            if files == 'all':
356                files = []
357                try:
358                    files = sorted(glob.glob(os.path.join(LOGS_dir, f"*.{data_format}")), 
359                                key=lambda x: int(os.path.basename(x).split('.')[0].split('profile')[1]))
360                except ValueError:
361                    files = sorted(glob.glob(os.path.join(LOGS_dir, f"*.{data_format}")))
362                files = [file.split('/')[-1] for file in files]
363                if len(files) == 0:
364                    raise ValueError(f"No {data_format} files found in LOGS directory.")
365                    
366            ## SPECIFIC FILES
367            elif type(files) == list or type(files) == str:
368                if type(files) == str:
369                    files = [files]
370                    gyre_input_params = [gyre_input_params]
371                if len(files) == 0:
372                    raise ValueError("No files provided.")
373                # else:
374                #     for file in files:
375                #         if not os.path.isfile(os.path.join(LOGS_dir, file)) and not os.path.isfile(file):
376                #             raise FileNotFoundError(f"File '{file}' does not exist.")
377                        
378            with open(f'{wdir}/gyre.log', 'a+') as f:
379                    f.write(f"Total {len(files)} profiles to be processed by GYRE.\n\n")
380            if parallel:
381                gyre_input_params = gyre_input_params if gyre_input_params is not None else repeat(None)
382                os.environ['HDF5_USE_FILE_LOCKING'] = 'FALSE'
383                ## copy gyre.in to gyre1.in, gyre2.in, etc. for parallel runs
384                for i, file in enumerate(files):
385                    num = file.split(".")[0]
386                    new_gyre_in = os.path.join(wdir, f"gyre{num}.in")
387                    shutil.copyfile(gyre_in, new_gyre_in)
388                # commands, wdir, 
389                # silent=True, runlog='', 
390                # status=None, filename="", 
391                # data_format="FGONG", parallel=False, gyre_in=None, 
392                # gyre_input_params=None, trace=None, env=None
393                args = (repeat(f'{gyre_ex} gyre.in'), repeat(LOGS_dir),
394                        repeat(silent), repeat(runlog),
395                        repeat(None), files, repeat(data_format),
396                        repeat(True), repeat(gyre_in),
397                        gyre_input_params, repeat(None), repeat(os.environ.copy()))
398                if n_cores is None:
399                    n_cores = psutil.cpu_count(logical=True) 
400                    Pool = mp.Pool
401                    with progress.Progress(*progress_columns) as progressbar:
402                        task = progressbar.add_task("[b i cyan3]Running GYRE...", total=len(files))
403                        n_processes = (n_cores//int(os.environ['OMP_NUM_THREADS']))
404                        with Pool(n_processes) as pool:
405                            gyre_in = os.path.abspath(gyre_in)
406                            try:
407                                for _ in pool.istarmap(ops_helper.run_subprocess, zip(*args)):
408                                    progressbar.advance(task)
409                            except Exception as e:
410                                print(traceback.format_exc())
411                                print(f"Error: {e}")
412                                pool.terminate()
413                else:
414                    try:
415                        from concurrent.futures import ThreadPoolExecutor
416                        n_processes = (n_cores//int(os.environ['OMP_NUM_THREADS']))
417                        with ThreadPoolExecutor(max_workers=n_processes) as executor:
418                            gyre_in = os.path.abspath(gyre_in)
419                            try:
420                                executor.map(ops_helper.run_subprocess, *args)
421                            except Exception as e:
422                                print(traceback.format_exc())
423                                print(f"Error: {e}")
424                                executor.shutdown(wait=False)
425                    except Exception as e:
426                        filenames = glob.glob(os.path.join(LOGS_dir, f"gyreprofile*.log"))
427                        with open(runlog, 'a+') as outfile:
428                            for fname in filenames:
429                                with open(fname) as infile:
430                                    for line in infile:
431                                        outfile.write(line)
432                                os.remove(fname)
433                        print(traceback.format_exc())
434                        print(f"Error: {e}")
435                filenames = glob.glob(os.path.join(LOGS_dir, f"gyreprofile*.log"))
436                with open(runlog, 'a+') as outfile:
437                    for fname in filenames:
438                        with open(fname) as infile:
439                            for line in infile:
440                                outfile.write(line)
441                        os.remove(fname)
442                res = True
443            else:
444                for i, file in enumerate(files):
445                    gyre_input_params_i = gyre_input_params[i] if gyre_input_params is not None else None
446                    res = ops_helper.run_subprocess(f'{gyre_ex} gyre.in', wdir=LOGS_dir, filename=file,
447                        silent=silent, runlog=runlog, status=None, gyre_in=gyre_in, 
448                        data_format=data_format, gyre_input_params=gyre_input_params_i)
449        else:
450            raise ValueError("Invalid input for argument 'files'")
451        if res is False:
452            print("GYRE run failed! Check runlog.")
453        else:
454            print("GYRE run complete!\n")
455        return res