<commit_before>import itertools
from checks import *
from pymc import *
from numpy import array, inf
import numpy

from scipy import integrate
from knownfailure import *


class Domain(object):
    def __init__(self, vals, dtype = None, edges=None, shape = None):
        avals = array(vals)

        if edges is None:
            edges = array(vals[0]), array(vals[-1])
            vals = vals[1:-1]
        if shape is None:
            shape = avals[0].shape


        self.vals = vals
        self.shape = shape

        self.lower, self.upper = edges
        self.dtype = avals.dtype

    def __neg__(self):
        return Domain([-v for v in self.vals], self.dtype, (-self.lower, -self.upper), self.shape)

def product(domains):
    return itertools.product(*[d.vals for d in domains])

R = Domain([-inf, -2.1, -1, -.01, .0, .01, 1, 2.1, inf])
Rplus = Domain([0, .01, .1, .9, .99, 1, 1.5, 2, 100, inf])
Rplusbig = Domain([0, .5, .9, .99, 1, 1.5, 2, 20, inf])
Unit = Domain([0, .001, .1, .5, .75, .99, 1])

Runif = Domain([-1, -.4, 0, .4, 1])
Rdunif = Domain([-10, 0, 10.])
Rplusunif = Domain([0, .5, inf])
Rplusdunif = Domain([2, 10, 100], 'int64')

I = Domain([-1000, -3, -2, -1, 0, 1, 2, 3, 1000], 'int64')

NatSmall = Domain([0, 3, 4, 5, 1000], 'int64')
Nat = Domain([0, 1, 2, 3, 2000], 'int64')
NatBig = Domain([0, 1, 2, 3, 5000, 50000], 'int64')

Bool = Domain([0, 0, 1, 1], 'int64')


Vec2small = Domain([
    [.1, 0.0],
    [-2.3, .1],
    [-2.3, 1.5],
    ],
    edges = ([ -25, -25], [25, 25]))

Vec3small = Domain([
    [.1, 0.0, 0],
    [-2.3, .1,1],
    [-2.3, 1.5, 2],
    ],
    edges = ([ -12, -12, -12], [12, 12,12]))

PdMatrix2 = Domain([
    np.eye(2),
    [[.5, .05],
     [.05, 4.5]]
    ],
    edges = (None,None))

PdMatrix3 = Domain([
    np.eye(3),
    [[.5, .1,0],
     [.1, 1, 0],
     [0, 0, 2.5]]
    ],
    edges = (None,None))




def test_unif():
    checkd(Uniform, Runif, {'lower': -Rplusunif, 'upper': Rplusunif})


def test_discrete_unif():
    checkd(DiscreteUniform, Rdunif,
           {'lower': -Rplusdunif, 'upper': Rplusdunif})


def test_flat():
    checkd(Flat, Runif, {}, checks = (check_dlogp,))


def test_normal():
    checkd(Normal, R, {'mu': R, 'tau': Rplus})


def test_beta():
    # TODO this fails with `Rplus`
    checkd(Beta, Unit, {'alpha': Rplusbig, 'beta': Rplusbig})


def test_exponential():
    checkd(Exponential, Rplus, {'lam': Rplus})


def test_geometric():
    checkd(Geometric, NatBig, {'p': Unit})


def test_negative_binomial():
    checkd(NegativeBinomial, Nat, {'mu': Rplusbig, 'alpha': Rplusbig})


def test_laplace():
    checkd(Laplace, R, {'mu': R, 'b': Rplus})

def test_lognormal():
    checkd(Lognormal, Rplus, {'mu': R, 'tau': Rplus}, (check_dlogp,))

def test_t():
    checkd(T, R, {'nu': Rplus, 'mu': R, 'lam': Rplus})


def test_cauchy():
    checkd(Cauchy, R, {'alpha': R, 'beta': Rplusbig})


def test_gamma():
    checkd(Gamma, Rplus, {'alpha': Rplusbig, 'beta': Rplusbig})


def test_tpos():
    checkd(Tpos, Rplus, {'nu': Rplus, 'mu': R, 'lam': Rplus}, checks=(check_dlogp,))


def test_binomial():
    checkd(Binomial, Nat, {'n': NatSmall, 'p': Unit})


def test_betabin():
    checkd(BetaBin, Nat, {'alpha': Rplus, 'beta': Rplus, 'n': NatSmall})


def test_bernoulli():
    checkd(Bernoulli, Bool, {'p': Unit})


def test_poisson():
    checkd(Poisson, Nat, {'mu': Rplus})


def test_constantdist():
    checkd(ConstantDist, I, {'c': I})


def test_zeroinflatedpoisson():
    checkd(ZeroInflatedPoisson, Nat, {'theta': Rplus, 'z': Bool})

def test_mvnormal2():
    checkd(MvNormal, Vec2small, {'mu': R, 'tau': PdMatrix2}, checks=(check_dlogp, check_int_to_1))

def test_mvnormal3():
    checkd(MvNormal, Vec3small, {'mu': R, 'tau': PdMatrix3}, checks=(check_dlogp, check_int_to_1))


def test_densitydist():
    def logp(x):
        return -log(2 * .5) - abs(x - .5) / .5

    checkd(DensityDist, R, {}, extra_args={'logp': logp})


def test_addpotential():
    with Model() as model:
        x = Normal('x', 1, 1)
        model.AddPotential(-x ** 2)

        check_dlogp(model, x, R, [])


<commit_msg>Disabled mvnormal3 test as it takes >10 mins.<commit_after>import unittest

import itertools
from checks import *
from pymc import *
from numpy import array, inf
import numpy

from scipy import integrate
from knownfailure import *


class Domain(object):
    def __init__(self, vals, dtype = None, edges=None, shape = None):
        avals = array(vals)

        if edges is None:
            edges = array(vals[0]), array(vals[-1])
            vals = vals[1:-1]
        if shape is None:
            shape = avals[0].shape


        self.vals = vals
        self.shape = shape

        self.lower, self.upper = edges
        self.dtype = avals.dtype

    def __neg__(self):
        return Domain([-v for v in self.vals], self.dtype, (-self.lower, -self.upper), self.shape)

def product(domains):
    return itertools.product(*[d.vals for d in domains])

R = Domain([-inf, -2.1, -1, -.01, .0, .01, 1, 2.1, inf])
Rplus = Domain([0, .01, .1, .9, .99, 1, 1.5, 2, 100, inf])
Rplusbig = Domain([0, .5, .9, .99, 1, 1.5, 2, 20, inf])
Unit = Domain([0, .001, .1, .5, .75, .99, 1])

Runif = Domain([-1, -.4, 0, .4, 1])
Rdunif = Domain([-10, 0, 10.])
Rplusunif = Domain([0, .5, inf])
Rplusdunif = Domain([2, 10, 100], 'int64')

I = Domain([-1000, -3, -2, -1, 0, 1, 2, 3, 1000], 'int64')

NatSmall = Domain([0, 3, 4, 5, 1000], 'int64')
Nat = Domain([0, 1, 2, 3, 2000], 'int64')
NatBig = Domain([0, 1, 2, 3, 5000, 50000], 'int64')

Bool = Domain([0, 0, 1, 1], 'int64')


Vec2small = Domain([
    [.1, 0.0],
    [-2.3, .1],
    [-2.3, 1.5],
    ],
    edges = ([ -25, -25], [25, 25]))

Vec3small = Domain([
    [.1, 0.0, 0],
    [-2.3, .1,1],
    [-2.3, 1.5, 2],
    ],
    edges = ([ -12, -12, -12], [12, 12,12]))

PdMatrix2 = Domain([
    np.eye(2),
    [[.5, .05],
     [.05, 4.5]]
    ],
    edges = (None,None))

PdMatrix3 = Domain([
    np.eye(3),
    [[.5, .1,0],
     [.1, 1, 0],
     [0, 0, 2.5]]
    ],
    edges = (None,None))




def test_unif():
    checkd(Uniform, Runif, {'lower': -Rplusunif, 'upper': Rplusunif})


def test_discrete_unif():
    checkd(DiscreteUniform, Rdunif,
           {'lower': -Rplusdunif, 'upper': Rplusdunif})


def test_flat():
    checkd(Flat, Runif, {}, checks = (check_dlogp,))


def test_normal():
    checkd(Normal, R, {'mu': R, 'tau': Rplus})


def test_beta():
    # TODO this fails with `Rplus`
    checkd(Beta, Unit, {'alpha': Rplusbig, 'beta': Rplusbig})


def test_exponential():
    checkd(Exponential, Rplus, {'lam': Rplus})


def test_geometric():
    checkd(Geometric, NatBig, {'p': Unit})


def test_negative_binomial():
    checkd(NegativeBinomial, Nat, {'mu': Rplusbig, 'alpha': Rplusbig})


def test_laplace():
    checkd(Laplace, R, {'mu': R, 'b': Rplus})

def test_lognormal():
    checkd(Lognormal, Rplus, {'mu': R, 'tau': Rplus}, (check_dlogp,))

def test_t():
    checkd(T, R, {'nu': Rplus, 'mu': R, 'lam': Rplus})


def test_cauchy():
    checkd(Cauchy, R, {'alpha': R, 'beta': Rplusbig})


def test_gamma():
    checkd(Gamma, Rplus, {'alpha': Rplusbig, 'beta': Rplusbig})


def test_tpos():
    checkd(Tpos, Rplus, {'nu': Rplus, 'mu': R, 'lam': Rplus}, checks=(check_dlogp,))


def test_binomial():
    checkd(Binomial, Nat, {'n': NatSmall, 'p': Unit})


def test_betabin():
    checkd(BetaBin, Nat, {'alpha': Rplus, 'beta': Rplus, 'n': NatSmall})


def test_bernoulli():
    checkd(Bernoulli, Bool, {'p': Unit})


def test_poisson():
    checkd(Poisson, Nat, {'mu': Rplus})


def test_constantdist():
    checkd(ConstantDist, I, {'c': I})


def test_zeroinflatedpoisson():
    checkd(ZeroInflatedPoisson, Nat, {'theta': Rplus, 'z': Bool})

def test_mvnormal2():
    checkd(MvNormal, Vec2small, {'mu': R, 'tau': PdMatrix2}, checks=(check_dlogp, check_int_to_1))

@unittest.skip('Takes too long for travis.')
def test_mvnormal3():
    checkd(MvNormal, Vec3small, {'mu': R, 'tau': PdMatrix3}, checks=(check_dlogp, check_int_to_1))


def test_densitydist():
    def logp(x):
        return -log(2 * .5) - abs(x - .5) / .5

    checkd(DensityDist, R, {}, extra_args={'logp': logp})


def test_addpotential():
    with Model() as model:
        x = Normal('x', 1, 1)
        model.AddPotential(-x ** 2)

        check_dlogp(model, x, R, [])



{'fraction_non_alphanumeric': 0.20263094376055998, 'fraction_numerical': 0.051291334781559256, 'mean_word_length': 3.1167411823149527, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4647251', 'n_tokens_mistral': 3681, 'n_tokens_neox': 3317, 'n_words': 886}