Merge pull request #71 from /issues/70

Issues/70

augur/analyze.py

@@ -3,7 +3,7 @@
 from augur.utils.diff_utils import five_pt_stencil
 from augur import generate
 from augur.utils.config_io import parse_config
-from firecrown.parameters import ParamsMap
+from augur.utils.theory_utils import compute_new_theory_vector
 from astropy.table import Table
 import warnings
 
@@ -81,6 +81,8 @@ def __init__(self, config, likelihood=None, tools=None, req_params=None,
         self.norm_step = norm_step
         # Get the fiducial cosmological parameters
         self.pars_fid = tools.get_ccl_cosmology().__dict__['_params_init_kwargs']
+        # CCL Factory placeholder (for newer firecrown)
+        self.cf = None
 
         # Load the relevant section of the configuration file
         self.config = config['fisher']
@@ -134,9 +136,8 @@ def __init__(self, config, likelihood=None, tools=None, req_params=None,
         # Normalize the pivot point given the sampling region
         if self.norm_step:
             self.norm = self.par_bounds[:, 1] - self.par_bounds[:, 0]
-            self.x = (self.x - self.par_bounds[:, 0]) * 1/self.norm
 
-    def f(self, x, labels, pars_fid, sys_fid):
+    def f(self, x, labels, pars_fid, sys_fid, donorm=False):
         """
         Auxiliary Function that returns a theory vector evaluated at x.
         Labels are the name of the parameters x (with the same length and order)
@@ -153,6 +154,8 @@ def f(self, x, labels, pars_fid, sys_fid):
         sys_fid: dict
             Dictionary containing the fiducial `systematic` (required) parameters
             for the likelihood.
+        norm: bool
+            If `True` it normalizes the input parameters vector (useful for derivatives).
         Returns:
         --------
         f_out : np.ndarray
@@ -165,8 +168,9 @@ def f(self, x, labels, pars_fid, sys_fid):
             if isinstance(x, list):
                 x = np.array(x)
             # If we normalize the sampling we need to undo the normalization
-            if self.norm_step:
+            if donorm:
                 x = self.norm * x + self.par_bounds[:, 0]
+
             if x.ndim == 1:
                 _pars = pars_fid.copy()
                 _sys_pars = sys_fid.copy()
@@ -177,14 +181,9 @@ def f(self, x, labels, pars_fid, sys_fid):
                         _sys_pars.update({labels[i]: x[i]})
                     else:
                         raise ValueError(f'Parameter name {labels[i]} not recognized!')
-                self.tools.reset()
-                self.lk.reset()
-                pmap = ParamsMap(_sys_pars)
-                cosmo = ccl.Cosmology(**_pars)
-                self.lk.update(pmap)
-                self.tools.update(pmap)
-                self.tools.prepare(cosmo)
-                f_out = self.lk.compute_theory_vector(self.tools)
+
+                f_out = self.compute_new_theory_vector(_sys_pars, _pars)
+
             elif x.ndim == 2:
                 f_out = []
                 for i in range(len(labels)):
@@ -198,14 +197,7 @@ def f(self, x, labels, pars_fid, sys_fid):
                             _sys_pars.update({labels[j]: xi[j]})
                         else:
                             raise ValueError(f'Parameter name {labels[j]} not recognized')
-                    self.tools.reset()
-                    self.lk.reset()
-                    pmap = ParamsMap(_sys_pars)
-                    cosmo = ccl.Cosmology(**_pars)
-                    self.lk.update(pmap)
-                    self.tools.update(pmap)
-                    self.tools.prepare(cosmo)
-                    f_out.append(self.lk.compute_theory_vector(self.tools))
+                    f_out.append(self.compute_new_theory_vector(_sys_pars, _pars))
             return np.array(f_out)
 
     def get_derivatives(self, force=False, method='5pt_stencil', step=None):
@@ -228,20 +220,29 @@ def get_derivatives(self, force=False, method='5pt_stencil', step=None):
         # Compute the derivatives with respect to the parameters in var_pars at x
         if (self.derivatives is None) or (force):
             if '5pt_stencil' in method:
+                if self.norm_step:
+                    x_here = (self.x - self.par_bounds[:, 0]) * 1/self.norm
+                else:
+                    x_here = self.x
                 self.derivatives = five_pt_stencil(lambda y: self.f(y, self.var_pars, self.pars_fid,
-                                                   self.req_params),
-                                                   self.x, h=step)
+                                                   self.req_params, donorm=self.norm_step),
+                                                   x_here, h=step)
             elif 'numdifftools' in method:
                 import numdifftools as nd
                 if 'numdifftools_kwargs' in self.config.keys():
                     ndkwargs = self.config['numdifftools_kwargs']
                 else:
                     ndkwargs = {}
+                if self.norm_step:
+                    x_here = (self.x - self.par_bounds[:, 0]) * 1/self.norm
+                else:
+                    x_here = self.x
                 jacobian_calc = nd.Jacobian(lambda y: self.f(y, self.var_pars, self.pars_fid,
-                                                             self.req_params),
+                                                             self.req_params,
+                                                             donorm=self.norm_step),
                                             step=step,
                                             **ndkwargs)
-                self.derivatives = jacobian_calc(self.x).T
+                self.derivatives = jacobian_calc(x_here).T
             else:
                 raise ValueError(f'Selected method: `{method}` is not available. \
                                  Please select 5pt_stencil or numdifftools.')
@@ -329,10 +330,31 @@ def get_fisher_bias(self):
                 else:
                     raise ValueError('bias_params is required if no biased_dv file is passed')
 
-                self.biased_cls = self.f(_x_here, _labels_here, _pars_here, _sys_here) \
-                    - self.data_fid
+                self.biased_cls = self.f(_x_here, _labels_here, _pars_here, _sys_here,
+                                         donorm=False) - self.data_fid
 
             Bj = np.einsum('l, lm, jm', self.biased_cls, self.lk.inv_cov, self.derivatives)
             bi = np.einsum('ij, j', np.linalg.inv(self.Fij), Bj)
             self.bi = bi
             return self.bi
+
+    def compute_new_theory_vector(self, _sys_pars, _pars):
+        """
+        Utility function to update the likelihood and modeling tool objects to use a new
+        set of parameters and compute a new theory prediction
+
+        Parameters:
+        -----------
+        _sys_pars : dict,
+            Dictionary containing the "systematic" modeling parameters.
+        _pars : dict,
+            Dictionary containing the cosmological parameters
+
+        Returns:
+        --------
+        f_out : ndarray,
+            Predicted data vector for the given input parameters _sys_pars, _pars.
+        """
+        f_out = compute_new_theory_vector(self.lk, self.tools, _sys_pars, _pars)
+
+        return f_out

augur/generate.py

@@ -13,9 +13,11 @@
 from augur.tracers.two_point import ZDistFromFile
 from augur.utils.cov_utils import get_gaus_cov, get_SRD_cov, get_noise_power
 from augur.utils.cov_utils import TJPCovGaus
+from augur.utils.theory_utils import compute_new_theory_vector
 from packaging.version import Version
 import firecrown
-if Version(firecrown.__version__) >= Version('1.8'):
+
+if Version(firecrown.__version__) >= Version('1.8.0a'):
     import firecrown.likelihood.weak_lensing as wl
     import firecrown.likelihood.number_counts as nc
     from firecrown.likelihood.two_point import TwoPoint
@@ -25,7 +27,6 @@
     import firecrown.likelihood.gauss_family.statistic.source.number_counts as nc
     from firecrown.likelihood.gauss_family.statistic.two_point import TwoPoint
     from firecrown.likelihood.gauss_family.gaussian import ConstGaussian
-from firecrown.modeling_tools import ModelingTools
 from firecrown.parameters import ParamsMap
 from augur.utils.config_io import parse_config
 
@@ -347,22 +348,13 @@ def generate(config, return_all_outputs=False, write_sacc=True):
     lk = ConstGaussian(statistics=stats)
     # Pass the correct binning/tracers
     lk.read(S)
-    # The newest version of firecrown requires a modeling tool rather than a cosmology
-    pt_calculator = ccl.nl_pt.EulerianPTCalculator(
-        with_NC=False,
-        with_IA=True,
-        log10k_min=-4,  # Leaving defaults for now
-        log10k_max=2,
-        nk_per_decade=20,
-        cosmo=cosmo
-    )
+
     cosmo.compute_nonlin_power()
-    tools = ModelingTools(pt_calculator=pt_calculator)
-    lk.update(sys_params)
-    tools.update(sys_params)
-    tools.prepare(cosmo)
-    # Run the likelihood (to get the theory)
-    lk.compute_loglike(tools)
+    _pars = cosmo.__dict__['_params_init_kwargs']
+    # Populate ModelingTools and likelihood
+    tools = firecrown.modeling_tools.ModelingTools()
+    _, lk, tools = compute_new_theory_vector(lk, tools, sys_params, _pars, return_all=True)
+
     # Get all bandpower windows before erasing the placeholder sacc
     win_dict = {}
     ell_dict = {}

augur/utils/theory_utils.py

@@ -0,0 +1,90 @@
+from packaging.version import Version
+import pyccl as ccl
+import firecrown
+from firecrown.parameters import ParamsMap
+
+
+def compute_new_theory_vector(lk, tools, _sys_pars, _pars, cf=None, return_all=False):
+    """
+    Utility function to update the likelihood and modeling tool objects to use a new
+    set of parameters and compute a new theory prediction
+
+    Parameters:
+    -----------
+    lk : firecrown.likelihood.Likelihood,
+        Likelihood object to use.
+    tools: firecrown.ModelingTools,
+        ModelingTools object to use.
+    _sys_pars : dict,
+        Dictionary containing the "systematic" modeling parameters.
+    _pars : dict,
+        Dictionary containing the cosmological parameters
+    cf : firecrown.CCLFactory,
+        If passed, CCLFactory object to use.
+    return_all : bool,
+        If `False` it will just return the predicted data vector. Else,
+        it will return the internal ModelingTools and Likelihood objects.
+
+    Returns:
+    --------
+    f_out : ndarray,
+        Predicted data vector for the given input parameters _sys_pars, _pars.
+    lk : firecrown.likelihood.Likelihood,
+        Modified likelihood object centered on _pars, and _sys_pars.
+    tools: firecrown.ModelingTools,
+        Modified tools object with fiducial values on _pars, _sys_pars.
+    """
+    lk.reset()
+    tools.reset()
+    # Firecrown without CCLFactory
+    if Version(firecrown.__version__) < Version('1.8.0a'):
+        pmap = ParamsMap(_sys_pars)
+        cosmo = ccl.Cosmology(**_pars)
+        lk.update(pmap)
+        tools.update(pmap)
+        tools.prepare(cosmo)
+        f_out = lk.compute_theory_vector(tools)
+        if return_all:
+            return f_out, lk, tools
+        else:
+            return f_out
+
+    else:
+        from firecrown.ccl_factory import CCLFactory
+        dict_all = {**_sys_pars, **_pars}
+        extra_dict = {}
+        if dict_all['A_s'] is None:
+            extra_dict['amplitude_parameter'] = 'sigma8'
+            dict_all.pop('A_s')
+        else:
+            extra_dict['amplitude_parameter'] = 'A_s'
+            dict_all.pop('sigma8')
+
+        extra_dict['mass_split'] = dict_all['mass_split']
+        dict_all.pop('mass_split')
+        if 'extra_parameters' in dict_all.keys():
+            if 'camb' in dict_all['extra_parameters'].keys():
+                extra_dict['camb_extra_params'] = dict_all['extra_parameters']['camb']
+                if 'kmin' in dict_all['extra_parameters']['camb'].keys():
+                    extra_dict['camb_extra_params'].pop('kmin')
+            dict_all.pop('extra_parameters')
+        keys = list(dict_all.keys())
+
+        # Remove None values
+        for key in keys:
+            if (dict_all[key] is None) or (dict_all[key] == 'None'):
+                dict_all.pop(key)
+        if cf is None:
+            cf = CCLFactory(**extra_dict)
+            tools = firecrown.modeling_tools.ModelingTools(ccl_factory=cf)
+            tools.reset()
+        pmap = ParamsMap(dict_all)
+        cf.update(pmap)
+        tools.update(pmap)
+        tools.prepare()
+        lk.update(pmap)
+        f_out = lk.compute_theory_vector(tools)
+        if return_all:
+            return f_out, lk, tools
+        else:
+            return f_out