asgardpy.config.operations

Complexity

Total Complexity

21

Size/Duplication

Total Lines	155
Duplicated Lines	0 %

Importance

Changes

0

6 Functions

Rating	Name	Duplication	Size	Complexity
A	compound_model_dict_converstion()	0	19	1
C	recursive_merge_dicts()	0	52	9
A	all_model_templates()	0	12	2
A	check_gammapy_model()	0	13	3
A	get_model_template()	0	12	3
A	deep_update()	0	12	3

"""
Main AsgardpyConfig Operations Module
"""

import logging
from collections.abc import Mapping
from pathlib import Path

import numpy as np
from gammapy.modeling.models import Models, SkyModel

__all__ = [
    "all_model_templates",
    "compound_model_dict_converstion",
    "get_model_template",
    "recursive_merge_dicts",
    "deep_update",
]

CONFIG_PATH = Path(__file__).resolve().parent

log = logging.getLogger(__name__)


def all_model_templates():
    """
    Collect all Template Models provided in Asgardpy, and their small tag names.
    """
    template_files = sorted(list(CONFIG_PATH.glob("model_templates/model_template*yaml")))

    all_tags = []
    for file in template_files:
        all_tags.append(file.name.split("_")[-1].split(".")[0])
    all_tags = np.array(all_tags)

    return all_tags, template_files


def get_model_template(spec_model_tag):
    """
    Read a particular template model yaml filename to create an AsgardpyConfig
    object.
    """
    all_tags, template_files = all_model_templates()
    new_model_file = None

    for file, tag in zip(template_files, all_tags, strict=True):
        if spec_model_tag == tag:
            new_model_file = file
    return new_model_file


def check_gammapy_model(gammapy_model):
    """
    For a given object type, try to read it as a Gammapy Models object.
    """
    if isinstance(gammapy_model, Models | SkyModel):
        models_gpy = Models(gammapy_model)
    else:
        try:
            models_gpy = Models.read(gammapy_model)
        except KeyError:
            raise TypeError("%s File cannot be read by Gammapy Models", gammapy_model) from KeyError

    return models_gpy


def recursive_merge_dicts(base_config, extra_config):
    """
    recursively merge two dictionaries.
    Entries in extra_config override entries in base_config. The built-in
    update function cannot be used for hierarchical dicts.

    Also for the case when there is a list of dicts involved, one has to be
    more careful. The extra_config may have longer list of dicts as compared
    with the base_config, in which case, the extra items are simply added to
    the merged final list.

    Combined here are 2 options from SO.

    See:
    http://stackoverflow.com/questions/3232943/update-value-of-a-nested-dictionary-of-varying-depth/3233356#3233356
    and also
    https://stackoverflow.com/questions/3232943/update-value-of-a-nested-dictionary-of-varying-depth/18394648#18394648

    Parameters
    ----------
    base_config : dict
        dictionary to be merged
    extra_config : dict
        dictionary to be merged
    Returns
    -------
    final_config : dict
        merged dict
    """
    final_config = base_config.copy()

    for key, value in extra_config.items():
        if key in final_config and isinstance(final_config[key], list):
            new_config = []

            for key_, value_ in zip(final_config[key], value, strict=False):
                key_ = recursive_merge_dicts(key_ or {}, value_)
                new_config.append(key_)

            # For example moving from a smaller list of model parameters to a
            # longer list.
            if len(final_config[key]) < len(extra_config[key]):
                for value_ in value[len(final_config[key]) :]:
                    new_config.append(value_)
            final_config[key] = new_config

        elif key in final_config and isinstance(final_config[key], dict):
            final_config[key] = recursive_merge_dicts(final_config.get(key) or {}, value)
        else:
            final_config[key] = value

    return final_config


def deep_update(d, u):
    """
    Recursively update a nested dictionary.

    Just like in Gammapy, taken from: https://stackoverflow.com/a/3233356/19802442
    """
    for k, v in u.items():
        if isinstance(v, Mapping):
            d[k] = deep_update(d.get(k, {}), v)
        else:
            d[k] = v
    return d


def compound_model_dict_converstion(dict):
    """
    Given a Gammapy CompoundSpectralModel as a dict object, convert it into
    an Asgardpy form.
    """
    ebl_abs = dict["model2"]
    ebl_abs["alpha_norm"] = ebl_abs["parameters"][0]["value"]
    ebl_abs["redshift"] = ebl_abs["parameters"][1]["value"]
    ebl_abs.pop("parameters", None)

    dict["type"] = dict["model1"]["type"]
    dict["parameters"] = dict["model1"]["parameters"]
    dict["ebl_abs"] = ebl_abs

    dict.pop("model1", None)
    dict.pop("model2", None)
    dict.pop("operator", None)

    return dict