diff --git a/pyproject.toml b/pyproject.toml index 174c4c6..639d0f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ python = ">=3.8, <3.11" astropy = "^5.2.2" autograd = "^1.5" celerite = "^0.4.2" +colorama = "^0.4.6" corner = "^2.2.1" emcee = "^3.1.4" matplotlib = "^3.7.1" diff --git a/src/pypetal/drw_rej/module.py b/src/pypetal/drw_rej/module.py index e27dc4e..5a3ed5a 100644 --- a/src/pypetal/drw_rej/module.py +++ b/src/pypetal/drw_rej/module.py @@ -7,7 +7,7 @@ from pypetal.drw_rej.utils import drw_flag from pypetal.utils import defaults -from pypetal.utils.petalio import write_data +from pypetal.utils.petalio import write_data, print_subheader #For multiprocessing @@ -75,22 +75,18 @@ def drw_rej_tot(cont_fname, line_fnames, line_names, output_dir, else: clip_str = clip - txt_str = """ -Performing DRW rejection ------------------------- -jitter: {} -nsig: {} -nwalker: {} -nburn: {} -nchain: {} -clip: {} -reject_data: {} -use_for_javelin: {} ------------------------- - """.format( jitter, nsig, nwalker, nburn, - nchain, clip_str, reject_data, use_for_javelin) - - print(txt_str) + print_dict = { + 'jitter': jitter, + 'nsig': nsig, + 'nwalker': nwalker, + 'nburn': nburn, + 'nchain': nchain, + 'clip': clip_str, + 'reject_data': reject_data, + 'use_for_javelin': use_for_javelin + } + + print_subheader('Performing DRW Rejection', 35, print_dict) sigmas = [] diff --git a/src/pypetal/pipeline.py b/src/pypetal/pipeline.py index 42f3a67..67cbb80 100644 --- a/src/pypetal/pipeline.py +++ b/src/pypetal/pipeline.py @@ -10,7 +10,7 @@ from pypetal.pyroa.module import pyroa_tot from pypetal.pyzdcf.module import pyzdcf_tot from pypetal.utils import defaults -from pypetal.utils.petalio import make_directories, write_data +from pypetal.utils.petalio import make_directories, write_data, print_header, print_error from pypetal.weighting.module import run_weighting_tot @@ -81,7 +81,8 @@ def run_pipeline(output_dir, arg2, output_dir = os.path.abspath(output_dir) + r'/' if arg2 is None: - raise Exception('Please provide a list of light curve filenames or the light curves themselves') + print_error('ERROR: Please provide a list of light curve filenames or the light curves themselves') + raise Exception if not isinstance(arg2[0], str): @@ -105,16 +106,6 @@ def run_pipeline(output_dir, arg2, - - if len(fnames) < 2: - print('ERROR: Requires at least two light curves to run pipeline.') - return {} - - if len(line_names) != len(fnames): - print('ERROR: Must have the same number of line names as light curves.') - return {} - - cont_fname = fnames[0] line_fnames = fnames[1:] @@ -124,6 +115,18 @@ def run_pipeline(output_dir, arg2, #Read in general kwargs general_kwargs = defaults.set_general(kwargs, fnames) + if general_kwargs['verbose']: + print_header('RUNNING PYPETAL') + + + if len(fnames) < 2: + print_error('ERROR: Requires at least two light curves to run pipeline.') + return {} + + if len(line_names) != len(fnames): + print_error('ERROR: Must have the same number of line names as light curves.') + return {} + #Get "reject_data" _, _, _, _, _, _, reject_data, _ = defaults.set_drw_rej(drw_rej_params, fnames) @@ -286,6 +289,9 @@ def run_pipeline(output_dir, arg2, if not isinstance(arg2[0], str): import shutil shutil.rmtree( output_dir + 'input_lcs/' ) + + if general_kwargs['verbose']: + print_header('RUN FINISHED') return tot_res @@ -353,7 +359,8 @@ def run_weighting(output_dir, line_names, if (not run_pyccf) & (not run_javelin) & (not run_pyroa): - raise Exception('ERROR: Either JAVELIN, pyCCF, or PyROA must be run before weighting can be done.') + print_error('ERROR: Either JAVELIN, pyCCF, or PyROA must be run before weighting can be done.') + raise Exception output_dir = os.path.abspath(output_dir) + r'/' @@ -364,6 +371,9 @@ def run_weighting(output_dir, line_names, #Read in general kwargs general_kwargs = defaults.set_general(kwargs, fnames) + if general_kwargs['verbose']: + print_header('RUNNING PYPETAL WEIGHTING') + #Get "interp" interp, _, _, _, _, _ = defaults.set_pyccf(pyccf_params) @@ -419,4 +429,8 @@ def run_weighting(output_dir, line_names, pyroa_params=pyroa_params, general_kwargs=general_kwargs, weighting_params=weighting_params) + if general_kwargs['verbose']: + print_header('RUN FINISHED') + + return res diff --git a/src/pypetal/pyccf/module.py b/src/pypetal/pyccf/module.py index e112218..3fd6fb8 100644 --- a/src/pypetal/pyccf/module.py +++ b/src/pypetal/pyccf/module.py @@ -6,7 +6,7 @@ from pypetal.pyccf.plotting import plot_pyccf_results from pypetal.pyccf.utils import get_pyccf_lags from pypetal.utils import defaults -from pypetal.utils.petalio import write_data +from pypetal.utils.petalio import write_data, print_subheader #For multiprocessing @@ -56,22 +56,18 @@ def pyccf_tot(cont_fname, line_fnames, line_names, output_dir, lag_bounds_str = 'array' else: lag_bounds_str = lag_bounds - - txt_str = """ -Running pyCCF ------------------ -lag_bounds: {} -interp: {} -nsim: {} -mcmode: {} -sigmode: {} -thres: {} -nbin: {} ------------------ - """.format( lag_bounds_str, interp, nsim, - mcmode, sigmode, thres, nbin) - - print(txt_str) + + print_dict = { + 'lag_bounds': lag_bounds_str, + 'interp': interp, + 'nsim': nsim, + 'mcmode': mcmode, + 'sigmode': sigmode, + 'thres': thres, + 'nbin': nbin + } + + print_subheader('Running pyCCF', 35, print_dict) res_tot = [] diff --git a/src/pypetal/pyccf/utils.py b/src/pypetal/pyccf/utils.py index 761791f..a9e31ab 100644 --- a/src/pypetal/pyccf/utils.py +++ b/src/pypetal/pyccf/utils.py @@ -3,6 +3,7 @@ import numpy as np import scipy.stats as sst +from pypetal.utils.petalio import print_warning #For multiprocessing @@ -612,7 +613,7 @@ def get_pyccf_lags(fname1, fname2, if np.any( np.abs(lag_bounds) > baseline ): - print('Lag bounds are larger than baseline') + print_warning('WARNING: Lag bounds are larger than baseline') #Run algorithm diff --git a/src/pypetal/pyroa/module.py b/src/pypetal/pyroa/module.py index 9d38681..7c4832a 100644 --- a/src/pypetal/pyroa/module.py +++ b/src/pypetal/pyroa/module.py @@ -4,6 +4,7 @@ pyroa_corner_plot, pyroa_trace_plot) from pypetal.pyroa.utils import run_pyroa from pypetal.utils import defaults +from pypetal.utils.petalio import print_subheader def pyroa_tot(cont_fname, line_fnames, line_names, output_dir, @@ -27,24 +28,20 @@ def pyroa_tot(cont_fname, line_fnames, line_names, output_dir, if verbose: - txt_str = """ -Running PyROA ----------------- -nburn: {} -nchain: {} -init_tau: {} -subtract_mean: {} -div_mean: {} -add_var: {} -delay_dist: {} -psi_types: {} -together: {} -objname: {} ----------------- - """.format(nburn, nchain, init_tau, subtract_mean, div_mean, add_var, - delay_dist, psi_types, together, objname) - - print(txt_str) + print_dict = { + 'nchain': nchain, + 'nburn': nburn, + 'init_tau': init_tau, + 'subtract_mean': subtract_mean, + 'div_mean': div_mean, + 'add_var': add_var, + 'delay_dist': delay_dist, + 'psi_types': psi_types, + 'together': together, + 'objname': objname + } + print_subheader('Running PyROA', 35, print_dict) + tot_fnames = np.hstack( [ [cont_fname], line_fnames ] ) diff --git a/src/pypetal/pyzdcf/module.py b/src/pypetal/pyzdcf/module.py index 287278a..1052908 100644 --- a/src/pypetal/pyzdcf/module.py +++ b/src/pypetal/pyzdcf/module.py @@ -10,6 +10,7 @@ import pypetal.pyzdcf.utils as utils from pypetal.pyzdcf.plotting import plot_pyzdcf_results from pypetal.utils import defaults +from pypetal.utils.petalio import print_subheader, print_error #For multiprocessing @@ -58,37 +59,31 @@ def pyzdcf_tot(cont_fname, line_fnames, line_names, output_dir, #------------------------------------------- if (run_plike) & (plike_dir is None): - print('Error: plike_dir must be specified if run_plike=True') - print('Skipping PLIKE') + print_error('ERROR: plike_dir must be specified if run_plike=True') + print_error('Skipping PLIKE') run_plike = False input_dir = os.path.dirname( os.path.realpath(cont_fname) ) for i in range(len(line_fnames)): if input_dir != os.path.dirname( os.path.realpath(line_fnames[i]) ): - print('ERROR: All light curve files must be in the same directory') + print_error('ERROR: All light curve files must be in the same directory') return {} if verbose: - txt_str = """ -Running pyZDCF ----------------------- -nsim: {} -minpts: {} -uniform_sampling: {} -omit_zero_lags: {} -sparse: {} -prefix: {} -run_plike: {} -plike_dir: {} ----------------------- - """.format( nsim, minpts, uniform_sampling, omit_zero_lags, - sparse, prefix, run_plike, plike_dir) - - print(txt_str) - + print_dict = { + 'nsim': nsim, + 'minpts': minpts, + 'uniform_sampling': uniform_sampling, + 'omit_zero_lags': omit_zero_lags, + 'sparse': sparse, + 'prefix': prefix, + 'run_plike': run_plike, + 'plike_dir': plike_dir + } + print_subheader('Running pyZDCF', 35, print_dict) input_dir += r'/' diff --git a/src/pypetal/pyzdcf/utils.py b/src/pypetal/pyzdcf/utils.py index ec096b3..2ffcdce 100644 --- a/src/pypetal/pyzdcf/utils.py +++ b/src/pypetal/pyzdcf/utils.py @@ -3,6 +3,7 @@ import numpy as np from pyzdcf import pyzdcf +from pypetal.utils.petalio import print_error def run_plike(dcf_fname, lag_bounds, plike_dir, verbose=False): @@ -52,7 +53,8 @@ def run_plike(dcf_fname, lag_bounds, plike_dir, verbose=False): #Make sure there are two lag bounds if len(lag_bounds) != 2: - print('Lag bounds: ', lag_bounds) + print_error('ERROR: Must provide two lag bounds.') + print_error('Lag bounds: ', lag_bounds) raise ValueError('Must provide two lag bounds.') diff --git a/src/pypetal/utils/detrending.py b/src/pypetal/utils/detrending.py index a0c666d..c478aa0 100644 --- a/src/pypetal/utils/detrending.py +++ b/src/pypetal/utils/detrending.py @@ -4,7 +4,7 @@ from linmix import LinMix from pypetal.utils.defaults import set_detrend -from pypetal.utils.petalio import write_data +from pypetal.utils.petalio import write_data, print_subheader mpl.rcParams['xtick.minor.visible'] = True mpl.rcParams['xtick.top'] = True @@ -220,20 +220,15 @@ def detrend_tot(output_dir, cont_fname, line_fnames, line_names, general_kwargs, else: parallelize = False - - txt = """ -Running detrending -------------------- -parallelize: {} -K: {} -nchains: {} -miniter: {} -maxiter: {} -------------------- - """.format(parallelize, K, nchain, miniter, maxiter) - if verbose: - print(txt) + print_dict = { + 'parallelize': parallelize, + 'K': K, + 'nchain': nchain, + 'miniter': miniter, + 'maxiter': maxiter + } + print_subheader('Running detrending', 35, print_dict) m_vals = [] diff --git a/src/pypetal/utils/load.py b/src/pypetal/utils/load.py index afce840..631dc91 100644 --- a/src/pypetal/utils/load.py +++ b/src/pypetal/utils/load.py @@ -5,6 +5,7 @@ from astropy.table import Table, vstack from pypetal.pyroa.utils import MyFit +from pypetal.utils.petalio import print_subheader, print_warning, print_error def get_line_names(main_dir): @@ -88,7 +89,7 @@ def get_modules(main_dir): n_pyccf = len( np.argwhere(has_pyccf).T[0] ) if ( n_pyccf != n_lnc ) & ( n_pyccf > 0 ): - print('pyCCF was not completed for all lines, so will assume run_pyccf=False') + print_warning('WARNING: pyCCF was not completed for all lines, so will assume run_pyccf=False') run_pyccf = False else: run_pyccf = np.any( has_pyccf ) @@ -105,7 +106,7 @@ def get_modules(main_dir): n_pyzdcf = len( np.argwhere(has_pyzdcf).T[0] ) if ( n_pyzdcf != n_lnc ) & ( n_pyzdcf > 0 ): - print('pyZDCF was not completed for all lines, so will assume run_pyzdcf=False') + print_warning('WARNING: pyZDCF was not completed for all lines, so will assume run_pyzdcf=False') run_pyzdcf = False else: run_pyzdcf = np.any( has_pyzdcf ) @@ -129,7 +130,7 @@ def get_modules(main_dir): n_pyroa = len( np.argwhere(has_pyroa).T[0] ) if ( n_pyroa != n_lnc ) & ( n_pyroa > 0 ): - print('PyROA was not completed for all lines, so will assume run_pyroa=False') + print_warning('WARNING: PyROA was not completed for all lines, so will assume run_pyroa=False') else: run_pyroa = np.any( has_pyroa ) @@ -152,7 +153,7 @@ def get_modules(main_dir): n_javelin = len( np.argwhere(has_javelin).T[0] ) if ( n_javelin != n_lnc ) & ( n_javelin > 0 ): - print('JAVELIN was not completed for all lines, so will assume run_javelin=False') + print_warning('WARNING: JAVELIN was not completed for all lines, so will assume run_javelin=False') else: run_javelin = np.any( has_javelin ) @@ -169,7 +170,7 @@ def get_modules(main_dir): n_weighting = len( np.argwhere(has_weighting).T[0] ) if ( n_weighting != n_lnc ) & ( n_weighting > 0 ): - print('Weighting was not completed for all lines, so will assume run_weighting=False') + print_warning('WARNING: Weighting was not completed for all lines, so will assume run_weighting=False') run_weighting = False else: run_weighting = np.any( has_weighting ) @@ -854,20 +855,18 @@ def load(main_dir, modules=None, verbose=False): run_javelin = ('javelin' in modules) run_weighting = ('weighting' in modules) - txt = """ -Prior pyPetal run ---------------------- -DRW Rejection: {} -pyCCF: {} -pyZDCF: {} -PyROA: {} -JAVELIN: {} -Weighting: {} ---------------------- -""".format(run_drw_rej, run_pyccf, run_pyzdcf, run_pyroa, run_javelin, run_weighting) - + if verbose: - print(txt) + print_dict = { + 'DRW Rejection': run_drw_rej, + 'pyCCF': run_pyccf, + 'pyZDCF': run_pyzdcf, + 'PyROA': run_pyroa, + 'JAVELIN': run_javelin, + 'Weighting': run_weighting + } + print_subheader('Prior pyPetal Run', print_dict) + drw_rej_res = {} diff --git a/src/pypetal/utils/petalio.py b/src/pypetal/utils/petalio.py index fb7dd7c..73ad739 100644 --- a/src/pypetal/utils/petalio.py +++ b/src/pypetal/utils/petalio.py @@ -6,6 +6,90 @@ from astropy.io import fits from astropy.table import Table +from colorama import init, Fore +init(autoreset=True) + +############################################## +# PRINTING + +mcolor = Fore.BLUE +acolor = Fore.GREEN +wcolor = Fore.YELLOW +ecolor = Fore.RED + +def print_hashstr(): + hashstr = "".join( ['#']*80 ) + print( mcolor + hashstr ) + return + +def print_linestr(n): + linestr = "".join( ['\u25AC']*n ) + print( mcolor + linestr ) + return + + +def print_header(header): + nhead = len(header) + midpt = int(80/2) + + hstart = midpt - int(nhead/2) - 2 + hend = midpt + (nhead-int(nhead/2)) + 2 + + nstart = hstart + nend = 80 - hend + assert nstart + 2 + nhead + 2 + nend == 80 + + hashstr1 = "".join( ['#']*nstart ) + hashstr2 = "".join( ['#']*nend ) + + print('') + print_hashstr() + print(mcolor + hashstr1, end='') + print(' ', end='') + print( acolor + header, end='') + print(' ', end='') + print(mcolor + hashstr2) + print_hashstr() + print('') + return + + +def print_subheader(subheader, n, d): + if n < len(subheader) + 5: + n = len(subheader) + 5 + + print_linestr(n) + print(acolor + subheader) + print_linestr(n) + print_subheader_dict(d) + print_linestr(n) + print('') + return + + +def print_subheader_dict(d): + n = 15 + 5 + + keys = list(d.keys()) + vals = list(d.values()) + for i in range(len(keys)): + print(keys[i] + ':', end='') + print( ''.join([' ']*(n-len(keys[i])-1)), end='' ) + print('{}'.format(vals[i])) + + return + +def print_warning(warning): + print(wcolor + '*** ', end='') + print(wcolor + warning) + return + +def print_error(error): + print(ecolor + '*** ', end='') + print(ecolor + error) + return + +############################################## def make_directories(output_dir, fnames, line_names, run_drw_rej, run_pyccf, run_pyzdcf, run_pyroa,