def test_subset_index():
desc = model_description("threecats['b'] ~ continuous + dummy")
comp = Model(
Intercept(),
Term(Variable("continuous")),
Term(Variable("dummy")),
)
comp.add_response(Response(Term(Variable("threecats", level="b"))))
assert desc == comp
def test_term():
desc = model_description("x")
comp = Model(Intercept(), Term(Variable("x")))
assert desc == comp
desc = model_description("term_name_abc")
comp = Model(Intercept(), Term(Variable("term_name_abc")))
assert desc == comp
desc = model_description("`$%!N4m3##!! NNN`")
comp = Model(Intercept(), Term(Variable("$%!N4m3##!! NNN")))
assert desc == comp
def visitLiteralExpr(self, expr):
if expr.value == 0:
return NegatedIntercept()
elif expr.value == 1:
return Intercept()
else:
return Term(Variable(expr.value))
def test_term_slash():
desc = model_description("x / y")
comp = Model(Intercept(), Term(Variable("x")),
Term(Variable("x"), Variable("y")))
assert desc == comp
with pytest.raises(TypeError):
model_description("x / 5")
desc = model_description("x / f(x)")
comp = Model(
Intercept(),
Term(Variable("x")),
Term(Variable("x"), Call(LazyCall("f", [LazyVariable("x")], {}))),
)
assert desc == comp
desc = model_description("x / y:z")
comp = Model(
Intercept(),
Term(Variable("x")),
Term(Variable("x"), Variable("y"), Variable("z")),
)
assert desc == comp
with pytest.raises(TypeError):
model_description("x / (1|g)")
desc = model_description("x / (z + y)")
comp = Model(
Intercept(),
Term(Variable("x")),
Term(Variable("x"), Variable("z")),
Term(Variable("x"), Variable("y")),
)
assert desc == comp
def test_term_remove():
desc = model_description("x - y")
comp = Model(Intercept(), Term(Variable("x")))
assert desc == comp
desc = model_description("x - 5")
comp = Model(Intercept(), Term(Variable("x")))
assert desc == comp
desc = model_description("x - f(x)")
comp = Model(Intercept(), Term(Variable("x")))
assert desc == comp
desc = model_description("x - y:z")
comp = Model(Intercept(), Term(Variable("x")))
assert desc == comp
desc = model_description("x - (1|g)")
comp = Model(Intercept(), Term(Variable("x")))
assert desc == comp
desc = model_description("x - (z + y)")
comp = Model(Intercept(), Term(Variable("x")))
assert desc == comp
def test_term_interaction():
desc = model_description("x:y")
comp = Model(Intercept(), Term(Variable("x"), Variable("y")))
assert desc == comp
with pytest.raises(TypeError):
model_description("x:5")
desc = model_description("x:f(x)")
comp = Model(
Intercept(),
Term(Variable("x"), Call(LazyCall("f", [LazyVariable("x")], {}))),
)
assert desc == comp
desc = model_description("x:y:z")
comp = Model(Intercept(), Term(Variable("x"), Variable("y"),
Variable("z")))
assert desc == comp
desc = model_description("x:y*z")
comp = Model(
Intercept(),
Term(Variable("x"), Variable("y")),
Term(Variable("z")),
Term(Variable("x"), Variable("y"), Variable("z")),
)
assert desc == comp
# Note the parenthesis, here `*` resolves earlier than `:`
desc = model_description("x:(y*z)")
comp = Model(
Intercept(),
Term(Variable("x"), Variable("y")),
Term(Variable("x"), Variable("z")),
Term(Variable("x"), Variable("y"), Variable("z")),
)
assert desc == comp
with pytest.raises(TypeError):
model_description("x:(1|g)")
desc = model_description("x:(z + y)")
comp = Model(
Intercept(),
Term(Variable("x"), Variable("z")),
Term(Variable("x"), Variable("y")),
)
assert desc == comp
def test_term_add():
desc = model_description("x + y")
comp = Model(Intercept(), Term(Variable("x")), Term(Variable("y")))
assert desc == comp
desc = model_description("x + 5")
comp = Model(Intercept(), Term(Variable("x")), Term(Variable(5)))
assert desc == comp
desc = model_description("x + f(x)")
comp = Model(
Intercept(),
Term(Variable("x")),
Term(Call(LazyCall("f", [LazyVariable("x")], {}))),
)
assert desc == comp
desc = model_description("x + y:z")
comp = Model(Intercept(), Term(Variable("x")),
Term(Variable("y"), Variable("z")))
assert desc == comp
desc = model_description("x + (1|g)")
comp = Model(Intercept(), Term(Variable("x")),
GroupSpecificTerm(Intercept(), Term(Variable("g"))))
assert desc == comp
desc = model_description("x + (z + y)")
comp = Model(Intercept(), Term(Variable("x")), Term(Variable("z")),
Term(Variable("y")))
assert desc == comp
def test_group_specific_interactions():
desc = model_description("0 + (a*b|h+g)")
comp = Model(
GroupSpecificTerm(expr=Intercept(), factor=Term(Variable("h"))),
GroupSpecificTerm(expr=Intercept(), factor=Term(Variable("g"))),
GroupSpecificTerm(expr=Term(Variable("a")),
factor=Term(Variable("h"))),
GroupSpecificTerm(expr=Term(Variable("a")),
factor=Term(Variable("g"))),
GroupSpecificTerm(expr=Term(Variable("b")),
factor=Term(Variable("h"))),
GroupSpecificTerm(expr=Term(Variable("b")),
factor=Term(Variable("g"))),
GroupSpecificTerm(expr=Term(Variable("a"), Variable("b")),
factor=Term(Variable("h"))),
GroupSpecificTerm(Term(Variable("a"), Variable("b")),
factor=Term(Variable("g"))),
)
assert desc == comp
desc = model_description("0 + (0 + a*b|h+g)")
comp = Model(
GroupSpecificTerm(expr=Term(Variable("a")),
factor=Term(Variable("h"))),
GroupSpecificTerm(expr=Term(Variable("a")),
factor=Term(Variable("g"))),
GroupSpecificTerm(expr=Term(Variable("b")),
factor=Term(Variable("h"))),
GroupSpecificTerm(expr=Term(Variable("b")),
factor=Term(Variable("g"))),
GroupSpecificTerm(expr=Term(Variable("a"), Variable("b")),
factor=Term(Variable("h"))),
GroupSpecificTerm(Term(Variable("a"), Variable("b")),
factor=Term(Variable("g"))),
)
assert desc == comp
def visitQuotedNameExpr(self, expr):
# Quoted names don't accept levels yet.
return Term(Variable(expr.expression.lexeme[1:-1]))
def visitVariableExpr(self, expr):
if expr.level:
level = expr.level.value
else:
level = None
return Term(Variable(expr.name.lexeme, level))
def visitCallExpr(self, expr):
# Delegates all the work to self._visitCallExpr, that works recursively through its args.
# It just wraps the result in a Term object.
return Term(Call(CallResolver(expr).resolve()))