Source code for fermipy.diffuse.model_manager

# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Classes and utilities that manage fitting models for diffuse analyses
"""
from __future__ import absolute_import, division, print_function

import os
from collections import OrderedDict

import yaml

from fermipy import utils
from fermipy.roi_model import MapCubeSource, IsoSource
from fermipy.diffuse.name_policy import NameFactory
from fermipy.diffuse.binning import Component
from fermipy.diffuse.diffuse_src_manager import GalpropMapManager,\
    DiffuseModelManager, make_diffuse_comp_info_dict
from fermipy.diffuse.catalog_src_manager import CatalogSourceManager, make_catalog_comp_dict
from fermipy.diffuse.source_factory import SourceFactory
from fermipy.diffuse.spectral import SpectralLibrary


[docs]class ModelComponent(object): """ Small helper class to tie a ModelComponentInfo to a spectrum """ def __init__(self, **kwargs): """ """ self.info = kwargs.get('info') self.spectrum = kwargs.get('spectrum') self.par_overrides = kwargs.get('par_overrides') self.edisp_disable = kwargs.get('edisp_disable', False) if self.par_overrides is not None: for parname, pardict in self.par_overrides.items(): try: self.spectrum['spectral_pars'][parname].update(pardict) except KeyError: errmsg = "Failed to update parameter %s "%parname errmsg += "in source %s of spectral type %s"%(self.info.source_name, self.spectrum['SpectrumType']) raise KeyError(errmsg)
[docs]class ModelInfo(object): """ Small helper class to keep track of a single fitting model """ def __init__(self, **kwargs): """ """ self.model_name = kwargs.get('model_name') self.model_components = kwargs.get('model_components', OrderedDict()) @property def component_names(self): """ Return the list of name of the components """ return sorted(self.model_components.keys()) def __getitem__(self, key): """ Return a single model info by name """ return self.model_components[key]
[docs] def items(self): """ Return the key, value pairs of model components """ return self.model_components.items()
[docs] def edisp_disable_list(self): """ Return the list of source for which energy dispersion should be turned off """ l = [] for model_comp in self.model_components.values(): if model_comp.edisp_disable: l += [model_comp.info.source_name] return l
[docs] def make_srcmap_manifest(self, components, name_factory): """ Build a yaml file that specfies how to make the srcmap files for a particular model Parameters ---------- components : list The binning components used in this analysis name_factory : `NameFactory` Object that handles naming conventions Returns a dictionary that contains information about where to find the source maps for each component of the model """ ret_dict = {} for comp in components: compkey = comp.make_key('{ebin_name}_{evtype_name}') zcut = "zmax%i" % comp.zmax name_keys = dict(modelkey=self.model_name, zcut=zcut, ebin=comp.ebin_name, mktime='none', psftype=comp.evtype_name, coordsys=comp.coordsys) outsrcmap = name_factory.merged_srcmaps(**name_keys) ccube = name_factory.ccube(**name_keys) src_dict = {} for comp_name, model_comp in self.model_components.items(): comp_info = model_comp.info model_type = comp_info.model_type name_keys['sourcekey'] = comp_name if model_type in ['CatalogSources']: #sourcekey = comp_info.comp_key sources = comp_info.source_names name_keys['sourcekey'] = comp_info.catalog_info.catalog_name elif model_type in ['CompositeSource']: #sourcekey = comp_info.sourcekey name_keys['sourcekey'] = comp_info.sourcekey sources = [comp_info.source_name] else: #sourcekey = comp_name sources = [comp_info.source_name] src_dict[comp_name] = dict(sourcekey=comp_name, srcmap_file=name_factory.srcmaps(**name_keys), source_names=sources) comp_dict = dict(outsrcmap=outsrcmap, ccube=ccube, source_dict=src_dict) ret_dict[compkey] = comp_dict return ret_dict
[docs] def make_model_rois(self, components, name_factory): """ Make the fermipy roi_model objects for each of a set of binning components """ ret_dict = {} # Figure out which sources need to be split by components master_roi_source_info = {} sub_comp_sources = {} for comp_name, model_comp in self.model_components.items(): comp_info = model_comp.info if comp_info.components is None: master_roi_source_info[comp_name] = model_comp else: sub_comp_sources[comp_name] = model_comp # Build the xml for the master master_roi = SourceFactory.make_roi(master_roi_source_info) master_xml_mdl = name_factory.master_srcmdl_xml( modelkey=self.model_name) print("Writing master ROI model to %s" % master_xml_mdl) master_roi.write_xml(master_xml_mdl) ret_dict['master'] = master_roi # Now deal with the components for comp in components: zcut = "zmax%i" % comp.zmax compkey = "%s_%s" % (zcut, comp.make_key( '{ebin_name}_{evtype_name}')) # name_keys = dict(zcut=zcut, # modelkey=self.model_name, # component=compkey) comp_roi_source_info = {} for comp_name, model_comp in sub_comp_sources.items(): comp_info = model_comp.info if comp_info.selection_dependent: key = comp.make_key('{ebin_name}_{evtype_name}') elif comp_info.moving: key = zcut info_clone = comp_info.components[key].clone_and_merge_sub(key) comp_roi_source_info[comp_name] =\ ModelComponent(info=info_clone, spectrum=model_comp.spectrum) # Build the xml for the component comp_roi = SourceFactory.make_roi(comp_roi_source_info) comp_xml_mdl = name_factory.comp_srcmdl_xml(modelkey=self.model_name, component=compkey) print("Writing component ROI model to %s" % comp_xml_mdl) comp_roi.write_xml(comp_xml_mdl) ret_dict[compkey] = comp_roi return ret_dict
[docs]class ModelManager(object): """ Small helper class to create fitting models and manager XML files for fermipy This class contains a 'library', which is a dictionary of all the source components: specifically it maps: sourcekey : `model_component.ModelComponentInfo` """ def __init__(self, **kwargs): """ Keyword arguments ------------------- basedir : str Top level directory for finding files """ self._name_factory = NameFactory(**kwargs) self._dmm = kwargs.get('DiffuseModelManager', DiffuseModelManager(**kwargs)) self._gmm = kwargs.get('GalpropMapManager', GalpropMapManager(**kwargs)) self._csm = kwargs.get('CatalogSourceManager', CatalogSourceManager(**kwargs)) self._library = {} self._models = {} self._spec_lib = SpectralLibrary({})
[docs] def read_model_yaml(self, modelkey): """ Read the yaml file for the diffuse components """ model_yaml = self._name_factory.model_yaml(modelkey=modelkey, fullpath=True) model = yaml.safe_load(open(model_yaml)) return model
@property def dmm(self): """ Return the DiffuseModelManager """ return self._dmm @property def gmm(self): """ Return the GalpropMapManager """ return self._gmm @property def csm(self): """ Return the CatalogSourceManager """ return self._csm
[docs] def make_library(self, diffuse_yaml, catalog_yaml, binning_yaml): """ Build up the library of all the components Parameters ---------- diffuse_yaml : str Name of the yaml file with the library of diffuse component definitions catalog_yaml : str Name of the yaml file width the library of catalog split definitions binning_yaml : str Name of the yaml file with the binning definitions """ ret_dict = {} #catalog_dict = yaml.safe_load(open(catalog_yaml)) components_dict = Component.build_from_yamlfile(binning_yaml) diffuse_ret_dict = make_diffuse_comp_info_dict(GalpropMapManager=self._gmm, DiffuseModelManager=self._dmm, library=diffuse_yaml, components=components_dict) catalog_ret_dict = make_catalog_comp_dict(library=catalog_yaml, CatalogSourceManager=self._csm) ret_dict.update(diffuse_ret_dict['comp_info_dict']) ret_dict.update(catalog_ret_dict['comp_info_dict']) self._library.update(ret_dict) return ret_dict
[docs] def make_model_info(self, modelkey): """ Build a dictionary with the information for a particular model. Parameters ---------- modelkey : str Key used to identify this particular model Return `ModelInfo` """ model = self.read_model_yaml(modelkey) sources = model['sources'] components = OrderedDict() spec_model_yaml = self._name_factory.fullpath(localpath=model['spectral_models']) self._spec_lib.update(yaml.safe_load(open(spec_model_yaml))) for source, source_info in sources.items(): model_type = source_info.get('model_type', None) par_overrides = source_info.get('par_overides', None) version = source_info['version'] spec_type = source_info['SpectrumType'] edisp_disable = source_info.get('edisp_disable', False) sourcekey = "%s_%s" % (source, version) if model_type == 'galprop_rings': comp_info_dict = self.gmm.diffuse_comp_info_dicts(version) def_spec_type = spec_type['default'] for comp_key, comp_info in comp_info_dict.items(): model_comp = ModelComponent(info=comp_info, spectrum=\ self._spec_lib[spec_type.get(comp_key, def_spec_type)], par_overrides=par_overrides, edisp_disable=edisp_disable) components[comp_key] = model_comp elif model_type == 'Catalog': comp_info_dict = self.csm.split_comp_info_dict(source, version) def_spec_type = spec_type['default'] for comp_key, comp_info in comp_info_dict.items(): model_comp = ModelComponent(info=comp_info, spectrum=\ self._spec_lib[spec_type.get(comp_key, def_spec_type)], par_overrides=par_overrides, edisp_disable=edisp_disable) components[comp_key] = model_comp else: comp_info = self.dmm.diffuse_comp_info(sourcekey) model_comp = ModelComponent(info=comp_info, spectrum=self._spec_lib[spec_type], par_overrides=par_overrides, edisp_disable=edisp_disable) components[sourcekey] = model_comp ret_val = ModelInfo(model_name=modelkey, model_components=components) self._models[modelkey] = ret_val return ret_val
[docs] def make_srcmap_manifest(self, modelkey, components, data): """Build a yaml file that specfies how to make the srcmap files for a particular model Parameters ---------- modelkey : str Key used to identify this particular model components : list The binning components used in this analysis data : str Path to file containing dataset definition """ try: model_info = self._models[modelkey] except KeyError: model_info = self.make_model_info(modelkey) self._name_factory.update_base_dict(data) outfile = os.path.join('analysis', 'model_%s' % modelkey, 'srcmap_manifest_%s.yaml' % modelkey) manifest = model_info.make_srcmap_manifest( components, self._name_factory) outdir = os.path.dirname(outfile) try: os.makedirs(outdir) except OSError: pass utils.write_yaml(manifest, outfile)
[docs] def make_fermipy_config_yaml(self, modelkey, components, data, **kwargs): """Build a fermipy top-level yaml configuration file Parameters ---------- modelkey : str Key used to identify this particular model components : list The binning components used in this analysis data : str Path to file containing dataset definition """ model_dir = os.path.join('analysis', 'model_%s' % modelkey) hpx_order = kwargs.get('hpx_order', 9) self._name_factory.update_base_dict(data) try: model_info = self._models[modelkey] except KeyError: model_info = self.make_model_info(modelkey) model_rois = model_info.make_model_rois(components, self._name_factory) #source_names = model_info.component_names master_xml_mdl = self._name_factory.master_srcmdl_xml(modelkey=modelkey, fullpath=True) master_data = dict(scfile=self._name_factory.ft2file(fullpath=True), cacheft1=False) master_binning = dict(projtype='HPX', roiwidth=360., binsperdec=4, coordsys='GAL', hpx_ordering_scheme="RING", hpx_order=hpx_order, hpx_ebin=True) # master_fileio = dict(outdir=model_dir, # logfile=os.path.join(model_dir, 'fermipy.log')) master_fileio = dict(logfile='fermipy.log') master_gtlike = dict(irfs=self._name_factory.irfs(**kwargs), edisp_disable=model_info.edisp_disable_list(), use_external_srcmap=True) master_selection = dict(glat=0., glon=0., radius=180.) master_model = dict(catalogs=[os.path.join(model_dir, master_xml_mdl)]) master_plotting = dict(label_ts_threshold=1e9) master = dict(data=master_data, binning=master_binning, fileio=master_fileio, selection=master_selection, gtlike=master_gtlike, model=master_model, plotting=master_plotting, components=[]) fermipy_dict = master #comp_rois = {} for comp in components: zcut = "zmax%i" % comp.zmax compkey = "%s_%s" % (zcut, comp.make_key( '{ebin_name}_{evtype_name}')) comp_roi = model_rois[compkey] name_keys = dict(zcut=zcut, modelkey=modelkey, component=compkey, mktime='none', coordsys=comp.coordsys, fullpath=True) comp_data = dict(ltcube=self._name_factory.ltcube(**name_keys)) comp_selection = dict(logemin=comp.log_emin, logemax=comp.log_emax, zmax=comp.zmax, evtype=comp.evtype) comp_binning = dict(enumbins=comp.enumbins, hpx_order=min(comp.hpx_order, hpx_order), coordsys=comp.coordsys) comp_gtlike = dict(srcmap=self._name_factory.merged_srcmaps(**name_keys), bexpmap=self._name_factory.bexpcube(**name_keys), use_external_srcmap=True) comp_diffuse_xml = [self._name_factory.comp_srcmdl_xml(**name_keys)] #comp_roi_source_info = {} #diffuse_srcs = [] #for src in comp_roi.diffuse_sources: # if isinstance(src, MapCubeSource): # diffuse_srcs.append(dict(name=src.name, file=src.mapcube)) # elif isinstance(src, IsoSource): # diffuse_srcs.append(dict(name=src.name, file=src.fileFunction)) # else: # pass comp_model = dict(diffuse_xml=comp_diffuse_xml) sub_dict = dict(data=comp_data, binning=comp_binning, selection=comp_selection, gtlike=comp_gtlike, model=comp_model) fermipy_dict['components'].append(sub_dict) # Add placeholder diffuse sources #fermipy_dict['model']['diffuse'] = diffuse_srcs outfile = os.path.join(model_dir, 'config.yaml') print("Writing fermipy config file %s" % outfile) utils.write_yaml(fermipy_dict, outfile) return fermipy_dict
[docs] @staticmethod def get_sub_comp_info(source_info, comp): """Build and return information about a sub-component for a particular selection """ sub_comps = source_info.get('components', None) if sub_comps is None: return source_info.copy() moving = source_info.get('moving', False) selection_dependent = source_info.get('selection_dependent', False) if selection_dependent: key = comp.make_key('{ebin_name}_{evtype_name}') elif moving: key = "zmax%i" % comp.zmax ret_dict = source_info.copy() ret_dict.update(sub_comps[key]) return ret_dict
def make_library(**kwargs): """Build and return a ModelManager object and fill the associated model library """ library_yaml = kwargs.pop('library', 'models/library.yaml') comp_yaml = kwargs.pop('comp', 'config/binning.yaml') basedir = kwargs.pop('basedir', os.path.abspath('.')) model_man = kwargs.get('ModelManager', ModelManager(basedir=basedir)) model_comp_dict = model_man.make_library(library_yaml, library_yaml, comp_yaml) return dict(model_comp_dict=model_comp_dict, ModelManager=model_man)