# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Covariance class."""
import numpy as np
import scipy
import matplotlib.pyplot as plt
from gammapy.utils.parallel import is_ray_initialized
from .parameter import Parameters
__all__ = ["Covariance"]
[docs]
class Covariance:
"""Parameter covariance class.
Parameters
----------
parameters : `~gammapy.modeling.Parameters`
Parameter list.
data : `~numpy.ndarray`
Covariance data array.
"""
def __init__(self, parameters, data=None):
self.parameters = parameters
if data is None:
data = np.diag([p.error**2 for p in self.parameters])
self._data = np.asanyarray(data, dtype=float)
@property
def shape(self):
"""Covariance shape."""
npars = len(self.parameters)
return npars, npars
@property
def data(self):
"""Covariance data as a `~numpy.ndarray`."""
return self._data
@data.setter
def data(self, value):
value = np.asanyarray(value)
npars = len(self.parameters)
shape = (npars, npars)
if value.shape != shape:
raise ValueError(
f"Invalid covariance shape: {value.shape}, expected {shape}"
)
self._data = value
@staticmethod
def _expand_factor_matrix(matrix, parameters):
"""Expand covariance matrix with zeros for frozen parameters."""
npars = len(parameters)
matrix_expanded = np.zeros((npars, npars))
mask_frozen = [par.frozen for par in parameters]
pars_index = [np.where(np.array(parameters) == p)[0][0] for p in parameters]
mask_duplicate = [pars_idx != idx for idx, pars_idx in enumerate(pars_index)]
mask = np.array(mask_frozen) | np.array(mask_duplicate)
free_parameters = ~(mask | mask[:, np.newaxis])
matrix_expanded[free_parameters] = matrix.ravel()
return matrix_expanded
[docs]
@classmethod
def from_factor_matrix(cls, parameters, matrix):
"""Set covariance from factor covariance matrix.
Used in the optimizer interface.
"""
npars = len(parameters)
if not matrix.shape == (npars, npars):
matrix = cls._expand_factor_matrix(matrix, parameters)
scales = [par.scale for par in parameters]
scale_matrix = np.outer(scales, scales)
data = scale_matrix * matrix
return cls(parameters, data=data)
[docs]
@classmethod
def from_stack(cls, covar_list):
"""Stack sub-covariance matrices from list.
Parameters
----------
covar_list : list of `Covariance`
List of sub-covariances.
Returns
-------
covar : `Covariance`
Stacked covariance.
"""
parameters = Parameters.from_stack([_.parameters for _ in covar_list])
covar = cls(parameters)
for subcovar in covar_list:
covar.set_subcovariance(subcovar)
return covar
[docs]
def get_subcovariance(self, parameters):
"""Get sub-covariance matrix.
Parameters
----------
parameters : `Parameters`
Sub list of parameters.
Returns
-------
covariance : `~numpy.ndarray`
Sub-covariance.
"""
idx = [self.parameters.index(par) for par in parameters]
data = self._data[np.ix_(idx, idx)]
return self.__class__(parameters=parameters, data=data)
[docs]
def set_subcovariance(self, covar):
"""Set sub-covariance matrix.
Parameters
----------
covar : `Covariance`
Sub-covariance.
"""
if is_ray_initialized():
# This copy is required to make the covariance setting work with ray
self._data = self._data.copy()
idx = [self.parameters.index(par) for par in covar.parameters]
if not np.allclose(self.data[np.ix_(idx, idx)], covar.data):
self.data[idx, :] = 0
self.data[:, idx] = 0
self._data[np.ix_(idx, idx)] = covar.data
[docs]
def plot_correlation(self, figsize=None, **kwargs):
"""Plot correlation matrix.
Parameters
----------
figsize : tuple, optional
Figure size. Default is None, which takes
(number_params*0.9, number_params*0.7).
**kwargs : dict
Keyword arguments passed to `~gammapy.visualization.plot_heatmap`.
Returns
-------
ax : `~matplotlib.axes.Axes`, optional
Matplotlib axes.
"""
from gammapy.visualization import annotate_heatmap, plot_heatmap
npars = len(self.parameters)
figsize = (npars * 0.9, npars * 0.7) if figsize is None else figsize
plt.figure(figsize=figsize)
ax = plt.gca()
kwargs.setdefault("cmap", "coolwarm")
names = self.parameters.names
im, cbar = plot_heatmap(
data=self.correlation,
col_labels=names,
row_labels=names,
ax=ax,
vmin=-1,
vmax=1,
cbarlabel="Correlation",
**kwargs,
)
annotate_heatmap(im=im)
return ax
@property
def correlation(self):
r"""Correlation matrix as a `numpy.ndarray`.
Correlation :math:`C` is related to covariance :math:`\Sigma` via:
.. math::
C_{ij} = \frac{ \Sigma_{ij} }{ \sqrt{\Sigma_{ii} \Sigma_{jj}} }
"""
err = np.sqrt(np.diag(self.data))
with np.errstate(invalid="ignore", divide="ignore"):
correlation = self.data / np.outer(err, err)
return np.nan_to_num(correlation)
@property
def scipy_mvn(self):
return scipy.stats.multivariate_normal(
self.parameters.value, self.data, allow_singular=True
)
def __str__(self):
return str(self.data)
def __array__(self):
return self.data
class CovarianceMixin:
"""Mixin class for covariance property on multi-components models"""
def _check_covariance(self):
if not self.parameters == self._covariance.parameters:
self._covariance = Covariance.from_stack(
[model.covariance for model in self._models]
)
@property
def covariance(self):
"""Covariance as a `~gammapy.modeling.Covariance` object."""
self._check_covariance()
for model in self._models:
self._covariance.set_subcovariance(model.covariance)
return self._covariance
@covariance.setter
def covariance(self, covariance):
self._check_covariance()
self._covariance.data = covariance
for model in self._models:
subcovar = self._covariance.get_subcovariance(model.covariance.parameters)
model.covariance = subcovar
