Merge pull request #237 from RasmusOrsoe/pisa_weight_fitter

Pisa weight fitter

examples/weight_fitter_example.py

@@ -0,0 +1,24 @@
+from graphnet.pisa.fitting import WeightFitter
+
+outdir = "/mnt/scratch/rasmus_orsoe/weight_test"
+database_path = "/mnt/scratch/rasmus_orsoe/databases/dev_lvl3_genie_burnsample_v5/data/dev_lvl3_genie_burnsample_v5.db"
+fitter = WeightFitter(database_path=database_path)
+
+pisa_config_dict = {
+    "reco_energy": {"num_bins": 8},
+    "reco_coszen": {"num_bins": 8},
+    "pid": {"bin_edges": [0, 0.5, 1]},
+    "true_energy": {"num_bins": 200},
+    "true_coszen": {"num_bins": 200},
+    "livetime": 10
+    * 0.01,  # set to 1% of 10 years - correspond to the size of the oscNext burn sample
+}
+# by calling fitter.fit_weights we get the weights returned pr. event. if add_to_database = True, a table will be added to the database
+weights = fitter.fit_weights(
+    outdir,
+    add_to_database=True,
+    weight_name="weight_livetime10_1percent",
+    pisa_config_dict=pisa_config_dict,
+)
+
+print(weights)

src/graphnet/pisa/fitting.py

@@ -15,6 +15,7 @@
 import io
 from configupdater import ConfigUpdater
 from contextlib import contextmanager
+from graphnet.data.sqlite import run_sql_code, save_to_sql
 
 mpl.use("pdf")
 plt.rc("font", family="serif")
@@ -74,6 +75,185 @@ def config_updater(
                 updater.write(configfile)
 
 
+class WeightFitter:
+    def __init__(
+        self,
+        database_path: str,
+        truth_table: str = "truth",
+        index_column: str = "event_no",
+        statistical_fit: bool = False,
+    ):
+        self._database_path = database_path
+        self._truth_table = truth_table
+        self._index_column = index_column
+        self._statistical_fit = statistical_fit
+
+    def fit_weights(
+        self,
+        config_outdir: str,
+        weight_name: str = None,
+        pisa_config_dict: dict = None,
+        add_to_database: bool = False,
+    ):
+        """Fits flux weights to each neutrino event in self._database_path. If statistical_fit = True, only statistical effects are accounted for. If True, certain systematic effects are included, but not hypersurfaces.
+
+        Args:
+            config_outdir (str): The outdir to store the configuration in.
+            weight_name (str, optional): The name of the weight. If add_to_database = True, this will be the name of the table. Defaults to None.
+            pisa_config_dict (dict, optional): The dictionary of pisa configurations. Can be used to change assumptions regarding the fit. Defaults to None.
+            add_to_database (bool, optional): If True, a table will be added to the database called weight_name with two columns; index_column, weight_name. Defaults to False.
+
+        Returns:
+            pandas.DataFrame: A dataframe with columns index_column, weight_name .
+        """
+        # if its a standard weight
+        if pisa_config_dict is None:
+            if isinstance(weight_name, str) is False:
+                print(weight_name)
+                weight_name = "pisa_weight_graphnet_standard"
+        # if it is a custom weight without name
+        elif pisa_config_dict is not None:
+            if isinstance(weight_name, str) is False:
+                weight_name = "pisa_custom_weight"
+        pisa_config_path = self._make_config(
+            config_outdir, weight_name, pisa_config_dict
+        )
+        model = Pipeline(pisa_config_path)
+        if self._statistical_fit == "True":
+            # Only free parameters will be [aeff_scale] - corresponding to a statistical fit
+            free_params = model.params.free.names
+            for free_param in free_params:
+                if free_param not in ["aeff_scale"]:
+                    model.params[free_param].is_fixed = True
+
+        # for stage in range(len(model.stages)):
+        model.stages[-1].apply_mode = "events"
+        model.stages[-1].calc_mode = "events"
+        model.run()
+        results = pd.DataFrame()
+        for container in model.data:
+            data = pd.DataFrame(container["event_no"], columns=["event_no"])
+            data[weight_name] = container["weights"]
+            results = results.append(data)
+
+        if add_to_database:
+            self._create_table(self._database_path, weight_name, results)
+            save_to_sql(results, weight_name, self._database_path)
+        return results.sort_values("event_no").reset_index(drop=True)
+
+    def _create_table(self, database, table_name, df):
+        """Creates a table.
+        Args:
+            pipeline_database (str): path to the pipeline database
+            df (str): pandas.DataFrame of combined predictions
+        """
+        query_columns = list()
+        for column in df.columns:
+            if column == "event_no":
+                type_ = "INTEGER PRIMARY KEY NOT NULL"
+            else:
+                type_ = "FLOAT"
+            query_columns.append(f"{column} {type_}")
+        query_columns = ", ".join(query_columns)
+
+        code = (
+            "PRAGMA foreign_keys=off;\n"
+            f"CREATE TABLE {table_name} ({query_columns});\n"
+            "PRAGMA foreign_keys=on;"
+        )
+        run_sql_code(database, code)
+        return
+
+    def _make_config(
+        self,
+        config_outdir,
+        weight_name,
+        pisa_config_dict,
+    ):
+        os.makedirs(config_outdir + "/" + weight_name, exist_ok=True)
+        if pisa_config_dict is None:
+            # Run on standard settings
+            pisa_config_dict = {
+                "reco_energy": {"num_bins": 8},
+                "reco_coszen": {"num_bins": 8},
+                "pid": {"bin_edges": [0, 0.5, 1]},
+                "true_energy": {"num_bins": 200},
+                "true_coszen": {"num_bins": 200},
+                "livetime": 10
+                * 0.01,  # set to 1% of 10 years - correspond to the size of the oscNext burn sample
+            }
+
+        pisa_config_dict["pipeline"] = self._database_path
+        pisa_config_dict["post_fix"] = None
+        pipeline_cfg_path = self._create_configs(
+            pisa_config_dict, config_outdir + "/" + weight_name
+        )
+        return pipeline_cfg_path
+
+    def _create_configs(self, config_dict, path):
+        # Update binning config
+        root = os.path.realpath(
+            os.path.join(os.getcwd(), os.path.dirname(__file__))
+        )
+        if config_dict["post_fix"] is not None:
+            config_name = "config%s" % config_dict["post_fix"]
+        else:
+            # config_dict["post_fix"] = '_pred'
+            config_name = "config"
+
+        with config_updater(
+            root
+            + "/resources/configuration_templates/binning_config_template.cfg",
+            "%s/binning_%s.cfg" % (path, config_name),
+            dummy_section="binning",
+        ) as updater:
+            updater["binning"][
+                "graphnet_dynamic_binning.reco_energy"
+            ].value = (
+                "{'num_bins':%s, 'is_log':True, 'domain':[0.5,55] * units.GeV, 'tex': r'E_{\\rm reco}'}"
+                % config_dict["reco_energy"]["num_bins"]
+            )  # noqa: W605
+            updater["binning"][
+                "graphnet_dynamic_binning.reco_coszen"
+            ].value = (
+                "{'num_bins':%s, 'is_lin':True, 'domain':[-1,1], 'tex':r'\\cos{\\theta}_{\\rm reco}'}"
+                % config_dict["reco_coszen"]["num_bins"]
+            )  # noqa: W605
+            updater["binning"]["graphnet_dynamic_binning.pid"].value = (
+                "{'bin_edges': %s, 'tex':r'{\\rm PID}'}"
+                % config_dict["pid"]["bin_edges"]
+            )  # noqa: W605
+            updater["binning"]["true_allsky_fine.true_energy"].value = (
+                "{'num_bins':%s, 'is_log':True, 'domain':[1,1000] * units.GeV, 'tex': r'E_{\\rm true}'}"
+                % config_dict["true_energy"]["num_bins"]
+            )  # noqa: W605
+            updater["binning"]["true_allsky_fine.true_coszen"].value = (
+                "{'num_bins':%s, 'is_lin':True, 'domain':[-1,1], 'tex':r'\\cos\,\\theta_{Z,{\\rm true}}'}"  # noqa: W605
+                % config_dict["true_coszen"]["num_bins"]
+            )  # noqa: W605
+
+        # Update pipeline config
+        with config_updater(
+            root
+            + "/resources/configuration_templates/pipeline_config_weight_template.cfg",
+            "%s/pipeline_%s.cfg" % (path, config_name),
+        ) as updater:
+            updater["pipeline"].add_before.comment(
+                "#include %s/binning_%s.cfg as binning" % (path, config_name)
+            )
+            updater["data.sqlite_loader"]["post_fix"].value = config_dict[
+                "post_fix"
+            ]
+            updater["data.sqlite_loader"]["database"].value = config_dict[
+                "pipeline"
+            ]
+            if "livetime" in config_dict.keys():
+                updater["aeff.aeff"]["param.livetime"].value = (
+                    "%s * units.common_year" % config_dict["livetime"]
+                )
+        return "%s/pipeline_%s.cfg" % (path, config_name)
+
+
 class ContourFitter:
     def __init__(
         self,

src/graphnet/pisa/resources/configuration_templates/pipeline_config_weight_template.cfg

@@ -0,0 +1,101 @@
+#include settings/osc/nufitv20.cfg as osc
+#include settings/osc/earth.cfg as earth
+[pipeline]
+order = data.sqlite_loader, flux.honda_ip, flux.barr_simple, osc.prob3, aeff.aeff
+param_selections = nh
+name = neutrinos
+output_binning = graphnet_dynamic_binning
+output_key = weights, errors
+
+[data.sqlite_loader]
+calc_mode = events
+apply_mode = events
+output_names = nue_cc, numu_cc, nutau_cc, nue_nc, numu_nc, nutau_nc, nuebar_cc, numubar_cc, nutaubar_cc, nuebar_nc, numubar_nc, nutaubar_nc
+post_fix = _pred
+database = /mnt/scratch/rasmus_orsoe/databases/oscillations/dev_lvl7_robustness_muon_neutrino_0000/pipelines/pipeline_oscillation_final/pipeline_oscillation_final.db
+
+[flux.honda_ip]
+calc_mode = true_allsky_fine
+apply_mode = events
+param.flux_table = flux/honda-2015-spl-solmin-aa.d
+
+[flux.barr_simple]
+calc_mode = true_allsky_fine
+apply_mode = events
+param.nu_nubar_ratio = 1.0 +/- 0.1
+param.nu_nubar_ratio.fixed = True
+param.nu_nubar_ratio.range = nominal + [-3., +3.] * sigma
+param.nue_numu_ratio = 1.0 +/- 0.05
+param.nue_numu_ratio.fixed = False
+param.nue_numu_ratio.range = nominal + [-0.5, +0.5]
+param.Barr_uphor_ratio = 0.0 +/- 1.0
+param.Barr_uphor_ratio.fixed = False
+param.Barr_uphor_ratio.range = nominal + [-3.0, +3.0]
+param.Barr_nu_nubar_ratio = 0.0 +/- 1.0
+param.Barr_nu_nubar_ratio.fixed = False
+param.Barr_nu_nubar_ratio.range = nominal + [-3.0, +3.0]
+param.delta_index = 0.0 +/- 0.1
+param.delta_index.fixed = False
+param.delta_index.range = nominal + [-5, +5] * sigma
+
+[osc.prob3]
+calc_mode = true_allsky_fine
+apply_mode = events
+param.earth_model = osc/PREM_12layer.dat
+param.YeI = ${earth:YeI}
+param.YeM = ${earth:YeM}
+param.YeO = ${earth:YeO}
+param.detector_depth = ${earth:detector_depth}
+param.prop_height = ${earth:prop_height}
+param.theta12 = ${osc:theta12}
+param.theta12.fixed = True
+param.nh.theta13 = ${osc:theta13_nh}
+param.nh.theta13.fixed = False
+param.nh.theta13.range = ${osc:theta13_nh.range}
+param.ih.theta13 = ${osc:theta13_ih}
+param.ih.theta13.fixed = False
+param.ih.theta13.range = ${osc:theta13_ih.range}
+param.nh.theta23 = ${osc:theta23_nh}
+param.nh.theta23.fixed = False
+param.nh.theta23.range = ${osc:theta23_nh.range}
+param.nh.theta23.prior = uniform
+param.ih.theta23 = ${osc:theta23_ih}
+param.ih.theta23.fixed = False
+param.ih.theta23.range = ${osc:theta23_ih.range}
+param.ih.theta23.prior = uniform
+param.nh.deltacp = 0.0 * units.degree
+param.nh.deltacp.fixed = True
+param.nh.deltacp.range = ${osc:deltacp_nh.range}
+param.nh.deltacp.prior = uniform
+param.ih.deltacp = 0.0 * units.degree
+param.ih.deltacp.fixed = True
+param.deltam21 = ${osc:deltam21}
+param.deltam21.fixed = True
+param.nh.deltam31 = ${osc:deltam31_nh}
+param.nh.deltam31.fixed = False
+param.nh.deltam31.prior = uniform
+param.nh.deltam31.range = [0.001, +0.007] * units.eV**2
+param.ih.deltam31 = ${osc:deltam31_ih}
+param.ih.deltam31.fixed = False
+param.ih.deltam31.prior = uniform
+param.ih.deltam31.range = [-0.007, -0.001] * units.eV**2
+
+[aeff.aeff]
+calc_mode = events
+apply_mode = events
+param.livetime = 10 * units.common_year
+param.aeff_scale = 1.0
+param.aeff_scale.fixed = False
+param.aeff_scale.prior = uniform
+param.aeff_scale.range = [0.,3.] * units.dimensionless
+param.nutau_cc_norm = 1.0
+param.nutau_cc_norm.fixed = True
+param.nutau_cc_norm.range = [0.2, 2.0] * units.dimensionless
+param.nutau_cc_norm.prior = uniform
+param.nutau_norm = 1.0
+param.nutau_norm.fixed = False
+param.nutau_norm.range = [-1.0, 8.5] * units.dimensionless
+param.nutau_norm.prior = uniform
+param.nu_nc_norm = 1.0 +/- 0.2
+param.nu_nc_norm.fixed = False
+param.nu_nc_norm.range = nominal + [-.5,+.5]