likelihood:
  solike.SimulatedLensingLikelihood:
    sim_number: 1

theory:
  classy:
    extra_args:
    output: lCl, tCl

params:
  n_s:
    prior:
      min: 0.8
      max: 1.2        

from cobaya.yaml import yaml_load_file
from cobaya.model import get_model

info = yaml_load_file('model.yaml')
model = get_model(info)

model.loglikes({'n_s': 0.965})[0]

import os
from pkg_resources import resource_filename

from .ps import BinnedPSLikelihood


class LensingLikelihood(BinnedPSLikelihood):
    class_options = {"kind": "pp"}


class SimulatedLensingLikelihood(LensingLikelihood):
    class_options = {
        "dataroot": resource_filename(
            "solike", "data/simulated_clkk_SO_Apr17_mv_nlkk_deproj0_SENS1_fsky_16000_iterOn_20191109"
        ),
        "cl_file": "simulated_clkk_SO_Apr17_mv_nlkk_deproj0_SENS1_fsky_16000_iterOn_20191109_sim_{:02d}_bandpowers.txt",
        "cov_file": "simulated_clkk_SO_Apr17_mv_nlkk_deproj0_SENS1_fsky_16000_iterOn_20191109_binned_covmat.txt",
        "sim_number": 0,
    }

    def initialize(self):
        self.datapath = os.path.join(self.dataroot, self.cl_file.format(self.sim_number))
        self.covpath = os.path.join(self.dataroot, self.cov_file)
        super().initialize()

import numpy as np

from .utils import binner
from .gaussian import GaussianLikelihood


class PSLikelihood(GaussianLikelihood):

    class_options = {"name": "TT", "kind": "tt", "lmax": 6000}

    def get_requirements(self):
        return {"Cl": {self.kind: self.lmax}}

    def _get_Cl(self):
        return self.theory.get_Cl(ell_factor=True)

    def _get_theory(self, **params_values):
        cl_theory = self.get_Cl()
        return cl_theory[self.kind]


class BinnedPSLikelihood(PSLikelihood):

    def _get_data(self):
        lefts, rights, bandpowers = np.loadtxt(self.datapath, 
                                               unpack=True)

        bin_centers = (lefts + rights) / 2
        self.bin_edges = np.append(lefts, [rights[-1]])

        return bin_centers, bandpowers

    def _get_theory(self, **params_values):
        cl_theory = self._get_Cl()
        _, theory = binner(cl_theory["ell"], 
                           cl_theory[self.kind], 
                           self.bin_edges)
        return theory

import numpy as np

from cobaya.likelihood import Likelihood

from .gaussian_data import GaussianData

class GaussianLikelihood(Likelihood):

    class_options = {
        "name": "Gaussian",
        "datapath": None,
        "covpath": None,
    }

    def initialize(self):
        x, y = self._get_data()
        cov = self._get_cov()
        self.data = GaussianData(self.name, x, y, cov)

    def _get_data(self):
        x, y = np.loadtxt(self.datapath, unpack=True)
        return x, y

    def _get_cov(self):
        cov = np.loadtxt(self.covpath)
        return cov

    def _get_theory(self, **kwargs):
        raise NotImplementedError

    def logp(self, **params_values):
        theory = self._get_theory(**params_values)
        return self.data.loglike(theory)

import holoviews as hv
from scipy.linalg import cholesky, LinAlgError


def multivariate_normal_logpdf(theory, data, cov, inv_cov, log_det):
    const = np.log(2 * np.pi) * (-len(data) / 2) + log_det * (-1 / 2)
    delta = data - theory
    return -0.5 * np.dot(delta, inv_cov.dot(delta)) + const


class GaussianData(object):
    """Named multivariate gaussian data
    """

    def __init__(self, name, x, y, cov):

        self.name = str(name)

        if not (len(x) == len(y) and cov.shape == (len(x), len(x))):
            raise ValueError(f"Incompatible shapes! x={x.shape}, y={y.shape}, cov={cov.shape}")

        self.x = x
        self.y = y
        self.cov = cov
        try:
            self.cholesky = cholesky(cov)
        except LinAlgError:
            raise ValueError("Covariance is not SPD!")
        self.inv_cov = np.linalg.inv(self.cov)
        self.log_det = np.linalg.slogdet(self.cov)[1]

    def __len__(self):
        return len(self.x)

    def loglike(self, theory):
        return multivariate_normal_logpdf(theory, self.y, self.cov, self.inv_cov, self.log_det)
     

import numpy as np
import pandas as pd


class PoissonData(object):
    """Poisson-process-generated data.

    Parameters
    ----------
    catalog : pd.DataFrame
        Catalog of observed data.
    columns : list
        Columns of catalog relevant for computing poisson rate.
    """

    def __init__(self, name, catalog, columns):
        self.name = str(name)

        self.catalog = pd.DataFrame(catalog)[columns]
        self.columns = columns

    def __len__(self):
        return len(self.catalog)

    def loglike(self, rate_fn, n_expected):
        """Computes log-likelihood of data under poisson process model

        rate_fn returns a rate as a function of self.columns 
        (must be able to take all of self.columns as keywords)

        n_expected is predicted total number
        """
        rate_densities = rate_fn(**{c: self.catalog[c] for c in self.columns})

        return -n_expected + sum(np.log(rate_densities))

     

import pandas as pd

from cobaya.likelihood import Likelihood

from .poisson_data import PoissonData


class PoissonLikelihood(Likelihood):
    class_options = {"name": "Poisson", "catalog_path": None, "columns": None}

    def initialize(self):
        catalog = self._get_catalog()
        if self.columns is None:
            self.columns = catalog.columns
        self.data = PoissonData(self.name, catalog, self.columns)

    def _get_catalog(self):
        return pd.read_csv(self.catalog_path)

    def _get_rate_fn(self, **kwargs):
        """Returns a callable rate function that takes each of 'columns' as kwargs.
        """
        raise NotImplementedError

    def _get_n_expected(self, **kwargs):
        """Computes and returns the integral of the rate function 
        """
        raise NotImplementedError

    def logp(self, **params_values):
        rate_fn = self._get_rate_fn(**params_values)
        n_expected = self._get_n_expected(**params_values)
        return self.data.loglike(rate_fn, n_expected)

import numpy as np

from solike.tests.test_mflike import cosmo_params, nuisance_params


def test_multi():
    info = {
        "likelihood": {
            "solike.gaussian.MultiGaussianLikelihood": {
                "components": ["solike.mflike.MFLike", "solike.SimulatedLensingLikelihood"],
                "options": [{"sim_id": 0}, {"sim_number": 1}],
                "stop_at_error": True,
            }
        },
        "theory": {"camb": {"extra_args": {"lens_potential_accuracy": 1}}},
        "params": {**cosmo_params, **nuisance_params},
    }

    info1 = {
        "likelihood": {"solike.mflike.MFLike": {"sim_id": 0},},
        "theory": {"camb": {"extra_args": {"lens_potential_accuracy": 1}}},
        "params": {**cosmo_params, **nuisance_params},
    }

    info2 = {
        "likelihood": {"solike.SimulatedLensingLikelihood": {"sim_number": 1},},
        "theory": {"camb": {"extra_args": {"lens_potential_accuracy": 1}}},
        "params": {**cosmo_params},
    }

    from cobaya.model import get_model

    model = get_model(info)
    model1 = get_model(info1)
    model2 = get_model(info2)

    logp = np.float(model.loglikes({}, cached=False)[0])
    assert np.isfinite(logp)

    assert np.isclose(
        logp, float(model1.loglikes({}, cached=False)[0] + model2.loglikes({}, cached=False)[0])
    )

class MultiGaussianLikelihood(GaussianLikelihood):
    class_options = {"components": None, "options": None, "cross_cov_path": None}

    def initialize(self):
        self.likelihoods = [get_likelihood(*kv) for kv in zip(self.components, self.options)]

        self.cross_cov = CrossCov.load(self.cross_cov_path)

        # Why doesn't merge_params_info() work here?
        all_params = [l.params for l in self.likelihoods if hasattr(l, "params")]
        if all_params:
            self.params = merge_info(*all_params)

        data_list = [l.data for l in self.likelihoods]
        self.data = MultiGaussianData(data_list, self.cross_cov)

    def initialize_with_provider(self, provider):
        for like in self.likelihoods:
            like.initialize_with_provider(provider)
        super().initialize_with_provider(provider)

    def _get_theory(self, **kwargs):
        return np.concatenate([l._get_theory(**kwargs) for l in self.likelihoods])

    def get_requirements(self):
        reqs = {}
        for like in self.likelihoods:
            reqs = recursive_update(reqs, like.get_requirements())
        return reqs