def test_condition_if():
source = '''
Y ~= norm(loc=0, scale=2)
if (Y > 1):
condition(Y < 1)
else:
condition(Y > -1)
'''
with pytest.raises(Exception):
compiler = SPPL_Compiler(source)
compiler.execute_module()
compiler = SPPL_Compiler(source.replace('Y > 1', 'Y > 0'))
namespace = compiler.execute_module()
assert isclose(namespace.model.prob((-1 < namespace.Y) < 1), 1)
def test_imports():
source = '''
Y ~= bernoulli(p=.5)
Z ~= choice({str(i): Fraction(1, 5) for i in range(5)})
X = array(5)
for i in range(5):
X[i] ~= Fraction(1,2) * Y
'''
compiler = SPPL_Compiler(source)
with pytest.raises(NameError):
compiler.execute_module()
compiler = SPPL_Compiler('from fractions import Fraction\n%s' % (source,))
namespace = compiler.execute_module()
for i in range(5):
assert isclose(namespace.model.prob(namespace.Z << {str(i)}), .2)
def sppl(self, line, cell):
if not line:
sys.stderr.write('specify model name after %%sppl')
return
if line in self.programs:
del self.programs[line]
compiler = SPPL_Compiler(cell, line)
namespace = compiler.execute_module()
self.programs[line] = Model(cell, compiler, namespace)
def test_condition_simple():
source = '''
Y ~= norm(loc=0, scale=2)
condition((0 < Y) < 2)
Z ~= binom(n=10, p=.2)
condition(Z == 0)
'''
compiler = SPPL_Compiler(source)
namespace = compiler.execute_module()
assert isclose(namespace.model.prob((0 < namespace.Y) < 2), 1)
def test_constrain_simple():
source = '''
Y ~= norm(loc=0, scale=2)
Z ~= binom(n=10, p=.2)
constrain({Z: 10, Y: 0})
'''
compiler = SPPL_Compiler(source)
namespace = compiler.execute_module()
assert isclose(namespace.model.prob((0 < namespace.Y) < 2), 0)
assert isclose(namespace.model.prob((0 <= namespace.Y) < 2), 1)
assert isclose(namespace.model.prob(namespace.Z << {10}), 1)
def test_switch_enumerate():
source = '''
Y ~= choice({'0': .25, '1': .5, '2': .25})
switch (Y) cases (i,j in enumerate(['0', '1', '2'])):
Z ~= atomic(loc=i+int(j))
'''
compiler = SPPL_Compiler(source)
namespace = compiler.execute_module()
assert isclose(namespace.model.prob(namespace.Z << {0}), .25)
assert isclose(namespace.model.prob(namespace.Z << {2}), .5)
assert isclose(namespace.model.prob(namespace.Z << {4}), .25)
def test_switch_shallow():
source = '''
Y ~= choice({'0': .25, '1': .5, '2': .25})
switch (Y) cases (i in ['0', '1', '2']):
Z ~= atomic(loc=int(i))
'''
compiler = SPPL_Compiler(source)
namespace = compiler.execute_module()
assert isclose(namespace.model.prob(namespace.Z << {0}), .25)
assert isclose(namespace.model.prob(namespace.Z << {1}), .5)
assert isclose(namespace.model.prob(namespace.Z << {2}), .25)
def test_constrain_if():
source = '''
Y ~= norm(loc=0, scale=2)
if (Y > 1):
constrain({Y: 2})
else:
constrain({Y: 0})
'''
compiler = SPPL_Compiler(source)
namespace = compiler.execute_module()
assert isclose(namespace.model.prob(namespace.Y > 0), 0.3085375387259868)
assert isclose(namespace.model.prob(namespace.Y < 0), 0)
assert isclose(namespace.model.prob(namespace.Y << {0}), 1-0.3085375387259868)
def test_render_sppl():
model = get_model()
sppl_code = render_sppl(model)
compiler = SPPL_Compiler(sppl_code.getvalue())
namespace = compiler.execute_module()
(X, Y) = (namespace.X, namespace.Y)
for i in range(5):
assert allclose(model.logprob(Y << {'0'}), [
model.logprob(Y << {str(i)}),
namespace.model.logprob(Y << {str(i)})
])
for i in range(4):
assert allclose(model.logprob(X << {i}),
namespace.model.logprob(X << {i}))
def test_constant_parameter():
source = '''
parameters = [1, 2, 3]
n_array = 2
Y = array(n_array + 1)
for i in range(n_array + 1):
Y[i] = randint(low=parameters[i], high=parameters[i]+1)
'''
compiler = SPPL_Compiler(source)
namespace = compiler.execute_module()
assert isclose(namespace.model.prob(namespace.Y[0] << {1}), 1)
assert isclose(namespace.model.prob(namespace.Y[1] << {2}), 1)
assert isclose(namespace.model.prob(namespace.Y[2] << {3}), 1)
with pytest.raises(AssertionError):
SPPL_Compiler('%sZ = "foo"\n' % (source,))
def test_ifexp():
source = '''
from fractions import Fraction
Y ~= choice({str(i): Fraction(1, 4) for i in range(4)})
Z ~= (
atomic(loc=0) if (Y in {'0', '1'}) else
atomic(loc=4) if (Y == '2') else
atomic(loc=6))
'''
compiler = SPPL_Compiler(source)
assert 'IfElse' in compiler.render_module()
namespace = compiler.execute_module()
assert isclose(namespace.model.prob(namespace.Z << {0}), .5)
assert isclose(namespace.model.prob(namespace.Z << {4}), .25)
assert isclose(namespace.model.prob(namespace.Z << {6}), .25)
def model_ifelse_non_exhuastive_compiled():
compiler = SPPL_Compiler('''
Nationality ~= choice({'India': 0.5, 'USA': 0.5})
Perfect ~= choice({'True': 0.01, 'False': 0.99})
if (Nationality == 'India') & (Perfect == 'False'):
GPA ~= uniform(loc=0, scale=10)
elif (Nationality == 'India') & (Perfect == 'True'):
GPA ~= atomic(loc=10)
elif (Nationality == 'USA') & (Perfect == 'False'):
GPA ~= uniform(loc=0, scale=4)
else:
GPA ~= atomic(loc=4)
''')
namespace = compiler.execute_module()
return namespace.model
def test_switch_nested():
source = '''
Y ~= randint(low=0, high=4)
W ~= randint(low=0, high=2)
switch (Y) cases (i in range(0, 5)):
Z ~= choice({str(i): 1})
switch (W) cases (i in range(0, 2)): V ~= atomic(loc=i)
'''
compiler = SPPL_Compiler(source)
namespace = compiler.execute_module()
assert isclose(namespace.model.prob(namespace.Z << {'0'}), .25)
assert isclose(namespace.model.prob(namespace.Z << {'1'}), .25)
assert isclose(namespace.model.prob(namespace.Z << {'2'}), .25)
assert isclose(namespace.model.prob(namespace.Z << {'3'}), .25)
assert isclose(namespace.model.prob(namespace.V << {0}), .5)
assert isclose(namespace.model.prob(namespace.V << {1}), .5)
def test_compile_all_constructs():
source = '''
X = array(10)
W = array(10)
Y = randint(low=1, high=2)
Z = bernoulli(p=0.1)
E = choice({'1': 0.3, '2': 0.7})
for i in range(1,5):
W[i] = uniform(loc=0, scale=2)
X[i] = bernoulli(p=0.5)
X[0] ~= gamma(a=1)
H = (X[0]**2 + 2*X[0] + 3)**(1, 10)
# Here is a comment, with indentation on next line
X[5] ~= 0.3*atomic(loc=0) | 0.4*atomic(loc=-1) | 0.3*atomic(loc=3)
if X[5] == 0:
X[7] = bernoulli(p=0.1)
X[8] = 1 + X[3]
elif (X[5]**2 == 1):
X[7] = bernoulli(p=0.2)
X[8] = 1 + X[3]
else:
if (X[3] in {0, 1}):
X[7] = bernoulli(p=0.2)
X[8] = 1 + X[3]
else:
X[7] = bernoulli(p=0.2)
X[8] = 1 + X[3]
'''
compiler = SPPL_Compiler(source)
namespace = compiler.execute_module()
model = namespace.model
assert isclose(model.prob(namespace.X[5] << {0}), .3)
assert isclose(model.prob(namespace.X[5] << {-1}), .4)
assert isclose(model.prob(namespace.X[5] << {3}), .3)
assert isclose(model.prob(namespace.E << {'1'}), .3)
assert isclose(model.prob(namespace.E << {'2'}), .7)