<commit_before>import pymc3 as pm
import numpy as np


__all__ = ['ObservableModel', 'ExpansionParameterModel']

def bn_cov(cov, Lambda_b, n, cov_dim):
    ones = np.ones((cov_dim, cov_dim))
    return pm.math.exp(pm.math.log(cov) - 2 * n * ones * pm.math.log(Lambda_b))


class ObservableModel(pm.Model):
    """A statistical model for the convergence pattern of observables in an EFT.

    Parameters
    ----------
    data   :
    inputs :
    name   :
    model  :

    """

    def __init__(self, coeff_data, inputs, index_list,
                 corr_length_obs_list=None,
                 expansion_parameter=None,
                 name='', model=None, **kwargs):
        # name will be prefix for all variables here
        # if no name specified for model there will be no prefix
        super(ObservableModel, self).__init__(name, model)

        noise = 1e-10
        cbar_lower = 1e-3
        cbar_upper = 1e3

        # The number of coefficient (functions)
        num_coeffs = len(index_list)

        # Domain scaling factors (account for scales other than [0, 1])
        domain_scales = np.amax(inputs, axis=0) - np.amin(inputs, axis=0)

        # Observations shape
        input_dim = len(inputs[0])
        # print(input_dim)

        assert len(coeff_data) == len(index_list), \
            "Indices must match number of coefficients"

        # Create a log-uniform distribution for cbar
        logcbar_lower = pm.math.log(cbar_lower)
        logcbar_upper = pm.math.log(cbar_upper)
        logcbar = pm.Uniform(
            'logcbar', logcbar_lower, logcbar_upper,
            # transform=None,
            # testval=pm.math.log(cbar_true)
            )
        self.cbar = pm.Deterministic('cbar', pm.math.exp(logcbar))

        cov = self.cbar**2

        # Gamma prior for length scales
        # Domain scaling: If X ~ Gamma(a, b), then cX ~ Gamma(a, b/c)
        # lengthscales = pm.Gamma("lengthscales",
        #                         alpha=2, beta=1/domain_scales[0],
        #                         observed=corr_length_obs_list,
        #                         shape=input_dim)

        # Let the covariance factor by dimension, each with ExpQuad
        for dim in range(input_dim):
            if corr_length_obs_list is None:
                ls_obs = None
            else:
                ls_obs = corr_length_obs_list[dim]
            # Gamma prior for length scales
            # Domain scaling: If X ~ Gamma(a, b), then cX ~ Gamma(a, b/c)
            # ls = pm.Gamma("length{}".format(dim),
            #               alpha=5, beta=5/domain_scales[dim]/3,
            #               observed=ls_obs
            #               )
            # print("ds", domain_scales[dim])
            # ls = pm.Normal("length{}".format(dim),
            #                mu=.5*domain_scales[dim],
            #                sd=.2*domain_scales[dim],
            #                # observed=ls_obs
            #                )
            ls = pm.Normal("length{}".format(dim),
                           mu=300,
                           sd=100,
                           # observed=ls_obs
                           )

            cov = cov * pm.gp.cov.ExpQuad(
               input_dim=input_dim, ls=ls, active_dims=[dim]
               )

        for n, cn in zip(index_list, coeff_data):
            scale = 1
            if expansion_parameter is not None:
                scale = (expansion_parameter.scale())**n

            scaled_cov = 1 / scale**2 * cov
            # ckvecset.append(
            #     pm.MvNormal(
            #         'c{}'.format(n),
            #         mu=0,
            #         cov=scaled_cov + cov_noise,
            #         shape=obs_shape,
            #         observed=cn
            #         )
            #     )
            gp = pm.gp.Marginal(cov_func=scaled_cov)
            cnobs = gp.marginal_likelihood(
                'c{}obs'.format(n),
                X=inputs,
                y=cn,
                noise=noise
                )

            pm.Deterministic("c{}".format(n), scale * cnobs)


class ExpansionParameterModel(pm.Model):
    """
    """

    def __init__(self, breakdown_eval, name='', model=None,
                 breakdown_lower=1, breakdown_upper=10000,
                 **kwargs):
        super(ExpansionParameterModel, self).__init__(name, model)

        # Define hard breakdown scale \Lambda_b
        logbreakdown_lower = pm.math.log(breakdown_lower)
        logbreakdown_upper = pm.math.log(breakdown_upper)
        self.logbreakdown = pm.Uniform(
            'logbreakdown', logbreakdown_lower, logbreakdown_upper
            )
        self.breakdown = pm.Deterministic(
            'breakdown',
            pm.math.exp(self.logbreakdown)
            )

        self.breakdown_eval = breakdown_eval

    def scale(self):
        return self.breakdown/self.breakdown_eval

# Obs = ObservableModel([], [], name="test")

# print(Obs['v1'], Obs.cbar_lower)

# print(Obs['v2'])

# print(Obs.model)

# x1, x2 = np.meshgrid(np.linspace(0,1,10), np.arange(1,4))
# X2 = np.concatenate((x1.reshape((30,1)), x2.reshape((30,1))), axis=1)

# print(x1, x2)
# print(X2)<commit_msg>Added complete control over priors through kwarg dictionaries.<commit_after>import pymc3 as pm
import numpy as np
import theano.tensor as tt
import collections


__all__ = ['ObservableModel', 'ExpansionParameterModel']

# def bn_cov(cov, Lambda_b, n, cov_dim):
#     ones = np.ones((cov_dim, cov_dim))
#     return pm.math.exp(pm.math.log(cov) - 2 * n * ones * pm.math.log(Lambda_b))


def update(d, u):
    """Update a nested dictionary d using u without overwriting subdicts"""
    for k, v in u.items():
        if isinstance(v, collections.Mapping):
            r = update(d.get(k, {}), v)
            d[k] = r
        else:
            d[k] = u[k]
    return d


class ObservableModel(pm.Model):
    """A statistical model for the convergence pattern of observables in an EFT.

    Parameters
    ----------
    data   :
    inputs :
    name   :
    model  :

    """

    # Could use defaults like "dist": pm.LogNormal
    ls_default_kwargs = {"name": "ls", "dist": pm.Lognormal}
    sd_default_kwargs = {"name": "sd", "dist": pm.Lognormal}
    cov_default_kwargs = {
        "cov": pm.gp.cov.ExpQuad,
        "noise": 1e-10,
        "sd": sd_default_kwargs,
        "ls": ls_default_kwargs
        }

    def __init__(self, coeff_data, inputs, index_list,
                 cov_kwargs={},
                 expansion_parameter=None,
                 name='', model=None, **kwargs):
        # name will be prefix for all variables here
        # if no name specified for model there will be no prefix
        super(ObservableModel, self).__init__(name, model)

        # Overwrite any default parameters
        self.cov_kwargs = self.cov_default_kwargs
        update(self.cov_kwargs, cov_kwargs)

        # The number of coefficient (functions)
        num_coeffs = len(index_list)

        # Observations shape
        self.input_dim = len(inputs[0])
        self.cov_kwargs["input_dim"] = self.input_dim

        assert len(coeff_data) == len(index_list), \
            "Indices must match number of coefficients"

        if "custom" in self.cov_kwargs:
            self.noise = self.cov_kwargs['noise']
            cov = self.cov_kwargs["custom"]
        else:
            cov = self.setup_covariance(**self.cov_kwargs)

        for n, cn in zip(index_list, coeff_data):
            scale = 1
            if expansion_parameter is not None:
                scale = (expansion_parameter.scale)**n

            scaled_cov = 1 / scale**2 * cov

            gp = pm.gp.Marginal(cov_func=scaled_cov)
            cnobs = gp.marginal_likelihood(
                'c{}obs'.format(n),
                X=inputs,
                y=cn,
                noise=self.noise
                )

            pm.Deterministic("c{}".format(n), scale * cnobs)

    def setup_hyperparameter(self, **kwargs):
        temp_kwargs = kwargs
        dist = temp_kwargs.pop("dist")
        return dist(**temp_kwargs)

    def setup_covariance(self, **kwargs):
        temp_kwargs = kwargs
        # Convert dicts to hyperparameters, i.e., pymc3 objects
        for k, v in temp_kwargs.items():
            if isinstance(v, dict):
                temp_kwargs[k] = self.setup_hyperparameter(**v)
        # Save for later: Put noise directly into gp
        self.noise = temp_kwargs.pop("noise")
        # Setup covariance function
        cov_func = temp_kwargs.pop("cov")
        self.sd = temp_kwargs.pop("sd")
        cov = self.sd**2 * cov_func(**temp_kwargs)
        return cov


class ExpansionParameterModel(pm.Model):
    """
    """

    breakdown_default_kwargs = {"name": "breakdown", "dist": pm.Lognormal}

    def __init__(self, breakdown_eval, breakdown_kwargs={}, name='',
                 model=None, **kwargs):
        super(ExpansionParameterModel, self).__init__(name, model)

        # Setup kwargs for breakdown scale random variable
        self.breakdown_kwargs = self.breakdown_default_kwargs
        # Issues can occur if sampling doesn't begin in a reasonable region
        self.breakdown_kwargs["testval"] = breakdown_eval
        # Override defaults
        self.breakdown_kwargs.update(breakdown_kwargs)

        # Setup random variable for the breakdown scale
        if "custom" in self.breakdown_kwargs:
            self.breakdown = self.breakdown_kwargs["custom"]
        else:
            self.breakdown = self.setup_hyperparameter(**self.breakdown_kwargs)

        self.breakdown_eval = breakdown_eval

        # The scaling factor for coefficients: c_n ~ scale^n
        self.scale = self.breakdown/self.breakdown_eval

    def setup_hyperparameter(self, **kwargs):
        temp_kwargs = kwargs
        dist = temp_kwargs.pop("dist")
        return dist(**temp_kwargs)
{'fraction_non_alphanumeric': 0.10950162074554295, 'fraction_numerical': 0.007394651539708265, 'mean_word_length': 1.890222482435597, 'pattern_counts': {'":': 11, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3773971', 'n_tokens_mistral': 3043, 'n_tokens_neox': 2889, 'n_words': 660}