Completed
Push — main ( 2edd10...a5707d )
by Chaitanya
28s queued 16s
created

asgardpy.config.generator.get_output_file_path()   A

Complexity

Conditions 4

Size

Total Lines 17
Code Lines 10

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 4
eloc 10
nop 2
dl 0
loc 17
rs 9.9
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 get_output_file_path(output_file, gammapy_model):
136
    """
137
    Conditions to return a Path variable of the output_file provided and if it
138
    is not provided, create a path in the model_templates sub-folder.
139
    """
140
    if not output_file:
141
        if isinstance(gammapy_model[0].spectral_model, CompoundSpectralModel):
142
            model_tag = gammapy_model[0].spectral_model.model1.tag[1] + "_ebl"
143
        else:
144
            model_tag = gammapy_model[0].spectral_model.tag[1]
145
146
        output_file = CONFIG_PATH / f"model_templates/model_template_{model_tag}.yaml"
147
        os.path.expandvars(output_file)
148
    else:
149
        if not isinstance(output_file, Path):
150
            output_file = Path(os.path.expandvars(output_file))
151
    return output_file
152
153
154
def write_asgardpy_model_to_file(gammapy_model, output_file=None, recursive_merge=True):
155
    """
156
    Read the Gammapy Models object and save it as AsgardpyConfig YAML file
157
    containing only the Model parameters, similar to the model templates
158
    available.
159
    """
160
    gammapy_model = check_gammapy_model(gammapy_model)
161
162
    asgardpy_config = gammapy_model_to_asgardpy_model_config(
163
        gammapy_model=gammapy_model[0],
164
        asgardpy_config_file=None,
165
        recursive_merge=recursive_merge,
166
    )
167
168
    output_file = get_output_file_path(output_file, gammapy_model)
169
170
    temp_ = asgardpy_config.model_dump(exclude_defaults=True)
171
    temp_["target"].pop("models_file", None)
172
173
    if isinstance(gammapy_model[0].spectral_model, CompoundSpectralModel):
174
        temp_["target"]["components"][0]["spectral"]["ebl_abs"]["filename"] = str(
175
            temp_["target"]["components"][0]["spectral"]["ebl_abs"]["filename"]
176
        )
177
    else:
178
        temp_["target"]["components"][0]["spectral"].pop("ebl_abs", None)
179
180
    yaml_ = yaml.dump(
181
        temp_,
182
        sort_keys=False,
183
        indent=4,
184
        width=80,
185
        default_flow_style=None,
186
    )
187
188
    output_file.write_text(yaml_)
189
190
191
# Combine everything!
192
class AsgardpyConfig(BaseConfig):
193
    """
194
    Asgardpy analysis configuration, based on Gammapy Analysis Config.
195
    """
196
197
    general: GeneralConfig = GeneralConfig()
198
199
    target: Target = Target()
200
201
    dataset3d: Dataset3DConfig = Dataset3DConfig()
202
    dataset1d: Dataset1DConfig = Dataset1DConfig()
203
204
    fit_params: FitConfig = FitConfig()
205
    flux_points_params: FluxPointsConfig = FluxPointsConfig()
206
207
    def __str__(self):
208
        """
209
        Display settings in pretty YAML format.
210
        """
211
        info = self.__class__.__name__ + "\n\n\t"
212
        data = self.to_yaml()
213
        data = data.replace("\n", "\n\t")
214
        info += data
215
        return info.expandtabs(tabsize=4)
216
217
    @classmethod
218
    def read(cls, path):
219
        """
220
        Reads from YAML file.
221
        """
222
        config = read_yaml(path)
223
        return AsgardpyConfig(**config)
224
225
    @classmethod
226
    def from_yaml(cls, config_str):
227
        """
228
        Create from YAML string.
229
        """
230
        settings = yaml.safe_load(config_str)
231
        return AsgardpyConfig(**settings)
232
233
    def write(self, path, overwrite=False):
234
        """
235
        Write to YAML file.
236
        """
237
        path = make_path(path)
238
        if path.exists() and not overwrite:
239
            raise OSError(f"File exists already: {path}")
240
        path.write_text(self.to_yaml())
241
242
    def to_yaml(self):
243
        """
244
        Convert to YAML string.
245
        """
246
        data = json.loads(self.model_dump_json())
247
        return yaml.dump(data, sort_keys=False, indent=4, width=80, default_flow_style=None)
248
249
    def set_logging(self):
250
        """
251
        Set logging config.
252
        Calls ``logging.basicConfig``, i.e. adjusts global logging state.
253
        """
254
        self.general.log.level = self.general.log.level.upper()
255
        logging.basicConfig(**self.general.log.model_dump())
256
        log.info("Setting logging config: %s", self.general.log.model_dump())
257
258
    def update(self, config=None, merge_recursive=False):
259
        """
260
        Update config with provided settings.
261
        Parameters
262
        ----------
263
        config : string dict or `AsgardpyConfig` object
264
            The other configuration settings provided in dict() syntax.
265
        merge_recursive : bool
266
            Perform a recursive merge from the other config onto the parent config.
267
268
        Returns
269
        -------
270
        config : `AsgardpyConfig` object
271
            Updated config object.
272
        """
273
        other = check_config(config)
274
275
        # Special case of when only updating target model parameters from a
276
        # separate file, where the name of the source is not provided.
277
        if other.target.components[0].name == "":
278
            merge_recursive = True
279
280
        if merge_recursive:
281
            config_new = recursive_merge_dicts(
282
                self.model_dump(exclude_defaults=True), other.model_dump(exclude_defaults=True)
283
            )
284
        else:
285
            config_new = deep_update(
286
                self.model_dump(exclude_defaults=True), other.model_dump(exclude_defaults=True)
287
            )
288
        return AsgardpyConfig(**config_new)
289