asgardpy.config.generator.write_asgardpy_model_to_file() - Code Metrics - Inspection of "Add LST1 DL4 Crab data" - chaimain/asgardpy - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — main (#184)

by Chaitanya

created 2024-10-17 20:47 UTC

write_asgardpy_model_to_file() B

↳ Parent: asgardpy.config.generator

Complexity

Conditions

Size

Total Lines	45
Code Lines	27

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	5
eloc	27
nop	3
dl	0
loc	45
rs	8.7653
c	0
b	0
f	0

"""
Main AsgardpyConfig Generator Module
"""

import json
import logging
import os
from enum import Enum
from pathlib import Path

import yaml
from gammapy.modeling.models import CompoundSpectralModel
from gammapy.utils.scripts import make_path, read_yaml

from asgardpy.analysis.step_base import AnalysisStepEnum
from asgardpy.base import BaseConfig, PathType
from asgardpy.config.operations import (
    CONFIG_PATH,
    check_gammapy_model,
    compound_model_dict_converstion,
    deep_update,
    recursive_merge_dicts,
)
from asgardpy.data import (
    Dataset1DConfig,
    Dataset3DConfig,
    FitConfig,
    FluxPointsConfig,
    Target,
)

__all__ = [
    "AsgardpyConfig",
    "GeneralConfig",
    "gammapy_model_to_asgardpy_model_config",
    "write_asgardpy_model_to_file",
]

log = logging.getLogger(__name__)


# Other general config params
class LogConfig(BaseConfig):
    """Config section for main logging information."""

    level: str = "info"
    filename: str = ""
    filemode: str = "w"
    format: str = ""
    datefmt: str = ""


class ParallelBackendEnum(str, Enum):
    """Config section for list of parallel processing backend methods."""

    multi = "multiprocessing"
    ray = "ray"


class GeneralConfig(BaseConfig):
    """Config section for general information for running AsgardpyAnalysis."""

    log: LogConfig = LogConfig()
    outdir: PathType = "None"
    n_jobs: int = 1
    parallel_backend: ParallelBackendEnum = ParallelBackendEnum.multi
    steps: list[AnalysisStepEnum] = []
    overwrite: bool = True
    stacked_dataset: bool = False


def check_config(config):
    """
    For a given object type, try to read it as an AsgardpyConfig object.
    """
    if isinstance(config, str | Path):
        if Path(config).is_file():
            AConfig = AsgardpyConfig.read(config)
        else:
            AConfig = AsgardpyConfig.from_yaml(config)
    elif isinstance(config, AsgardpyConfig):
        AConfig = config
    else:
        raise TypeError(f"Invalid type: {config}")

    return AConfig


def gammapy_model_to_asgardpy_model_config(gammapy_model, asgardpy_config_file=None, recursive_merge=True):
    """
    Read the Gammapy Models object and save it as AsgardpyConfig object.

    The gammapy_model object may be a YAML config filename/path/object or a
    Gammapy Models object itself.

    Return
    ------
    asgardpy_config: `asgardpy.config.generator.AsgardpyConfig`
        Updated AsgardpyConfig object
    """

    models_gpy = check_gammapy_model(gammapy_model)

    models_gpy_dict = models_gpy.to_dict()

    if not asgardpy_config_file:
        asgardpy_config = AsgardpyConfig()  # Default object
        # Remove any name values in the model dict
        models_gpy_dict["components"][0].pop("datasets_names", None)
        models_gpy_dict["components"][0].pop("name", None)
    else:
        asgardpy_config = check_config(asgardpy_config_file)

    # For EBL part only
    if "model1" in models_gpy_dict["components"][0]["spectral"].keys():
        models_gpy_dict["components"][0]["spectral"] = compound_model_dict_converstion(
            models_gpy_dict["components"][0]["spectral"]
        )

    asgardpy_config_target_dict = asgardpy_config.model_dump()["target"]

    if recursive_merge:
        temp_target_dict = recursive_merge_dicts(asgardpy_config_target_dict, models_gpy_dict)
    else:
        # Use when there are nans present in the other config file, which are
        # the defaults in Gammapy, but NOT in Asgardpy.
        # E.g. test data Fermi-3fhl-crab model file
        temp_target_dict = deep_update(asgardpy_config_target_dict, models_gpy_dict)

    asgardpy_config.target = temp_target_dict

    return asgardpy_config


def write_asgardpy_model_to_file(gammapy_model, output_file=None, recursive_merge=True):
    """
    Read the Gammapy Models object and save it as AsgardpyConfig YAML file
    containing only the Model parameters, similar to the model templates
    available.
    """
    gammapy_model = check_gammapy_model(gammapy_model)

    asgardpy_config = gammapy_model_to_asgardpy_model_config(
        gammapy_model=gammapy_model[0],
        asgardpy_config_file=None,
        recursive_merge=recursive_merge,
    )

    if not output_file:
        if isinstance(gammapy_model[0].spectral_model, CompoundSpectralModel):
            model_tag = gammapy_model[0].spectral_model.model1.tag[1] + "_ebl"
        else:
            model_tag = gammapy_model[0].spectral_model.tag[1]

        output_file = CONFIG_PATH / f"model_templates/model_template_{model_tag}.yaml"
        os.path.expandvars(output_file)
    else:
        if not isinstance(output_file, Path):
            output_file = Path(os.path.expandvars(output_file))

    temp_ = asgardpy_config.model_dump(exclude_defaults=True)
    temp_["target"].pop("models_file", None)

    if isinstance(gammapy_model[0].spectral_model, CompoundSpectralModel):
        temp_["target"]["components"][0]["spectral"]["ebl_abs"]["filename"] = str(
            temp_["target"]["components"][0]["spectral"]["ebl_abs"]["filename"]
        )
    else:
        temp_["target"]["components"][0]["spectral"].pop("ebl_abs", None)

    yaml_ = yaml.dump(
        temp_,
        sort_keys=False,
        indent=4,
        width=80,
        default_flow_style=None,
    )

    output_file.write_text(yaml_)


# Combine everything!
class AsgardpyConfig(BaseConfig):
    """
    Asgardpy analysis configuration, based on Gammapy Analysis Config.
    """

    general: GeneralConfig = GeneralConfig()

    target: Target = Target()

    dataset3d: Dataset3DConfig = Dataset3DConfig()
    dataset1d: Dataset1DConfig = Dataset1DConfig()

    fit_params: FitConfig = FitConfig()
    flux_points_params: FluxPointsConfig = FluxPointsConfig()

    def __str__(self):
        """
        Display settings in pretty YAML format.
        """
        info = self.__class__.__name__ + "\n\n\t"
        data = self.to_yaml()
        data = data.replace("\n", "\n\t")
        info += data
        return info.expandtabs(tabsize=4)

    @classmethod
    def read(cls, path):
        """
        Reads from YAML file.
        """
        config = read_yaml(path)
        return AsgardpyConfig(**config)

    @classmethod
    def from_yaml(cls, config_str):
        """
        Create from YAML string.
        """
        settings = yaml.safe_load(config_str)
        return AsgardpyConfig(**settings)

    def write(self, path, overwrite=False):
        """
        Write to YAML file.
        """
        path = make_path(path)
        if path.exists() and not overwrite:
            raise OSError(f"File exists already: {path}")
        path.write_text(self.to_yaml())

    def to_yaml(self):
        """
        Convert to YAML string.
        """
        # Here using `dict()` instead of `json()` would be more natural.
        # We should change this once pydantic adds support for custom encoders
        # to `dict()`. See https://github.com/samuelcolvin/pydantic/issues/1043
        data = json.loads(self.model_dump_json())
        return yaml.dump(data, sort_keys=False, indent=4, width=80, default_flow_style=None)

    def set_logging(self):
        """
        Set logging config.
        Calls ``logging.basicConfig``, i.e. adjusts global logging state.
        """
        self.general.log.level = self.general.log.level.upper()
        logging.basicConfig(**self.general.log.model_dump())
        log.info("Setting logging config: %s", self.general.log.model_dump())

    def update(self, config=None, merge_recursive=False):
        """
        Update config with provided settings.
        Parameters
        ----------
        config : string dict or `AsgardpyConfig` object
            The other configuration settings provided in dict() syntax.
        merge_recursive : bool
            Perform a recursive merge from the other config onto the parent config.

        Returns
        -------
        config : `AsgardpyConfig` object
            Updated config object.
        """
        other = check_config(config)

        # Special case of when only updating target model parameters from a
        # separate file, where the name of the source is not provided.
        if other.target.components[0].name == "":
            merge_recursive = True

        if merge_recursive:
            config_new = recursive_merge_dicts(
                self.model_dump(exclude_defaults=True), other.model_dump(exclude_defaults=True)
            )
        else:
            config_new = deep_update(
                self.model_dump(exclude_defaults=True), other.model_dump(exclude_defaults=True)
            )
        return AsgardpyConfig(**config_new)


1			"""
2			Main AsgardpyConfig Generator Module
3			"""
4
5			import json
6			import logging
7			import os
8			from enum import Enum
9			from pathlib import Path
10
11			import yaml
12			from gammapy.modeling.models import CompoundSpectralModel
13			from gammapy.utils.scripts import make_path, read_yaml
14
15			from asgardpy.analysis.step_base import AnalysisStepEnum
16			from asgardpy.base import BaseConfig, PathType
17			from asgardpy.config.operations import (
18			CONFIG_PATH,
19			check_gammapy_model,
20			compound_model_dict_converstion,
21			deep_update,
22			recursive_merge_dicts,
23			)
24			from asgardpy.data import (
25			Dataset1DConfig,
26			Dataset3DConfig,
27			FitConfig,
28			FluxPointsConfig,
29			Target,
30			)
31
32			__all__ = [
33			"AsgardpyConfig",
34			"GeneralConfig",
35			"gammapy_model_to_asgardpy_model_config",
36			"write_asgardpy_model_to_file",
37			]
38
39			log = logging.getLogger(__name__)
40
41
42			# Other general config params
43			class LogConfig(BaseConfig):
44			"""Config section for main logging information."""
45
46			level: str = "info"
47			filename: str = ""
48			filemode: str = "w"
49			format: str = ""
50			datefmt: str = ""
51
52
53			class ParallelBackendEnum(str, Enum):
54			"""Config section for list of parallel processing backend methods."""
55
56			multi = "multiprocessing"
57			ray = "ray"
58
59
60			class GeneralConfig(BaseConfig):
61			"""Config section for general information for running AsgardpyAnalysis."""
62
63			log: LogConfig = LogConfig()
64			outdir: PathType = "None"
65			n_jobs: int = 1
66			parallel_backend: ParallelBackendEnum = ParallelBackendEnum.multi
67			steps: list[AnalysisStepEnum] = []
68			overwrite: bool = True
69			stacked_dataset: bool = False
70
71
72			def check_config(config):
73			"""
74			For a given object type, try to read it as an AsgardpyConfig object.
75			"""
76			if isinstance(config, str \| Path):
77			if Path(config).is_file():
78			AConfig = AsgardpyConfig.read(config)
79			else:
80			AConfig = AsgardpyConfig.from_yaml(config)
81			elif isinstance(config, AsgardpyConfig):
82			AConfig = config
83			else:
84			raise TypeError(f"Invalid type: {config}")
85
86			return AConfig
87
88
89			def gammapy_model_to_asgardpy_model_config(gammapy_model, asgardpy_config_file=None, recursive_merge=True):
90			"""
91			Read the Gammapy Models object and save it as AsgardpyConfig object.
92
93			The gammapy_model object may be a YAML config filename/path/object or a
94			Gammapy Models object itself.
95
96			Return
97			------
98			asgardpy_config: `asgardpy.config.generator.AsgardpyConfig`
99			Updated AsgardpyConfig object
100			"""
101
102			models_gpy = check_gammapy_model(gammapy_model)
103
104			models_gpy_dict = models_gpy.to_dict()
105
106			if not asgardpy_config_file:
107			asgardpy_config = AsgardpyConfig() # Default object
108			# Remove any name values in the model dict
109			models_gpy_dict["components"][0].pop("datasets_names", None)
110			models_gpy_dict["components"][0].pop("name", None)
111			else:
112			asgardpy_config = check_config(asgardpy_config_file)
113
114			# For EBL part only
115			if "model1" in models_gpy_dict["components"][0]["spectral"].keys():
116			models_gpy_dict["components"][0]["spectral"] = compound_model_dict_converstion(
117			models_gpy_dict["components"][0]["spectral"]
118			)
119
120			asgardpy_config_target_dict = asgardpy_config.model_dump()["target"]
121
122			if recursive_merge:
123			temp_target_dict = recursive_merge_dicts(asgardpy_config_target_dict, models_gpy_dict)
124			else:
125			# Use when there are nans present in the other config file, which are
126			# the defaults in Gammapy, but NOT in Asgardpy.
127			# E.g. test data Fermi-3fhl-crab model file
128			temp_target_dict = deep_update(asgardpy_config_target_dict, models_gpy_dict)
129
130			asgardpy_config.target = temp_target_dict
131
132			return asgardpy_config
133
134
135			def write_asgardpy_model_to_file(gammapy_model, output_file=None, recursive_merge=True):
136			"""
137			Read the Gammapy Models object and save it as AsgardpyConfig YAML file
138			containing only the Model parameters, similar to the model templates
139			available.
140			"""
141			gammapy_model = check_gammapy_model(gammapy_model)
142
143			asgardpy_config = gammapy_model_to_asgardpy_model_config(
144			gammapy_model=gammapy_model[0],
145			asgardpy_config_file=None,
146			recursive_merge=recursive_merge,
147			)
148
149			if not output_file:
150			if isinstance(gammapy_model[0].spectral_model, CompoundSpectralModel):
151			model_tag = gammapy_model[0].spectral_model.model1.tag[1] + "_ebl"
152			else:
153			model_tag = gammapy_model[0].spectral_model.tag[1]
154
155			output_file = CONFIG_PATH / f"model_templates/model_template_{model_tag}.yaml"
156			os.path.expandvars(output_file)
157			else:
158			if not isinstance(output_file, Path):
159			output_file = Path(os.path.expandvars(output_file))
160
161			temp_ = asgardpy_config.model_dump(exclude_defaults=True)
162			temp_["target"].pop("models_file", None)
163
164			if isinstance(gammapy_model[0].spectral_model, CompoundSpectralModel):
165			temp_["target"]["components"][0]["spectral"]["ebl_abs"]["filename"] = str(
166			temp_["target"]["components"][0]["spectral"]["ebl_abs"]["filename"]
167			)
168			else:
169			temp_["target"]["components"][0]["spectral"].pop("ebl_abs", None)
170
171			yaml_ = yaml.dump(
172			temp_,
173			sort_keys=False,
174			indent=4,
175			width=80,
176			default_flow_style=None,
177			)
178
179			output_file.write_text(yaml_)
180
181
182			# Combine everything!
183			class AsgardpyConfig(BaseConfig):
184			"""
185			Asgardpy analysis configuration, based on Gammapy Analysis Config.
186			"""
187
188			general: GeneralConfig = GeneralConfig()
189
190			target: Target = Target()
191
192			dataset3d: Dataset3DConfig = Dataset3DConfig()
193			dataset1d: Dataset1DConfig = Dataset1DConfig()
194
195			fit_params: FitConfig = FitConfig()
196			flux_points_params: FluxPointsConfig = FluxPointsConfig()
197
198			def __str__(self):
199			"""
200			Display settings in pretty YAML format.
201			"""
202			info = self.__class__.__name__ + "\n\n\t"
203			data = self.to_yaml()
204			data = data.replace("\n", "\n\t")
205			info += data
206			return info.expandtabs(tabsize=4)
207
208			@classmethod
209			def read(cls, path):
210			"""
211			Reads from YAML file.
212			"""
213			config = read_yaml(path)
214			return AsgardpyConfig(**config)
215
216			@classmethod
217			def from_yaml(cls, config_str):
218			"""
219			Create from YAML string.
220			"""
221			settings = yaml.safe_load(config_str)
222			return AsgardpyConfig(**settings)
223
224			def write(self, path, overwrite=False):
225			"""
226			Write to YAML file.
227			"""
228			path = make_path(path)
229			if path.exists() and not overwrite:
230			raise OSError(f"File exists already: {path}")
231			path.write_text(self.to_yaml())
232
233			def to_yaml(self):
234			"""
235			Convert to YAML string.
236			"""
237			# Here using `dict()` instead of `json()` would be more natural.
238			# We should change this once pydantic adds support for custom encoders
239			# to `dict()`. See https://github.com/samuelcolvin/pydantic/issues/1043
240			data = json.loads(self.model_dump_json())
241			return yaml.dump(data, sort_keys=False, indent=4, width=80, default_flow_style=None)
242
243			def set_logging(self):
244			"""
245			Set logging config.
246			Calls ``logging.basicConfig``, i.e. adjusts global logging state.
247			"""
248			self.general.log.level = self.general.log.level.upper()
249			logging.basicConfig(**self.general.log.model_dump())
250			log.info("Setting logging config: %s", self.general.log.model_dump())
251
252			def update(self, config=None, merge_recursive=False):
253			"""
254			Update config with provided settings.
255			Parameters
256			----------
257			config : string dict or `AsgardpyConfig` object
258			The other configuration settings provided in dict() syntax.
259			merge_recursive : bool
260			Perform a recursive merge from the other config onto the parent config.
261
262			Returns
263			-------
264			config : `AsgardpyConfig` object
265			Updated config object.
266			"""
267			other = check_config(config)
268
269			# Special case of when only updating target model parameters from a
270			# separate file, where the name of the source is not provided.
271			if other.target.components[0].name == "":
272			merge_recursive = True
273
274			if merge_recursive:
275			config_new = recursive_merge_dicts(
276			self.model_dump(exclude_defaults=True), other.model_dump(exclude_defaults=True)
277			)
278			else:
279			config_new = deep_update(
280			self.model_dump(exclude_defaults=True), other.model_dump(exclude_defaults=True)
281			)
282			return AsgardpyConfig(**config_new)
283

chaimain / asgardpy

Pull Request — main (#184)

write_asgardpy_model_to_file() B

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like