Exemplo n.º 1
0
    def setup(self):

        if self.output_irf_file.absolute().exists():
            if self.overwrite:
                self.log.warning(f"Overwriting {self.output_irf_file}")
                self.output_irf_file.unlink()
            else:
                raise ToolConfigurationError(
                    f"Output file {self.output_irf_file} already exists,"
                    " use --overwrite to overwrite"
                )

        filename = self.output_irf_file.name
        if not (filename.endswith('.fits') or filename.endswith('.fits.gz')):
            raise ValueError(
                f"{filename} is not a correct compressed FITS file name"
                "(use .fits or .fits.gz)."
                )

        if self.input_proton_dl2 and self.input_electron_dl2 is not Undefined:
            self.only_gamma_irf = False
        else:
            self.only_gamma_irf = True

        self.event_sel = EventSelector(parent=self)
        self.cuts = DL3Cuts(parent=self)
        self.data_bin = DataBinning(parent=self)

        self.mc_particle = {
            "gamma": {
                "file": self.input_gamma_dl2,
                "target_spectrum": CRAB_MAGIC_JHEAP2015,
            },
        }
        Provenance().add_input_file(self.input_gamma_dl2)

        self.t_obs = self.irf_obs_time * u.hour

        # Read and update MC information
        if not self.only_gamma_irf:
            self.mc_particle["proton"] = {
                "file":  self.input_proton_dl2,
                "target_spectrum": IRFDOC_PROTON_SPECTRUM,
            }

            self.mc_particle["electron"] = {
                "file": self.input_electron_dl2,
                "target_spectrum": IRFDOC_ELECTRON_SPECTRUM,
            }

            Provenance().add_input_file(self.input_proton_dl2)
            Provenance().add_input_file(self.input_electron_dl2)

        self.provenance_log = self.output_irf_file.parent / (
                self.name + ".provenance.log"
        )
    def setup(self):

        self.filename_dl3 = dl2_to_dl3_filename(self.input_dl2,
                                                compress=self.gzip)
        self.provenance_log = self.output_dl3_path / (self.name +
                                                      ".provenance.log")

        Provenance().add_input_file(self.input_dl2)

        self.event_sel = EventSelector(parent=self)
        self.cuts = DL3Cuts(parent=self)

        self.output_file = self.output_dl3_path.absolute() / self.filename_dl3
        if self.output_file.exists():
            if self.overwrite:
                self.log.warning(f"Overwriting {self.output_file}")
                self.output_file.unlink()
            else:
                raise ToolConfigurationError(
                    f"Output file {self.output_file} already exists,"
                    " use --overwrite to overwrite")
        if not (self.source_ra or self.source_dec):
            self.source_pos = SkyCoord.from_name(self.source_name)
        elif bool(self.source_ra) != bool(self.source_dec):
            raise ToolConfigurationError(
                "Either provide both RA and DEC values for the source or none")
        else:
            self.source_pos = SkyCoord(ra=self.source_ra, dec=self.source_dec)

        self.log.debug(f"Output DL3 file: {self.output_file}")

        try:
            with fits.open(self.input_irf) as hdul:
                self.use_energy_dependent_gh_cuts = (
                    "GH_CUT" not in hdul["EFFECTIVE AREA"].header)
        except:
            raise ToolConfigurationError(
                f"{self.input_irf} does not have EFFECTIVE AREA HDU, "
                " to check for global cut information in the Header value")

        if self.source_dep:
            with fits.open(self.input_irf) as hdul:
                self.use_energy_dependent_alpha_cuts = (
                    "AL_CUT" not in hdul["EFFECTIVE AREA"].header)
Exemplo n.º 3
0
def test_event_selection():
    evt_fil = EventSelector()

    data_t = QTable({
        "a": u.Quantity([1, 2, 3], unit=u.kg),
        "b": u.Quantity([np.nan, 2.2, 3.2], unit=u.m),
        "c": u.Quantity([1, 3, np.inf], unit=u.s),
    })

    evt_fil.filters = dict(a=[0, 2.5], b=[0, 3], c=[0, 4])
    evt_fil.finite_params = ["b"]

    data_t = evt_fil.filter_cut(data_t)
    data_t_df = evt_fil.filter_cut(data_t.to_pandas())

    np.testing.assert_array_equal(
        data_t_df, pd.DataFrame({
            "a": [2],
            "b": [2.2],
            "c": [3]
        }))

    np.testing.assert_array_equal(
        data_t,
        QTable({
            "a": u.Quantity([2], unit=u.kg),
            "b": u.Quantity([2.2], unit=u.m),
            "c": u.Quantity([3], unit=u.s),
        }),
    )
class IRFFITSWriter(Tool):
    name = "IRFFITSWriter"
    description = __doc__
    example = """
    To generate IRFs from MC gamma only, using default cuts/binning:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)
        --overwrite

    Or to generate all 4 IRFs, using default cuts/binning:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -p /path/to/DL2_MC_proton_file.h5
        -e /path/to/DL2_MC_electron_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)

    Or use a config file for cuts and binning information:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)
        --config /path/to/config.json

    Or pass the selection cuts from command-line:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)
        --global-gh-cut 0.9
        --global-theta-cut 0.2
        --irf-obs-time 50

    Or use energy-dependent cuts based on a gamma efficiency:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)
        --energy-dependent-gh
        --energy-dependent-theta
        --gh-efficiency 0.95
        --theta-containment 0.68

    Or generate source-dependent IRFs
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -o /path/to/irf.fits.gz
        --point-like
        --global-gh-cut 0.9
        --global-alpha-cut 10
        --source-dep

    """

    input_gamma_dl2 = traits.Path(help="Input MC gamma DL2 file",
                                  allow_none=True,
                                  exists=True,
                                  directory_ok=False,
                                  file_ok=True).tag(config=True)

    input_proton_dl2 = traits.Path(help="Input MC proton DL2 file",
                                   allow_none=True,
                                   exists=True,
                                   directory_ok=False,
                                   file_ok=True).tag(config=True)

    input_electron_dl2 = traits.Path(help="Input MC electron DL2 file",
                                     allow_none=True,
                                     exists=True,
                                     directory_ok=False,
                                     file_ok=True).tag(config=True)

    output_irf_file = traits.Path(
        help="IRF output file",
        allow_none=True,
        directory_ok=False,
        file_ok=True,
        default_value="./irf.fits.gz",
    ).tag(config=True)

    irf_obs_time = traits.Float(
        help="Observation time for IRF in hours",
        default_value=50,
    ).tag(config=True)

    point_like = traits.Bool(
        help="True for point_like IRF, False for Full Enclosure",
        default_value=False,
    ).tag(config=True)

    energy_dependent_gh = traits.Bool(
        help="True for applying energy-dependent gammaness cuts",
        default_value=False,
    ).tag(config=True)

    energy_dependent_theta = traits.Bool(
        help="True for applying energy-dependent theta cuts",
        default_value=False,
    ).tag(config=True)

    energy_dependent_alpha = traits.Bool(
        help="True for applying energy-dependent alpha cuts",
        default_value=False,
    ).tag(config=True)

    overwrite = traits.Bool(
        help="If True, overwrites existing output file without asking",
        default_value=False,
    ).tag(config=True)

    source_dep = traits.Bool(
        help="True for source-dependent analysis",
        default_value=False,
    ).tag(config=True)

    classes = [EventSelector, DL3Cuts, DataBinning]

    aliases = {
        ("g", "input-gamma-dl2"): "IRFFITSWriter.input_gamma_dl2",
        ("p", "input-proton-dl2"): "IRFFITSWriter.input_proton_dl2",
        ("e", "input-electron-dl2"): "IRFFITSWriter.input_electron_dl2",
        ("o", "output-irf-file"): "IRFFITSWriter.output_irf_file",
        "irf-obs-time": "IRFFITSWriter.irf_obs_time",
        "global-gh-cut": "DL3Cuts.global_gh_cut",
        "gh-efficiency": "DL3Cuts.gh_efficiency",
        "theta-containment": "DL3Cuts.theta_containment",
        "global-theta-cut": "DL3Cuts.global_theta_cut",
        "alpha-containment": "DL3Cuts.alpha_containment",
        "global-alpha-cut": "DL3Cuts.global_alpha_cut",
        "allowed-tels": "DL3Cuts.allowed_tels",
        "overwrite": "IRFFITSWriter.overwrite",
    }

    flags = {
        "point-like": (
            {
                "IRFFITSWriter": {
                    "point_like": True
                }
            },
            "Point like IRFs will be produced, otherwise Full Enclosure",
        ),
        "overwrite": (
            {
                "IRFFITSWriter": {
                    "overwrite": True
                }
            },
            "overwrites output file",
        ),
        "source-dep": (
            {
                "IRFFITSWriter": {
                    "source_dep": True
                }
            },
            "Source-dependent analysis will be performed",
        ),
        "energy-dependent-gh": (
            {
                "IRFFITSWriter": {
                    "energy_dependent_gh": True
                }
            },
            "Uses energy-dependent cuts for gammaness",
        ),
        "energy-dependent-theta": (
            {
                "IRFFITSWriter": {
                    "energy_dependent_theta": True
                }
            },
            "Uses energy-dependent cuts for theta",
        ),
        "energy-dependent-alpha": (
            {
                "IRFFITSWriter": {
                    "energy_dependent_alpha": True
                }
            },
            "Uses energy-dependent cuts for alpha",
        ),
    }

    def setup(self):

        if self.output_irf_file.absolute().exists():
            if self.overwrite:
                self.log.warning(f"Overwriting {self.output_irf_file}")
                self.output_irf_file.unlink()
            else:
                raise ToolConfigurationError(
                    f"Output file {self.output_irf_file} already exists,"
                    " use --overwrite to overwrite")

        filename = self.output_irf_file.name
        if not (filename.endswith('.fits') or filename.endswith('.fits.gz')):
            raise ValueError(
                f"{filename} is not a correct compressed FITS file name"
                "(use .fits or .fits.gz).")

        if self.input_proton_dl2 and self.input_electron_dl2 is not Undefined:
            self.only_gamma_irf = False
        else:
            self.only_gamma_irf = True

        self.event_sel = EventSelector(parent=self)
        self.cuts = DL3Cuts(parent=self)
        self.data_bin = DataBinning(parent=self)

        self.mc_particle = {
            "gamma": {
                "file": self.input_gamma_dl2,
                "target_spectrum": CRAB_MAGIC_JHEAP2015,
            },
        }
        Provenance().add_input_file(self.input_gamma_dl2)

        self.t_obs = self.irf_obs_time * u.hour

        # Read and update MC information
        if not self.only_gamma_irf:
            self.mc_particle["proton"] = {
                "file": self.input_proton_dl2,
                "target_spectrum": IRFDOC_PROTON_SPECTRUM,
            }

            self.mc_particle["electron"] = {
                "file": self.input_electron_dl2,
                "target_spectrum": IRFDOC_ELECTRON_SPECTRUM,
            }

            Provenance().add_input_file(self.input_proton_dl2)
            Provenance().add_input_file(self.input_electron_dl2)

        self.provenance_log = self.output_irf_file.parent / (self.name +
                                                             ".provenance.log")

    def start(self):

        for particle_type, p in self.mc_particle.items():
            self.log.info(f"Simulated {particle_type.title()} Events:")
            p["events"], p["simulation_info"] = read_mc_dl2_to_QTable(
                p["file"])

            p["mc_type"] = check_mc_type(p["file"])

            self.log.debug(
                f"Simulated {p['mc_type']} {particle_type.title()} Events:")

            # Calculating event weights for Background IRF
            if particle_type != "gamma":
                p["simulated_spectrum"] = PowerLaw.from_simulation(
                    p["simulation_info"], self.t_obs)

                p["events"]["weight"] = calculate_event_weights(
                    p["events"]["true_energy"],
                    p["target_spectrum"],
                    p["simulated_spectrum"],
                )

            if not self.source_dep:
                for prefix in ("true", "reco"):
                    k = f"{prefix}_source_fov_offset"
                    p["events"][k] = calculate_source_fov_offset(p["events"],
                                                                 prefix=prefix)

                # calculate theta / distance between reco and assumed source position
                p["events"]["theta"] = calculate_theta(
                    p["events"],
                    assumed_source_az=p["events"]["true_az"],
                    assumed_source_alt=p["events"]["true_alt"],
                )

            else:
                # Alpha cut is applied for source-dependent analysis.
                # To adapt source-dependent analysis to pyirf codes,
                # true position is set as reco position for survived events
                # after alpha cut
                p["events"][
                    "true_source_fov_offset"] = calculate_source_fov_offset(
                        p["events"], prefix="true")
                p["events"]["reco_source_fov_offset"] = p["events"][
                    "true_source_fov_offset"]

        self.log.debug(p["simulation_info"])
        gammas = self.mc_particle["gamma"]["events"]

        # Binning of parameters used in IRFs
        true_energy_bins = self.data_bin.true_energy_bins()
        reco_energy_bins = self.data_bin.reco_energy_bins()
        migration_bins = self.data_bin.energy_migration_bins()
        source_offset_bins = self.data_bin.source_offset_bins()

        gammas = self.event_sel.filter_cut(gammas)
        gammas = self.cuts.allowed_tels_filter(gammas)

        if self.energy_dependent_gh:
            self.gh_cuts_gamma = self.cuts.energy_dependent_gh_cuts(
                gammas, reco_energy_bins)
            gammas = self.cuts.apply_energy_dependent_gh_cuts(
                gammas, self.gh_cuts_gamma)
            self.log.info(
                f"Using gamma efficiency of {self.cuts.gh_efficiency}")
        else:
            gammas = self.cuts.apply_global_gh_cut(gammas)
            self.log.info("Using a global gammaness cut of "
                          f"{self.cuts.global_gh_cut}")

        if self.point_like:
            if not self.source_dep:
                if self.energy_dependent_theta:
                    self.theta_cuts = self.cuts.energy_dependent_theta_cuts(
                        gammas,
                        reco_energy_bins,
                    )
                    gammas = self.cuts.apply_energy_dependent_theta_cuts(
                        gammas, self.theta_cuts)
                    self.log.info("Using a containment region for theta of "
                                  f"{self.cuts.theta_containment}")
                else:
                    gammas = self.cuts.apply_global_theta_cut(gammas)
                    self.log.info(
                        "Using a global Theta cut of "
                        f"{self.cuts.global_theta_cut} for point-like IRF")
            else:
                if self.energy_dependent_alpha:
                    self.alpha_cuts = self.cuts.energy_dependent_alpha_cuts(
                        gammas,
                        reco_energy_bins,
                    )
                    gammas = self.cuts.apply_energy_dependent_alpha_cuts(
                        gammas, self.alpha_cuts)
                    self.log.info("Using a containment region for alpha of "
                                  f"{self.cuts.alpha_containment} %")
                else:
                    gammas = self.cuts.apply_global_alpha_cut(gammas)
                    self.log.info(
                        'Using a global Alpha cut of '
                        f'{self.cuts.global_alpha_cut} for point like IRF')

        if self.mc_particle["gamma"]["mc_type"] in [
                "point_like", "ring_wobble"
        ]:
            mean_fov_offset = round(
                gammas["true_source_fov_offset"].mean().to_value(), 1)
            fov_offset_bins = [mean_fov_offset - 0.1, mean_fov_offset + 0.1
                               ] * u.deg
            self.log.info('Single offset for point like gamma MC')
        else:
            fov_offset_bins = self.data_bin.fov_offset_bins()
            self.log.info('Multiple offset for diffuse gamma MC')

            if self.energy_dependent_theta:
                fov_offset_bins = [
                    round(gammas["true_source_fov_offset"].min().to_value(),
                          1),
                    round(gammas["true_source_fov_offset"].max().to_value(), 1)
                ] * u.deg
                self.log.info("For RAD MAX, the full FoV is used")

        if not self.only_gamma_irf:
            background = table.vstack([
                self.mc_particle["proton"]["events"],
                self.mc_particle["electron"]["events"]
            ])

            if self.energy_dependent_gh:
                background = self.cuts.apply_energy_dependent_gh_cuts(
                    background, self.gh_cuts_gamma)
            else:
                background = self.cuts.apply_global_gh_cut(background)

            background = self.event_sel.filter_cut(background)
            background = self.cuts.allowed_tels_filter(background)

            background_offset_bins = self.data_bin.bkg_fov_offset_bins()

        # For a global gh/theta cut, only a header value is added.
        # For energy-dependent cuts, along with GADF specified RAD_MAX HDU,
        # a new HDU is created, GH_CUTS which is based on RAD_MAX table

        # NOTE: The GH_CUTS HDU is just for provenance and is not supported
        # by GADF or used by any Science Tools
        extra_headers = {
            "TELESCOP": "CTA-N",
            "INSTRUME": "LST-" + " ".join(map(str, self.cuts.allowed_tels)),
            "FOVALIGN": "RADEC",
        }
        if self.point_like:
            self.log.info("Generating point_like IRF HDUs")
        else:
            self.log.info("Generating Full-Enclosure IRF HDUs")

        # Updating the HDU headers with the gammaness and theta cuts/efficiency
        if not self.energy_dependent_gh:
            extra_headers["GH_CUT"] = self.cuts.global_gh_cut

        else:
            extra_headers["GH_EFF"] = (self.cuts.gh_efficiency,
                                       "gamma/hadron efficiency")

        if self.point_like:
            if not self.source_dep:
                if self.energy_dependent_theta:
                    extra_headers["TH_CONT"] = (
                        self.cuts.theta_containment,
                        "Theta containment region in percentage")
                else:
                    extra_headers["RAD_MAX"] = (self.cuts.global_theta_cut,
                                                'deg')
            else:
                if self.energy_dependent_alpha:
                    extra_headers["AL_CONT"] = (
                        self.cuts.alpha_containment,
                        "Alpha containment region in percentage")
                else:
                    extra_headers["AL_CUT"] = (self.cuts.global_alpha_cut,
                                               'deg')

        # Write HDUs
        self.hdus = [
            fits.PrimaryHDU(),
        ]

        with np.errstate(invalid="ignore", divide="ignore"):
            if self.mc_particle["gamma"]["mc_type"] in [
                    "point_like", "ring_wobble"
            ]:
                self.effective_area = effective_area_per_energy(
                    gammas,
                    self.mc_particle["gamma"]["simulation_info"],
                    true_energy_bins,
                )
                self.hdus.append(
                    create_aeff2d_hdu(
                        # add one dimension for single FOV offset
                        self.effective_area[..., np.newaxis],
                        true_energy_bins,
                        fov_offset_bins,
                        point_like=self.point_like,
                        extname="EFFECTIVE AREA",
                        **extra_headers,
                    ))
            else:
                self.effective_area = effective_area_per_energy_and_fov(
                    gammas,
                    self.mc_particle["gamma"]["simulation_info"],
                    true_energy_bins,
                    fov_offset_bins,
                )
                self.hdus.append(
                    create_aeff2d_hdu(
                        self.effective_area,
                        true_energy_bins,
                        fov_offset_bins,
                        point_like=self.point_like,
                        extname="EFFECTIVE AREA",
                        **extra_headers,
                    ))

        self.log.info("Effective Area HDU created")
        self.edisp = energy_dispersion(
            gammas,
            true_energy_bins,
            fov_offset_bins,
            migration_bins,
        )
        self.hdus.append(
            create_energy_dispersion_hdu(
                self.edisp,
                true_energy_bins,
                migration_bins,
                fov_offset_bins,
                point_like=self.point_like,
                extname="ENERGY DISPERSION",
                **extra_headers,
            ))
        self.log.info("Energy Dispersion HDU created")

        if not self.only_gamma_irf:
            self.background = background_2d(
                background,
                reco_energy_bins=reco_energy_bins,
                fov_offset_bins=background_offset_bins,
                t_obs=self.t_obs,
            )
            self.hdus.append(
                create_background_2d_hdu(
                    self.background.T,
                    reco_energy_bins,
                    background_offset_bins,
                    extname="BACKGROUND",
                    **extra_headers,
                ))
            self.log.info("Background HDU created")

        if not self.point_like:
            self.psf = psf_table(
                gammas,
                true_energy_bins,
                fov_offset_bins=fov_offset_bins,
                source_offset_bins=source_offset_bins,
            )
            self.hdus.append(
                create_psf_table_hdu(
                    self.psf,
                    true_energy_bins,
                    source_offset_bins,
                    fov_offset_bins,
                    extname="PSF",
                    **extra_headers,
                ))
            self.log.info("PSF HDU created")

        if self.energy_dependent_gh:
            # Create a separate temporary header
            gh_header = fits.Header()
            gh_header["CREATOR"] = f"lstchain v{__version__}"
            gh_header["DATE"] = Time.now().utc.iso

            for k, v in extra_headers.items():
                gh_header[k] = v

            self.hdus.append(
                fits.BinTableHDU(self.gh_cuts_gamma,
                                 header=gh_header,
                                 name="GH_CUTS"))
            self.log.info("GH CUTS HDU added")

        if self.energy_dependent_theta and self.point_like:
            if not self.source_dep:
                self.hdus.append(
                    create_rad_max_hdu(self.theta_cuts["cut"][:, np.newaxis],
                                       reco_energy_bins, fov_offset_bins,
                                       **extra_headers))
                self.log.info("RAD MAX HDU added")

        if self.energy_dependent_alpha and self.source_dep:
            # Create a separate temporary header
            alpha_header = fits.Header()
            alpha_header["CREATOR"] = f"lstchain v{__version__}"
            alpha_header["DATE"] = Time.now().utc.iso

            for k, v in extra_headers.items():
                alpha_header[k] = v

            self.hdus.append(
                fits.BinTableHDU(self.alpha_cuts,
                                 header=gh_header,
                                 name="AL_CUTS"))
            self.log.info("ALPHA CUTS HDU added")

    def finish(self):

        fits.HDUList(self.hdus).writeto(self.output_irf_file,
                                        overwrite=self.overwrite)
        Provenance().add_output_file(self.output_irf_file)
class DataReductionFITSWriter(Tool):
    name = "DataReductionFITSWriter"
    description = __doc__
    example = """
    To generate DL3 file from an observed data DL2 file, using default cuts:
    > lstchain_create_dl3_file
        -d /path/to/DL2_data_file.h5
        -o /path/to/DL3/file/
        --input-irf /path/to/irf.fits.gz
        --source-name Crab
        --source-ra 83.633deg
        --source-dec 22.01deg

    Or use a config file for the cuts:
    > lstchain_create_dl3_file
        -d /path/to/DL2_data_file.h5
        -o /path/to/DL3/file/
        --input-irf /path/to/irf.fits.gz
        --source-name Crab
        --source-ra 83.633deg
        --source-dec 22.01deg
        --overwrite
        --config /path/to/config.json

    Or pass the selection cuts from command-line:
    > lstchain_create_dl3_file
        -d /path/to/DL2_data_file.h5
        -o /path/to/DL3/file/
        --input-irf /path/to/irf.fits.gz
        --source-name Crab
        --source-ra 83.633deg
        --source-dec 22.01deg
        --global-gh-cut 0.9
        --overwrite

    Or generate source-dependent DL3 files
    > lstchain_create_dl3_file
        -d /path/to/DL2_data_file.h5
        -o /path/to/DL3/file/
        --input-irf /path/to/irf.fits.gz
        --source-name Crab
        --source-dep
        --overwrite
    """

    input_dl2 = traits.Path(help="Input data DL2 file",
                            exists=True,
                            directory_ok=False,
                            file_ok=True).tag(config=True)

    output_dl3_path = traits.Path(help="DL3 output filedir",
                                  directory_ok=True,
                                  file_ok=False).tag(config=True)

    input_irf = traits.Path(
        help="Compressed FITS file of IRFs",
        exists=True,
        directory_ok=False,
        file_ok=True,
    ).tag(config=True)

    source_name = traits.Unicode(help="Name of Source").tag(config=True)

    source_ra = traits.Unicode(help="RA position of the source").tag(
        config=True)

    source_dec = traits.Unicode(help="DEC position of the source").tag(
        config=True)

    overwrite = traits.Bool(
        help="If True, overwrites existing output file without asking",
        default_value=False,
    ).tag(config=True)

    source_dep = traits.Bool(
        help="If True, source-dependent analysis will be performed.",
        default_value=False,
    ).tag(config=True)

    classes = [EventSelector, DL3Cuts]

    aliases = {
        ("d", "input-dl2"): "DataReductionFITSWriter.input_dl2",
        ("o", "output-dl3-path"): "DataReductionFITSWriter.output_dl3_path",
        "input-irf": "DataReductionFITSWriter.input_irf",
        "global-gh-cut": "DL3Cuts.global_gh_cut",
        "source-name": "DataReductionFITSWriter.source_name",
        "source-ra": "DataReductionFITSWriter.source_ra",
        "source-dec": "DataReductionFITSWriter.source_dec",
    }

    flags = {
        "overwrite": (
            {
                "DataReductionFITSWriter": {
                    "overwrite": True
                }
            },
            "overwrite output file if True",
        ),
        "source-dep": (
            {
                "DataReductionFITSWriter": {
                    "source_dep": True
                }
            },
            "source-dependent analysis if True",
        ),
    }

    def setup(self):

        self.filename_dl3 = dl2_to_dl3_filename(self.input_dl2)
        self.provenance_log = self.output_dl3_path / (self.name +
                                                      ".provenance.log")

        Provenance().add_input_file(self.input_dl2)

        self.event_sel = EventSelector(parent=self)
        self.cuts = DL3Cuts(parent=self)

        self.output_file = self.output_dl3_path.absolute() / self.filename_dl3
        if self.output_file.exists():
            if self.overwrite:
                self.log.warning(f"Overwriting {self.output_file}")
                self.output_file.unlink()
            else:
                raise ToolConfigurationError(
                    f"Output file {self.output_file} already exists,"
                    " use --overwrite to overwrite")
        if not (self.source_ra or self.source_dec):
            self.source_pos = SkyCoord.from_name(self.source_name)
        elif bool(self.source_ra) != bool(self.source_dec):
            raise ToolConfigurationError(
                "Either provide both RA and DEC values for the source or none")
        else:
            self.source_pos = SkyCoord(ra=self.source_ra, dec=self.source_dec)

        self.log.debug(f"Output DL3 file: {self.output_file}")

        try:
            with fits.open(self.input_irf) as hdul:
                self.use_energy_dependent_cuts = (
                    "GH_CUT" not in hdul["EFFECTIVE AREA"].header)
        except:
            raise ToolConfigurationError(
                f"{self.input_irf} does not have EFFECTIVE AREA HDU, "
                " to check for global cut information in the Header value")

    def apply_srcindep_gh_cut(self):
        ''' apply gammaness cut '''
        self.data = self.event_sel.filter_cut(self.data)

        if self.use_energy_dependent_cuts:
            self.energy_dependent_gh_cuts = QTable.read(self.input_irf,
                                                        hdu="GH_CUTS")

            self.data = self.cuts.apply_energy_dependent_gh_cuts(
                self.data, self.energy_dependent_gh_cuts)
            self.log.info("Using gamma efficiency of "
                          f"{self.energy_dependent_gh_cuts.meta['GH_EFF']}")
        else:
            with fits.open(self.input_irf) as hdul:
                self.cuts.global_gh_cut = hdul[1].header["GH_CUT"]
            self.data = self.cuts.apply_global_gh_cut(self.data)
            self.log.info(f"Using global G/H cut of {self.cuts.global_gh_cut}")

    def apply_srcdep_gh_alpha_cut(self):
        ''' apply gammaness and alpha cut for source-dependent analysis '''
        srcdep_assumed_positions = get_srcdep_assumed_positions(self.input_dl2)

        for i, srcdep_pos in enumerate(srcdep_assumed_positions):
            data_temp = read_data_dl2_to_QTable(self.input_dl2,
                                                srcdep_pos=srcdep_pos)

            data_temp = self.event_sel.filter_cut(data_temp)

            if self.use_energy_dependent_cuts:
                self.energy_dependent_gh_cuts = QTable.read(self.input_irf,
                                                            hdu="GH_CUTS")

                data_temp = self.cuts.apply_energy_dependent_gh_cuts(
                    data_temp, self.energy_dependent_gh_cuts)
            else:
                with fits.open(self.input_irf) as hdul:
                    self.cuts.global_gh_cut = hdul[1].header["GH_CUT"]
                data_temp = self.cuts.apply_global_gh_cut(data_temp)

            with fits.open(self.input_irf) as hdul:
                self.cuts.global_alpha_cut = hdul[1].header["AL_CUT"]
            data_temp = self.cuts.apply_global_alpha_cut(data_temp)

            # set expected source positions as reco positions
            set_expected_pos_to_reco_altaz(data_temp)

            if i == 0:
                self.data = data_temp
            else:
                self.data = vstack([self.data, data_temp])

    def start(self):

        if not self.source_dep:
            self.data = read_data_dl2_to_QTable(self.input_dl2)
        else:
            self.data = read_data_dl2_to_QTable(self.input_dl2, 'on')
        self.effective_time, self.elapsed_time = get_effective_time(self.data)
        self.run_number = run_info_from_filename(self.input_dl2)[1]

        if not self.source_dep:
            self.apply_srcindep_gh_cut()
        else:
            self.apply_srcdep_gh_alpha_cut()

        self.data = add_icrs_position_params(self.data, self.source_pos)

        self.log.info("Generating event list")
        self.events, self.gti, self.pointing = create_event_list(
            data=self.data,
            run_number=self.run_number,
            source_name=self.source_name,
            source_pos=self.source_pos,
            effective_time=self.effective_time.value,
            elapsed_time=self.elapsed_time.value,
        )

        self.hdulist = fits.HDUList(
            [fits.PrimaryHDU(), self.events, self.gti, self.pointing])

        irf = fits.open(self.input_irf)
        self.log.info("Adding IRF HDUs")

        for irf_hdu in irf[1:]:
            self.hdulist.append(irf_hdu)

    def finish(self):
        self.hdulist.writeto(self.output_file, overwrite=self.overwrite)

        Provenance().add_output_file(self.output_file)
class IRFFITSWriter(Tool):
    name = "IRFFITSWriter"
    description = __doc__
    example = """
    To generate IRFs from MC gamma only, using default cuts/binning:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)
        --overwrite

    Or to generate all 4 IRFs, using default cuts/binning:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -p /path/to/DL2_MC_proton_file.h5
        -e /path/to/DL2_MC_electron_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)

    Or use a config file for cuts and binning information:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)
        --config /path/to/config.json

    Or pass the selection cuts from command-line:
    > lstchain_create_irf_files
        -g /path/to/DL2_MC_gamma_file.h5
        -o /path/to/irf.fits.gz
        --point-like (Only for point_like IRFs)
        --fixed-gh-cut 0.9
        --fixed-theta-cut 0.2
        --irf-obs-time 50
    """

    input_gamma_dl2 = traits.Path(
        help="Input MC gamma DL2 file",
        exists=True,
        directory_ok=False,
        file_ok=True
    ).tag(config=True)

    input_proton_dl2 = traits.Path(
        help="Input MC proton DL2 file",
        exists=True,
        directory_ok=False,
        file_ok=True
    ).tag(config=True)

    input_electron_dl2 = traits.Path(
        help="Input MC electron DL2 file",
        exists=True,
        directory_ok=False,
        file_ok=True
    ).tag(config=True)

    output_irf_file = traits.Path(
        help="IRF output file",
        directory_ok=False,
        file_ok=True,
        default_value="./irf.fits.gz",
    ).tag(config=True)

    irf_obs_time = traits.Float(
        help="Observation time for IRF in hours",
        default_value=50,
    ).tag(config=True)

    point_like = traits.Bool(
        help="True for point_like IRF, False for Full Enclosure",
        default_value=False,
    ).tag(config=True)

    overwrite = traits.Bool(
        help="If True, overwrites existing output file without asking",
        default_value=False,
    ).tag(config=True)

    classes = [EventSelector, DL3FixedCuts, DataBinning]

    aliases = {
        ("g", "input-gamma-dl2"): "IRFFITSWriter.input_gamma_dl2",
        ("p", "input-proton-dl2"): "IRFFITSWriter.input_proton_dl2",
        ("e", "input-electron-dl2"): "IRFFITSWriter.input_electron_dl2",
        ("o", "output-irf-file"): "IRFFITSWriter.output_irf_file",
        "irf-obs-time": "IRFFITSWriter.irf_obs_time",
        "fixed-gh-cut": "DL3FixedCuts.fixed_gh_cut",
        "fixed-theta-cut": "DL3FixedCuts.fixed_theta_cut",
        "allowed-tels": "DL3FixedCuts.allowed_tels",
        "overwrite": "IRFFITSWriter.overwrite",
    }

    flags = {
        "point-like": (
            {"IRFFITSWriter": {"point_like": True}},
            "Point like IRFs will be produced, otherwise Full Enclosure",
        ),
        "overwrite": (
            {"IRFFITSWriter": {"overwrite": True}},
            "overwrites output file",
        )
    }

    def setup(self):

        if self.output_irf_file.absolute().exists():
            if self.overwrite:
                self.log.warning(f"Overwriting {self.output_irf_file}")
                self.output_irf_file.unlink()
            else:
                raise ToolConfigurationError(
                    f"Output file {self.output_irf_file} already exists,"
                    " use --overwrite to overwrite"
                )

        filename = self.output_irf_file.name
        if not (filename.endswith('.fits') or filename.endswith('.fits.gz')):
            raise ValueError("f{filename} is not a correct compressed FITS file name (use .fits or .fits.gz).")

        if self.input_proton_dl2 and self.input_electron_dl2 is not None:
            self.only_gamma_irf = False
        else:
            self.only_gamma_irf = True

        self.event_sel = EventSelector(parent=self)
        self.fixed_cuts = DL3FixedCuts(parent=self)
        self.data_bin = DataBinning(parent=self)

        self.mc_particle = {
            "gamma": {
                "file": str(self.input_gamma_dl2),
                "target_spectrum": CRAB_MAGIC_JHEAP2015,
            },
        }
        Provenance().add_input_file(self.input_gamma_dl2)

        self.t_obs = self.irf_obs_time * u.hour

        # Read and update MC information
        if not self.only_gamma_irf:
            self.mc_particle["proton"] = {
                "file": str(self.input_proton_dl2),
                "target_spectrum": IRFDOC_PROTON_SPECTRUM,
            }

            self.mc_particle["electron"] = {
                "file": str(self.input_electron_dl2),
                "target_spectrum": IRFDOC_ELECTRON_SPECTRUM,
            }

            Provenance().add_input_file(self.input_proton_dl2)
            Provenance().add_input_file(self.input_electron_dl2)

        self.provenance_log = self.output_irf_file.parent / (
            self.name + ".provenance.log"
        )

    def start(self):

        for particle_type, p in self.mc_particle.items():
            self.log.info(f"Simulated {particle_type.title()} Events:")
            p["events"], p["simulation_info"] = read_mc_dl2_to_QTable(p["file"])

            if p["simulation_info"].viewcone.value == 0.0:
                p["mc_type"] = "point_like"
            else:
                p["mc_type"] = "diffuse"

            self.log.debug(f"Simulated {p['mc_type']} {particle_type.title()} Events:")

            # Calculating event weights for Background IRF
            if particle_type != "gamma":
                p["simulated_spectrum"] = PowerLaw.from_simulation(
                    p["simulation_info"], self.t_obs
                )

                p["events"]["weight"] = calculate_event_weights(
                    p["events"]["true_energy"],
                    p["target_spectrum"],
                    p["simulated_spectrum"],
                )

            for prefix in ("true", "reco"):
                k = f"{prefix}_source_fov_offset"
                p["events"][k] = calculate_source_fov_offset(p["events"], prefix=prefix)
            # calculate theta / distance between reco and assumed source position
            p["events"]["theta"] = calculate_theta(
                p["events"],
                assumed_source_az=p["events"]["true_az"],
                assumed_source_alt=p["events"]["true_alt"],
            )
            self.log.debug(p["simulation_info"])

        gammas = self.mc_particle["gamma"]["events"]

        self.log.info(f"Using fixed G/H cut of {self.fixed_cuts.fixed_gh_cut}")

        gammas = self.event_sel.filter_cut(gammas)
        gammas = self.fixed_cuts.allowed_tels_filter(gammas)
        gammas = self.fixed_cuts.gh_cut(gammas)

        if self.point_like:
            gammas = self.fixed_cuts.theta_cut(gammas)
            self.log.info('Theta cuts applied for point like IRF')

        # Binning of parameters used in IRFs
        true_energy_bins = self.data_bin.true_energy_bins()
        reco_energy_bins = self.data_bin.reco_energy_bins()
        migration_bins = self.data_bin.energy_migration_bins()
        source_offset_bins = self.data_bin.source_offset_bins()

        if self.mc_particle["gamma"]["mc_type"] == "point_like":
            mean_fov_offset = round(gammas["true_source_fov_offset"].mean().to_value(), 1)
            fov_offset_bins = [mean_fov_offset - 0.1, mean_fov_offset + 0.1] * u.deg
            self.log.info('Single offset for point like gamma MC')
        else:
            fov_offset_bins = self.data_bin.fov_offset_bins()
            self.log.info('Multiple offset for diffuse gamma MC')

        if not self.only_gamma_irf:
            background = table.vstack(
                [
                    self.mc_particle["proton"]["events"],
                    self.mc_particle["electron"]["events"],
                ]
            )

            background = self.event_sel.filter_cut(background)
            background = self.fixed_cuts.allowed_tels_filter(background)
            background = self.fixed_cuts.gh_cut(background)

            background_offset_bins = self.data_bin.bkg_fov_offset_bins()

        # For a fixed gh/theta cut, only a header value is added.
        # For energy dependent cuts, a new HDU should be created
        # GH_CUT and FOV_CUT are temporary non-standard header data
        extra_headers = {
            "TELESCOP": "CTA-N",
            "INSTRUME": "LST-" + " ".join(map(str, self.fixed_cuts.allowed_tels)),
            "FOVALIGN": "RADEC",
            "GH_CUT": self.fixed_cuts.fixed_gh_cut,
        }
        if self.point_like:
            self.log.info("Generating point_like IRF HDUs")
            extra_headers["RAD_MAX"] = str(self.fixed_cuts.fixed_theta_cut * u.deg)
        else:
            self.log.info("Generating Full-Enclosure IRF HDUs")

        # Write HDUs
        self.hdus = [fits.PrimaryHDU(), ]

        with np.errstate(invalid="ignore", divide="ignore"):
            if self.mc_particle["gamma"]["mc_type"] == "point_like":
                self.effective_area = effective_area_per_energy(
                    gammas,
                    self.mc_particle["gamma"]["simulation_info"],
                    true_energy_bins,
                )
                self.hdus.append(
                    create_aeff2d_hdu(
                        # add one dimension for single FOV offset
                        self.effective_area[..., np.newaxis],
                        true_energy_bins,
                        fov_offset_bins,
                        point_like=self.point_like,
                        extname="EFFECTIVE AREA",
                        **extra_headers,
                    )
                )
            else:
                self.effective_area = effective_area_per_energy_and_fov(
                    gammas,
                    self.mc_particle["gamma"]["simulation_info"],
                    true_energy_bins,
                    fov_offset_bins,
                )
                self.hdus.append(
                    create_aeff2d_hdu(
                        self.effective_area,
                        true_energy_bins,
                        fov_offset_bins,
                        point_like=self.point_like,
                        extname="EFFECTIVE AREA",
                        **extra_headers,
                    )
                )

        self.log.info("Effective Area HDU created")
        self.edisp = energy_dispersion(
            gammas,
            true_energy_bins,
            fov_offset_bins,
            migration_bins,
        )
        self.hdus.append(
            create_energy_dispersion_hdu(
                self.edisp,
                true_energy_bins,
                migration_bins,
                fov_offset_bins,
                point_like=self.point_like,
                extname="ENERGY DISPERSION",
                **extra_headers,
            )
        )
        self.log.info("Energy Dispersion HDU created")

        if not self.only_gamma_irf:
            self.background = background_2d(
                background,
                reco_energy_bins=reco_energy_bins,
                fov_offset_bins=background_offset_bins,
                t_obs=self.t_obs,
            )
            self.hdus.append(
                create_background_2d_hdu(
                    self.background.T,
                    reco_energy_bins,
                    background_offset_bins,
                    extname="BACKGROUND",
                    **extra_headers,
                )
            )
            self.log.info("Background HDU created")

        if not self.point_like:
            self.psf = psf_table(
                gammas,
                true_energy_bins,
                fov_offset_bins=fov_offset_bins,
                source_offset_bins=source_offset_bins,
            )
            self.hdus.append(
                create_psf_table_hdu(
                    self.psf,
                    true_energy_bins,
                    source_offset_bins,
                    fov_offset_bins,
                    extname="PSF",
                    **extra_headers,
                )
            )
            self.log.info("PSF HDU created")

    def finish(self):

        fits.HDUList(self.hdus).writeto(self.output_irf_file, overwrite=self.overwrite)
        Provenance().add_output_file(self.output_irf_file)