PR313 follow-up: proposed config and arguments format

cosipy/pipeline/src/fitting.py

@@ -49,7 +49,7 @@ def build_spectrum(model, pars, par_minvalues, par_maxvalues):
 
 
 
-def get_fit_results(sou,bk,resp_path,ori,l,b,souname,bkname,spectrum):
+def get_fit_results(sou,bk,resp_path,ori,bkname,model):
     """
     Fits the spectrum to the data
     """
@@ -68,39 +68,34 @@ def get_fit_results(sou,bk,resp_path,ori,l,b,souname,bkname,spectrum):
                     nuisance_param=bkg_par,
                     earth_occ=True)  # background parameter
 
-    source = PointSource(souname,  # Name of source (arbitrary, but needs to be unique)
-                         l=l,  # Longitude (deg)
-                         b=b,  # Latitude (deg)
-                         spectral_shape=spectrum)  # Spectral model
-
-    model = Model(source)
     cosi.set_model(model)
     plugins = DataList(cosi)
     like = JointLikelihood(model, plugins, verbose=False)
     like.fit()
     results = like.results
-    expectation = cosi._expected_counts[souname]
+
+    # FIXME: Do no rely on protected variables. _expected_counts is not guaranteed to
+    #   survive refactoring. Revisit this after the interfaces refactoring
+    expectation = None
+    for source_expectation in cosi._expected_counts.values():
+        if expectation is None:
+            expectation = source_expectation.copy()
+        else:
+            expectation += source_expectation
+
     tot_exp_counts=expectation.project('Em').todense().contents + (bkg_par.value * bk.project('Em').todense().contents)
     #
     #
     return results,tot_exp_counts
 
-def get_fit_par(results,souname,bkname):
+def get_fit_par(results):
     """
-    Extracts dictionaries from the fit results
-    par_sou,epar_sou,par_bk,epar_bk
+    Extracts a dictionary whose keys are the free parameters of the model,
+    and values are a tuple with the median and standard deviation.
     """
-    par_sou = {par.name: results.get_variates(par.path).median
-               for par in results.optimized_model[souname].parameters.values()
-               if par.free}
-    #
-    epar_sou = {par.name: results.get_variates(par.path).std
-                for par in results.optimized_model[souname].parameters.values()
-                if par.free}
 
-    par_bk=results.get_variates(bkname).median
-    epar_bk = results.get_variates(bkname).std
-    return(par_sou,epar_sou,par_bk,epar_bk)
+    return {par_name: (results.get_variates(par.path).median, results.get_variates(par.path).std)
+              for par_name, par in results.optimized_model.free_parameters.items()}
 
 def get_fit_fluxes(results):
     """

cosipy/pipeline/src/io.py

@@ -1,7 +1,7 @@
 from cosipy import BinnedData
 from cosipy.spacecraftfile import SpacecraftFile
 
-
+from astropy.time import Time
 
 
 def load_binned_data(yaml_path,data_path):
@@ -19,23 +19,21 @@ def load_ori(ori_path):
     ori = SpacecraftFile.parse_from_file(ori_path)
     return ori
 
-def tslice_binned_data(data, tmin, tmax):
+def tslice_binned_data(data, tmin:Time, tmax:Time):
     """Slice a binned dataset in time"""
-    idx_tmin = np.where(data.axes['Time'].edges.value >= tmin)[0][0]
-    idx_tmax_all = np.where(data.axes['Time'].edges.value <= tmax)
+    idx_tmin = np.where(data.axes['Time'].edges.value >= tmin.unix)[0][0]
+    idx_tmax_all = np.where(data.axes['Time'].edges.value <= tmax.unix)
     y = len(idx_tmax_all[0]) - 1
-    idx_tmax = np.where(data.axes['Time'].edges.value <= tmax)[0][y]
+    idx_tmax = np.where(data.axes['Time'].edges.value <= tmax.unix)[0][y]
     tsliced_data = data.slice[{'Time': slice(idx_tmin, idx_tmax)}]
     return tsliced_data
 
 
 
 
-def tslice_ori(ori, tmin, tmax):
+def tslice_ori(ori, tmin:Time, tmax:Time):
     """
     Slices time for the orientation file
     """
-    ori_min = Time(tmin, format='unix')
-    ori_max = Time(tmax, format='unix')
-    tsliced_ori = ori.source_interval(ori_min, ori_max)
+    tsliced_ori = ori.source_interval(tmin, tmax)
     return tsliced_ori

cosipy/pipeline/task/task.py

@@ -1,140 +1,134 @@
+import logging
+logging.basicConfig(format='%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s',
+                    datefmt='%Y-%m-%d %H:%M:%S',
+                    level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 import argparse, textwrap
-import os
+
 from yayc import Configurator
-from threeML import Band
-from cosipy import test_data
+
 from cosipy.pipeline.src.io import *
 from cosipy.pipeline.src.preprocessing import *
 from cosipy.pipeline.src.fitting import *
 from cosipy.pipeline.src.plotting import *
-from cosipy.pipeline.src.fitting import MODEL_IDS
-from astropy import units as u
-from astropy.io.misc import yaml
 
+from pathlib import Path
+
+from astromodels.core.model_parser import ModelParser
 
 def cosi_threemlfit(argv=None):
     # Parse arguments from commandline
     apar = argparse.ArgumentParser(
         usage=textwrap.dedent(
             """
-            %(prog)s [--help] <command> [<args>] <filename> [<options>]
+            %(prog)s [--help] --config /path/to/config/file <command> [<options>]
             """),
         description=textwrap.dedent(
             """
             Fits a source at l,b and optionally in a time window tstart-tstop using the given model.
-            Data, response and orientation files should be in the same indir.
-            Outputs fit results in a fits file and a pdf plot of the fits.
+            Data, response and orientation files paths in the config file should be relative to the config file.
+            Outputs fit results in a fits file and a pdf plot of the fits. The fitted parameter
+            are also printed to stdout.
             """),
         formatter_class=argparse.RawTextHelpFormatter)
 
     apar.add_argument('--config',
-                      help="Path to .yaml file listing all the parameters.See example in test_data.")
-    apar.add_argument('--indir',
-                      help="Optional. Path to a folder containing data, orientation and response files. Default is cosipy/test_data")
-    apar.add_argument('--odir',
-                      help="Optional. Path to a folder where to save outputs. Default is cosipy/test_data")
-    apar.add_argument('--sou_data',
-                      help="Required. Name of the file containing the source.")
-    apar.add_argument('--sou_yaml',
-                      help="Required. Name of the .yaml file containing the binning information of the source file")
-    apar.add_argument('--bk_data',
-                      help="Required. Name of the background file.")
-    apar.add_argument('--bk_yaml',
-                      help="Required. Name of the .yaml file containing the binning information of the bk file")
-    apar.add_argument('--ori',
-                      help="Required. Name of the orientation file.")
-    apar.add_argument('--resp',
-                      help="Required. Name of the response file.")
-    apar.add_argument('--l',
-                      help="Required. Galactic longitude of the source.")
-    apar.add_argument('--b',
-                      help="Required. Galactic latitude of the source.")
-    apar.add_argument('--tstart',
-                      help="Optional. Starting time of time window to be fitted.")
-    apar.add_argument('--tstop',
-                      help="Optional. Ending time of time window to be fitted.")
-    apar.add_argument('--souname',
-                      help="Required. Name of the source, to be used internally to threeml.")
-    apar.add_argument('--bkname',
-                      help="Required. Name of the background, to be used internally to threeml.")
-    apar.add_argument('--model_id',
-                      help="Required. ID of the model to be fitted. Dictionary of models in pipeline/fitting")
-    apar.add_argument('--par_streams',
-                      help="Required. List of start values of parameters as astropy quantity streams. Streams can be generated with the method yaml.dump. Order should be as in threeml.")
-    apar.add_argument('--par_minvalues',
-                      help="Optional. List of min values of the parameters.")
-    apar.add_argument('--par_maxvalues',
-                      help="Optional. List of max values of the parameters.")
+                      help="Path to .yaml file listing all the parameters.See example in test_data.",
+                      required=True)
+    apar.add_argument("--config_group", default='threemlfit',
+                      help="Path withing the config file with the tutorials information")
+    apar.add_argument("--override", nargs='*',
+                      help="Override config parameters. e.g. \"section:param_int = 2\" \"section:param_string = b\"")
+    apar.add_argument("--tstart", type = float,
+                      help="Start time of the signal (unix seconds)")
+    apar.add_argument("--tstop", type=float,
+                      help="Stop time of the signal (unix seconds)")
+    apar.add_argument('-o','--output-dir',
+                      help="Output directory. Current working directory by default")
+    apar.add_argument('--log-level', default='info',
+                      help='Set the logging level (debug, info, warning, error, critical)')
+    apar.add_argument('--overwrite', action='store_true', default=False,
+                      help='Overwrite outputs. Otherwise, if a file with the same name already exists, it will throw an error.')
 
     args = apar.parse_args(argv)
 
-    # Config
-    if args.config is None:
-        config = Configurator()
-    else:
-        config = Configurator.open(args.config)
-    #
-    indir=config.get("indir")
-    odir=config.get("odir")
-    sou_data=config.get("sou_data")
-    sou_yaml=config.get("sou_yaml")
-    bk_data = config.get("bk_data")
-    bk_yaml= config.get("bk_yaml")
-    ori =config.get("ori")
-    resp=config.get("resp")
-    l=config.get("l")
-    b=config.get("b")
-    tstart=config.get("tstart")
-    tstop=config.get("tstop")
-    souname=config.get("souname")
-    bkname=config.get("bkname")
-    model_id=config.get("model_id")
-    par_streams=config.get("par_streams")
-    par_minvalues=config.get("par_minvalues")
-    par_maxvalues=config.get('par_maxvalues')
-    #
-    if indir is None:
-        print(indir)
-        indir=str(test_data.path)
-    if odir is None:
-        print(odir)
-        odir=str(test_data.path)
-    #
-    print(indir)
-    print(odir)
-    #
-    sou_data_path=os.path.join(indir,sou_data)
-    sou_yaml_path=os.path.join(indir,sou_yaml)
-    bk_data_path=os.path.join(indir,bk_data)
-    bk_yaml_path=os.path.join(indir,bk_yaml)
-    resp_path=os.path.join(indir,resp)
-    ori_path=os.path.join(indir,ori)
-    #
-    sou_binned_data = load_binned_data(sou_yaml_path, sou_data_path)
+    # Logger
+    logger.setLevel(level=args.log_level.upper())
+
+    # config file
+    full_config = Configurator.open(args.config)
+    config = Configurator(full_config[args.config_group])
+    config.config_path = full_config.config_path
+
+    # General overrides
+    if args.override is not None:
+        config.override(*args.override)
+
+    # Other specific convenience overrides
+    if args.tstart:
+        config["cuts:kwargs:tstart"] = args.tstart
+
+    if args.tstop:
+        config["cuts:kwargs:tstop"] = args.tstop
+
+    # Default output
+    odir = Path.cwd() if not args.output_dir else Path(args.output_dir)
+
+    # Parse model
+    model = ModelParser(model_dict = config['model']).get_model()
+
+    # Parse input files from config file
+    data_path = config.absolute_path(config["data:args"][0])
+    yaml_path = config.absolute_path(config["data:kwargs:input_yaml"])
+    binned_data = load_binned_data(yaml_path, data_path)
+
+    bk_data_path = config.absolute_path(config["background:args"][0])
+    bk_yaml_path = config.absolute_path(config["background:kwargs:input_yaml"])
     bk_binned_data = load_binned_data(bk_yaml_path, bk_data_path)
-    ori = load_ori(ori_path)
-            #
+
+    resp_path = config.absolute_path(config["response:args"][0])
+
+    ori = load_ori(config.absolute_path(config["sc_file"]))
+
+    # Slice time, if needed
+    tstart = config.get("cuts:kwargs:tstart")
+    tstop = config.get("cuts:kwargs:tstop")
+
     if tstart is not None and tstop is not None:
-        sou_sliced_data=tslice_binned_data(sou_binned_data, tstart, tstop)
-        sou_binned_data=sou_sliced_data
+
+        tstart = Time(tstart, format='unix')
+        tstop = Time(tstop, format='unix')
+
+        sliced_data=tslice_binned_data(binned_data, tstart, tstop)
+        binned_data=sliced_data
         bk_sliced_data = tslice_binned_data(bk_binned_data, tstart - 100, tstop + 100)
         bk_binned_data=bk_sliced_data
         ori_sliced = tslice_ori(ori, tstart, tstop)
         ori=ori_sliced
-            #
-    pars = [yaml.load(p) for p in par_streams]
-    spectrum=build_spectrum(MODEL_IDS[model_id], pars, par_minvalues, par_maxvalues)
-    results, cts_exp = get_fit_results(sou_binned_data, bk_binned_data, resp_path, ori, l, b,
-                                                   souname, bkname, spectrum)
+
+    # Calculation
+    results, cts_exp = get_fit_results(binned_data, bk_binned_data, resp_path, ori,
+                                                   "cosi_bkg", model)
+
+    # Results
     results.display()
-    results.write_to(str(odir + "/" + souname + ".fits"), overwrite=True)
-                #
-    pars_sou, epars_sou, par_bk, epar_bk = get_fit_par(results, souname, bkname)
+    results.write_to(odir/"results.fits", overwrite=args.overwrite)
+
+    print("Median and errors:")
+    fitted_par_err = get_fit_par(results)
+    for par_name,(par_median,par_err) in fitted_par_err.items():
+        print(f"{par_name} = {par_median:.2e} +/- {par_err:.2e}")
+
+    print("Total flux:")
     fl, el_fl, eh_fl = get_fit_fluxes(results)
     print("flux=%f +%f -%f" % (fl, el_fl, eh_fl))
-                #
-    figname = str(odir + "/fit_"+souname+".pdf")
-    plot_fit(sou_binned_data, cts_exp, figname)
-#
 
+    plot_filename = odir/"raw_spectrum.pdf"
+    if plot_filename.exists() and not args.overwrite:
+        raise RuntimeError(f"{plot_filename} already exists. If you mean to replace it then use --overwrite.")
+
+    plot_fit(binned_data, cts_exp, plot_filename)
 
+if __name__ == "__main__":
+    cosi_threemlfit()