Completed
Push — main ( 4744cf...d735dd )
by Chaitanya
32s queued 23s
created

write_asgardpy_model_to_file()   B

Complexity

Conditions 5

Size

Total Lines 45
Code Lines 27

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 5
eloc 27
nop 3
dl 0
loc 45
rs 8.7653
c 0
b 0
f 0
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