Added colored output

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"

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 = []

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

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 = []

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

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 ] )
 

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'/'

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.')
 
 

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 = []