asgardpy.data.target.set_models()

Complexity

Conditions

15

Size

Total Lines	82
Code Lines	39

Duplication

Lines	0
Ratio	0 %

Importance

Changes

0

"""
def extend_bkg_models(models, all_datasets, datasets_with_fov_bkg_model):
    """ """
    if len(datasets_with_fov_bkg_model) > 0:
        # For extending a Background Model for a given 3D dataset name
        bkg_models = []

        for dataset in all_datasets:
            if dataset.name in datasets_with_fov_bkg_model:
                # Check if it is of MapDataset or MapDatasetOnOff type and
                # no associated background model exists already.
                if "MapDataset" in dataset.tag and dataset.background_model is None:
                    bkg_models.append(FoVBackgroundModel(dataset_name=f"{dataset.name}-bkg"))

        models.extend(bkg_models)

    return models


def update_models_datasets_names(models, source_name, datasets_name_list):
    """ """
    if len(models) > 1:
        models[source_name].datasets_names = datasets_name_list

        # Check if FoVBackgroundModel is provided
        bkg_model_name = [m_name for m_name in models.names if "bkg" in m_name]

        if len(bkg_model_name) > 0:
            # Remove the -bkg part of the name of the FoVBackgroundModel to get
            # the appropriate datasets name.
            for bkg_name in bkg_model_name:
                models[bkg_name].datasets_names = [bkg_name[:-4]]
    else:
        models.datasets_names = datasets_name_list

    return models


def set_models(
    config_target,
    datasets,
    datasets_name_list=None,
    models=None,
):
    """
    Set models on given Datasets.

    Parameters
    ----------
    config_target: `AsgardpyConfig.target`
        AsgardpyConfig containing target information.
    datasets: `gammapy.datasets.Datasets`
        Datasets object
    dataset_name_list: list
        List of datasets_names to be used on the Models, before assigning them
        to the given datasets.
    models : `~gammapy.modeling.models.Models` or str of file location or None
        Models object or YAML models string

    Returns
    -------
    datasets: `gammapy.datasets.Datasets`
        Datasets object with Models assigned.
    """
    # Have some checks on argument types
    if isinstance(models, DatasetModels | list):
        models = Models(models)
    elif isinstance(models, PathType):
        models = Models.read(models)
    elif models is None:
        models = Models(models)
    else:
        raise TypeError(f"Invalid type: {type(models)}")

    if len(models) == 0:
        if config_target.components[0].name != "":
            models = read_models_from_asgardpy_config(config_target)
        else:
            raise ValueError("No input for Models provided for the Target source!")
    else:
        models = apply_selection_mask_to_models(
            list_sources=models,
            target_source=config_target.source_name,
            roi_radius=config_target.roi_selection.roi_radius,