def test_apply_loop_invariant_optimisation_constant_evaluation():
variables = {
'v1': Variable('v1', scalar=False),
'v2': Variable('v2', scalar=False),
'i1': Variable('i1', scalar=False, dtype=int),
'N': Constant('N', 10),
's1': Variable('s1', scalar=True, dtype=float),
's2': Variable('s2', scalar=True, dtype=float),
'exp': DEFAULT_FUNCTIONS['exp']
}
statements = [
Statement('v1', '=', 'v1 * (1 + 2 + 3)', '', np.float),
Statement('v1', '=', 'exp(N)*v1', '', np.float),
Statement('v1', '=', 'exp(0)*v1', '', np.float),
]
scalar, vector = optimise_statements([], statements, variables)
# exp(N) should be pulled out of the vector statements, the rest should be
# evaluated in place
assert len(scalar) == 1
assert scalar[0].expr == 'exp(N)'
assert len(vector) == 3
expr = vector[0].expr.replace(' ', '')
assert expr == '_lio_1*v1' or 'v1*_lio_1'
expr = vector[1].expr.replace(' ', '')
assert expr == '6.0*v1' or 'v1*6.0'
assert vector[2].expr == 'v1'
def test_apply_loop_invariant_optimisation_no_optimisation():
variables = {
'v1': Variable('v1', scalar=False),
'v2': Variable('v2', scalar=False),
'N': Constant('N', 10),
's1': Variable('s1', scalar=True, dtype=float),
's2': Variable('s2', scalar=True, dtype=float),
'rand': DEFAULT_FUNCTIONS['rand']
}
statements = [
# This hould not be simplified to 0!
Statement('v1', '=', 'rand() - rand()', '', np.float),
Statement('v1', '=', '3*rand() - 3*rand()', '', np.float),
Statement('v1', '=', '3*rand() - ((1+2)*rand())', '', np.float),
# This should not pull out rand()*N
Statement('v1', '=', 's1*rand()*N', '', np.float),
Statement('v1', '=', 's2*rand()*N', '', np.float),
# This is not important mathematically, but it would change the numbers
# that are generated
Statement('v1', '=', '0*rand()*N', '', np.float),
Statement('v1', '=', '0/rand()*N', '', np.float)
]
scalar, vector = optimise_statements([], statements, variables)
for vs in vector[:3]:
assert vs.expr.count(
'rand()'
) == 2, 'Expression should still contain two rand() calls, but got ' + str(
vs)
for vs in vector[3:]:
assert vs.expr.count(
'rand()'
) == 1, 'Expression should still contain a rand() call, but got ' + str(
vs)
def test_repeated_subexpressions():
variables = {
'a':
Subexpression(name='a',
dtype=np.float32,
owner=FakeGroup(variables={}),
device=None,
expr='2*z'),
'x':
Variable(name='x'),
'y':
Variable(name='y'),
'z':
Variable(name='z')
}
# subexpression a (referring to z) is used twice, but can be reused the
# second time (no change to z)
code = '''
x = a
y = a
'''
scalar_stmts, vector_stmts = make_statements(code, variables, np.float32)
assert len(scalar_stmts) == 0
assert [stmt.var for stmt in vector_stmts] == ['a', 'x', 'y']
assert vector_stmts[0].constant
code = '''
x = a
z *= 2
'''
scalar_stmts, vector_stmts = make_statements(code, variables, np.float32)
assert len(scalar_stmts) == 0
assert [stmt.var for stmt in vector_stmts] == ['a', 'x', 'z']
# Note that we currently do not mark the subexpression as constant in this
# case, because its use after the "z *=2" line would actually redefine it.
# Our algorithm is currently not smart enough to detect that it is actually
# not used afterwards
# a refers to z, therefore we have to redefine a after z changed, and a
# cannot be constant
code = '''
x = a
z *= 2
y = a
'''
scalar_stmts, vector_stmts = make_statements(code, variables, np.float32)
assert len(scalar_stmts) == 0
assert [stmt.var for stmt in vector_stmts] == ['a', 'x', 'z', 'a', 'y']
assert not any(stmt.constant for stmt in vector_stmts)
def test_apply_loop_invariant_optimisation_integer():
variables = {
'v': Variable('v', scalar=False),
'N': Constant('N', 10),
'b': Variable('b', scalar=True, dtype=int),
'c': Variable('c', scalar=True, dtype=int),
'd': Variable('d', scalar=True, dtype=int),
'y': Variable('y', scalar=True, dtype=float),
'z': Variable('z', scalar=True, dtype=float),
'w': Variable('w', scalar=True, dtype=float),
}
statements = [
Statement('v', '=', 'v % (2*3*N)', '', np.float32),
# integer version doesn't get rewritten but float version does
Statement('a', ':=', 'b//(c//d)', '', int),
Statement('x', ':=', 'y/(z/w)', '', float),
]
scalar, vector = optimise_statements([], statements, variables)
assert len(scalar) == 3
assert np.issubdtype(scalar[0].dtype, np.signedinteger)
assert scalar[0].var == '_lio_1'
expr = scalar[0].expr.replace(' ', '')
assert expr == '6*N' or expr == 'N*6'
assert np.issubdtype(scalar[1].dtype, np.signedinteger)
assert scalar[1].var == '_lio_2'
expr = scalar[1].expr.replace(' ', '')
assert expr == 'b//(c//d)'
assert np.issubdtype(scalar[2].dtype, np.floating)
assert scalar[2].var == '_lio_3'
expr = scalar[2].expr.replace(' ', '')
assert expr == '(y*w)/z' or expr == '(w*y)/z'
def test_analyse_identifiers():
'''
Test that the analyse_identifiers function works on a simple clear example.
'''
code = '''
a = b+c
d = e+f
'''
known = {
'b': Variable(name='b'),
'c': Variable(name='c'),
'd': Variable(name='d'),
'g': Variable(name='g')
}
defined, used_known, dependent = analyse_identifiers(code, known)
assert 'a' in defined # There might be an additional constant added by the
# loop-invariant optimisation
assert used_known == {'b', 'c', 'd'}
assert dependent == {'e', 'f'}
def test_apply_loop_invariant_optimisation():
variables = {
'v': Variable('v', scalar=False),
'w': Variable('w', scalar=False),
'dt': Constant('dt', dimensions=second.dim, value=0.1 * ms),
'tau': Constant('tau', dimensions=second.dim, value=10 * ms),
'exp': DEFAULT_FUNCTIONS['exp']
}
statements = [
Statement('v', '=', 'dt*w*exp(-dt/tau)/tau + v*exp(-dt/tau)', '',
np.float32),
Statement('w', '=', 'w*exp(-dt/tau)', '', np.float32)
]
scalar, vector = optimise_statements([], statements, variables)
# The optimisation should pull out at least exp(-dt / tau)
assert len(scalar) >= 1
assert np.issubdtype(scalar[0].dtype, np.floating)
assert scalar[0].var == '_lio_1'
assert len(vector) == 2
assert all('_lio_' in stmt.expr for stmt in vector)
def test_apply_loop_invariant_optimisation_boolean():
variables = {
'v1':
Variable('v1', scalar=False),
'v2':
Variable('v2', scalar=False),
'N':
Constant('N', 10),
'b':
Variable('b', scalar=True, dtype=bool),
'c':
Variable('c', scalar=True, dtype=bool),
'int':
DEFAULT_FUNCTIONS['int'],
'foo':
Function(lambda x: None,
arg_units=[Unit(1)],
return_unit=Unit(1),
arg_types=['boolean'],
return_type='float',
stateless=False)
}
# The calls for "foo" cannot be pulled out, since foo is marked as stateful
statements = [
Statement('v1', '=', '1.0*int(b and c)', '', np.float32),
Statement('v1', '=', '1.0*foo(b and c)', '', np.float32),
Statement('v2', '=', 'int(not b and True)', '', np.float32),
Statement('v2', '=', 'foo(not b and True)', '', np.float32)
]
scalar, vector = optimise_statements([], statements, variables)
assert len(scalar) == 4
assert scalar[0].expr == '1.0 * int(b and c)'
assert scalar[1].expr == 'b and c'
assert scalar[2].expr == 'int((not b) and True)'
assert scalar[3].expr == '(not b) and True'
assert len(vector) == 4
assert vector[0].expr == '_lio_1'
assert vector[1].expr == 'foo(_lio_2)'
assert vector[2].expr == '_lio_3'
assert vector[3].expr == 'foo(_lio_4)'
def test_nested_subexpressions():
'''
This test checks that code translation works with nested subexpressions.
'''
code = '''
x = a + b + c
c = 1
x = a + b + c
d = 1
x = a + b + c
'''
variables = {
'a':
Subexpression(name='a',
dtype=np.float32,
owner=FakeGroup(variables={}),
device=None,
expr='b*b+d'),
'b':
Subexpression(name='b',
dtype=np.float32,
owner=FakeGroup(variables={}),
device=None,
expr='c*c*c'),
'c':
Variable(name='c'),
'd':
Variable(name='d'),
}
scalar_stmts, vector_stmts = make_statements(code, variables, np.float32)
assert len(scalar_stmts) == 0
evalorder = ''.join(stmt.var for stmt in vector_stmts)
# This is the order that variables ought to be evaluated in (note that
# previously this test did not expect the last "b" evaluation, because its
# value did not change (c was not changed). We have since removed this
# subexpression caching, because it did not seem to apply in practical
# use cases)
assert evalorder == 'baxcbaxdbax'
def test_write_to_subexpression():
variables = {
'a':
Subexpression(name='a',
dtype=np.float32,
owner=FakeGroup(variables={}),
device=None,
expr='2*z'),
'z':
Variable(name='z')
}
# Writing to a subexpression is not allowed
code = 'a = z'
with pytest.raises(SyntaxError):
make_statements(code, variables, np.float32)
def test_get_identifiers_recursively():
'''
Test finding identifiers including subexpressions.
'''
variables = {
'sub1':
Subexpression(name='sub1',
dtype=np.float32,
expr='sub2 * z',
owner=FakeGroup(variables={}),
device=None),
'sub2':
Subexpression(name='sub2',
dtype=np.float32,
expr='5 + y',
owner=FakeGroup(variables={}),
device=None),
'x':
Variable(name='x')
}
identifiers = get_identifiers_recursively(['_x = sub1 + x'], variables)
assert identifiers == {'x', '_x', 'y', 'z', 'sub1', 'sub2'}
def check_units_statements(code, variables):
'''
Check the units for a series of statements. Setting a model variable has to
use the correct unit. For newly introduced temporary variables, the unit
is determined and used to check the following statements to ensure
consistency.
Parameters
----------
code : str
The statements as a (multi-line) string
variables : dict of `Variable` objects
The information about all variables used in `code` (including
`Constant` objects for external variables)
Raises
------
KeyError
In case on of the identifiers cannot be resolved.
DimensionMismatchError
If an unit mismatch occurs during the evaluation.
'''
variables = dict(variables)
# Avoid a circular import
from angela2.codegen.translation import analyse_identifiers
newly_defined, _, unknown = analyse_identifiers(code, variables)
if len(unknown):
raise AssertionError(
('Encountered unknown identifiers, this should '
'not happen at this stage. Unknown identifiers: %s' % unknown))
code = re.split(r'[;\n]', code)
for line in code:
line = line.strip()
if not len(line):
continue # skip empty lines
varname, op, expr, comment = parse_statement(line)
if op in ('+=', '-=', '*=', '/=', '%='):
# Replace statements such as "w *=2" by "w = w * 2"
expr = '{var} {op_first} {expr}'.format(var=varname,
op_first=op[0],
expr=expr)
op = '='
elif op == '=':
pass
else:
raise AssertionError('Unknown operator "%s"' % op)
expr_unit = parse_expression_dimensions(expr, variables)
if varname in variables:
expected_unit = variables[varname].dim
fail_for_dimension_mismatch(expr_unit,
expected_unit,
('The right-hand-side of code '
'statement "%s" does not have the '
'expected unit {expected}') % line,
expected=expected_unit)
elif varname in newly_defined:
# note the unit for later
variables[varname] = Variable(name=varname,
dimensions=get_dimensions(expr_unit),
scalar=False)
else:
raise AssertionError(('Variable "%s" is neither in the variables '
'dictionary nor in the list of undefined '
'variables.' % varname))
def test_determination():
'''
Test the determination of suitable state updaters.
'''
# To save some typing
apply_stateupdater = StateUpdateMethod.apply_stateupdater
eqs = Equations('dv/dt = -v / (10*ms) : 1')
# Just make sure that state updaters know about the two state variables
variables = {'v': Variable(name='v'), 'w': Variable(name='w')}
# all methods should work for these equations.
# First, specify them explicitly (using the object)
for integrator in (
linear,
euler,
exponential_euler, #TODO: Removed "independent" here due to the issue in sympy 0.7.4
rk2,
rk4,
heun,
milstein):
with catch_logs() as logs:
returned = apply_stateupdater(eqs, variables, method=integrator)
assert len(logs) == 0, 'Got %d unexpected warnings: %s' % (
len(logs), str([l[2] for l in logs]))
# Equation with multiplicative noise, only milstein and heun should work
eqs = Equations('dv/dt = -v / (10*ms) + v*xi*second**-.5: 1')
for integrator in (linear, independent, euler, exponential_euler, rk2,
rk4):
with pytest.raises(UnsupportedEquationsException):
apply_stateupdater(eqs, variables, integrator)
for integrator in (heun, milstein):
with catch_logs() as logs:
returned = apply_stateupdater(eqs, variables, method=integrator)
assert len(logs) == 0, 'Got %d unexpected warnings: %s' % (
len(logs), str([l[2] for l in logs]))
# Arbitrary functions (converting equations into abstract code) should
# always work
my_stateupdater = lambda eqs, vars, options: 'x_new = x'
with catch_logs() as logs:
returned = apply_stateupdater(eqs, variables, method=my_stateupdater)
# No warning here
assert len(logs) == 0
# Specification with names
eqs = Equations('dv/dt = -v / (10*ms) : 1')
for name, integrator in [
('exact', exact),
('linear', linear),
('euler', euler),
#('independent', independent), #TODO: Removed "independent" here due to the issue in sympy 0.7.4
('exponential_euler', exponential_euler),
('rk2', rk2),
('rk4', rk4),
('heun', heun),
('milstein', milstein)
]:
with catch_logs() as logs:
returned = apply_stateupdater(eqs, variables, method=name)
# No warning here
assert len(logs) == 0
# Now all except heun and milstein should refuse to work
eqs = Equations('dv/dt = -v / (10*ms) + v*xi*second**-.5: 1')
for name in [
'linear', 'exact', 'independent', 'euler', 'exponential_euler',
'rk2', 'rk4'
]:
with pytest.raises(UnsupportedEquationsException):
apply_stateupdater(eqs, variables, method=name)
# milstein should work
with catch_logs() as logs:
apply_stateupdater(eqs, variables, method='milstein')
assert len(logs) == 0
# heun should work
with catch_logs() as logs:
apply_stateupdater(eqs, variables, method='heun')
assert len(logs) == 0
# non-existing name
with pytest.raises(ValueError):
apply_stateupdater(eqs, variables, method='does_not_exist')
# Automatic state updater choice should return linear for linear equations,
# euler for non-linear, non-stochastic equations and equations with
# additive noise, heun for equations with multiplicative noise
# Because it is somewhat fragile, the "independent" state updater is not
# included in this list
all_methods = [
'linear', 'exact', 'exponential_euler', 'euler', 'heun', 'milstein'
]
eqs = Equations('dv/dt = -v / (10*ms) : 1')
with catch_logs(log_level=logging.INFO) as logs:
apply_stateupdater(eqs, variables, all_methods)
assert len(logs) == 1
assert ('linear' in logs[0][2]) or ('exact' in logs[0][2])
# This is conditionally linear
eqs = Equations('''dv/dt = -(v + w**2)/ (10*ms) : 1
dw/dt = -w/ (10*ms) : 1''')
with catch_logs(log_level=logging.INFO) as logs:
apply_stateupdater(eqs, variables, all_methods)
assert len(logs) == 1
assert 'exponential_euler' in logs[0][2]
# # Do not test for now
# eqs = Equations('dv/dt = sin(t) / (10*ms) : 1')
# assert apply_stateupdater(eqs, variables) is independent
eqs = Equations('dv/dt = -sqrt(v) / (10*ms) : 1')
with catch_logs(log_level=logging.INFO) as logs:
apply_stateupdater(eqs, variables, all_methods)
assert len(logs) == 1
assert "'euler'" in logs[0][2]
eqs = Equations('dv/dt = -v / (10*ms) + 0.1*second**-.5*xi: 1')
with catch_logs(log_level=logging.INFO) as logs:
apply_stateupdater(eqs, variables, all_methods)
assert len(logs) == 1
assert "'euler'" in logs[0][2]
eqs = Equations('dv/dt = -v / (10*ms) + v*0.1*second**-.5*xi: 1')
with catch_logs(log_level=logging.INFO) as logs:
apply_stateupdater(eqs, variables, all_methods)
assert len(logs) == 1
assert "'heun'" in logs[0][2]
def test_apply_loop_invariant_optimisation_simplification():
variables = {
'v1': Variable('v1', scalar=False),
'v2': Variable('v2', scalar=False),
'i1': Variable('i1', scalar=False, dtype=int),
'N': Constant('N', 10)
}
statements = [
# Should be simplified to 0.0
Statement('v1', '=', 'v1 - v1', '', np.float),
Statement('v1', '=', 'N*v1 - N*v1', '', np.float),
Statement('v1', '=', 'v1*N * 0', '', np.float),
Statement('v1', '=', 'v1 * 0', '', np.float),
Statement('v1', '=', 'v1 * 0.0', '', np.float),
Statement('v1', '=', '0.0 / (v1*N)', '', np.float),
# Should be simplified to 0
Statement('i1', '=', 'i1*N * 0', '', np.int),
Statement('i1', '=', '0 * i1', '', np.int),
Statement('i1', '=', '0 * i1*N', '', np.int),
Statement('i1', '=', 'i1 * 0', '', np.int),
# Should be simplified to v1*N
Statement('v2', '=', '0 + v1*N', '', np.float),
Statement('v2', '=', 'v1*N + 0.0', '', np.float),
Statement('v2', '=', 'v1*N - 0', '', np.float),
Statement('v2', '=', 'v1*N - 0.0', '', np.float),
Statement('v2', '=', '1 * v1*N', '', np.float),
Statement('v2', '=', '1.0 * v1*N', '', np.float),
Statement('v2', '=', 'v1*N / 1.0', '', np.float),
Statement('v2', '=', 'v1*N / 1', '', np.float),
# Should be simplified to i1
Statement('i1', '=', 'i1*1', '', int),
Statement('i1', '=', 'i1//1', '', int),
Statement('i1', '=', 'i1+0', '', int),
Statement('i1', '=', '0+i1', '', int),
Statement('i1', '=', 'i1-0', '', int),
# Should *not* be simplified (because it would change the type,
# important for integer division, for example)
Statement('v1', '=', 'i1*1.0', '', float),
Statement('v1', '=', '1.0*i1', '', float),
Statement('v1', '=', 'i1/1.0', '', float),
Statement('v1', '=', 'i1/1', '', float),
Statement('v1', '=', 'i1+0.0', '', float),
Statement('v1', '=', '0.0+i1', '', float),
Statement('v1', '=', 'i1-0.0', '', float),
## Should *not* be simplified, flooring division by 1 changes the value
Statement('v1', '=', 'v2//1.0', '', float),
Statement('i1', '=', 'i1//1.0', '', float) # changes type
]
scalar, vector = optimise_statements([], statements, variables)
assert len(scalar) == 0
for s in vector[:6]:
assert s.expr == '0.0'
for s in vector[6:10]:
assert s.expr == '0', s.expr # integer
for s in vector[10:18]:
expr = s.expr.replace(' ', '')
assert expr == 'v1*N' or expr == 'N*v1'
for s in vector[18:23]:
expr = s.expr.replace(' ', '')
assert expr == 'i1'
for s in vector[23:27]:
expr = s.expr.replace(' ', '')
assert expr == '1.0*i1' or expr == 'i1*1.0' or expr == 'i1/1.0'
for s in vector[27:30]:
expr = s.expr.replace(' ', '')
assert expr == '0.0+i1' or expr == 'i1+0.0'
for s in vector[30:31]:
expr = s.expr.replace(' ', '')
assert expr == 'v2//1.0' or expr == 'v2//1'
for s in vector[31:]:
expr = s.expr.replace(' ', '')
assert expr == 'i1//1.0'
def analyse_identifiers(code, variables, recursive=False):
'''
Analyses a code string (sequence of statements) to find all identifiers by type.
In a given code block, some variable names (identifiers) must be given as inputs to the code
block, and some are created by the code block. For example, the line::
a = b+c
This could mean to create a new variable a from b and c, or it could mean modify the existing
value of a from b or c, depending on whether a was previously known.
Parameters
----------
code : str
The code string, a sequence of statements one per line.
variables : dict of `Variable`, set of names
Specifiers for the model variables or a set of known names
recursive : bool, optional
Whether to recurse down into subexpressions (defaults to ``False``).
Returns
-------
newly_defined : set
A set of variables that are created by the code block.
used_known : set
A set of variables that are used and already known, a subset of the
``known`` parameter.
unknown : set
A set of variables which are used by the code block but not defined by
it and not previously known. Should correspond to variables in the
external namespace.
'''
if isinstance(variables, Mapping):
known = set(k for k, v in variables.items()
if not isinstance(k, AuxiliaryVariable))
else:
known = set(variables)
variables = dict(
(k, Variable(name=k, dtype=np.float64)) for k in known)
known |= STANDARD_IDENTIFIERS
scalar_stmts, vector_stmts = make_statements(code,
variables,
np.float64,
optimise=False)
stmts = scalar_stmts + vector_stmts
defined = set(stmt.var for stmt in stmts if stmt.op == ':=')
if len(stmts) == 0:
allids = set()
elif recursive:
if not isinstance(variables, Mapping):
raise TypeError('Have to specify a variables dictionary.')
allids = get_identifiers_recursively(
[stmt.expr
for stmt in stmts], variables) | {stmt.var
for stmt in stmts}
else:
allids = set.union(
*[get_identifiers(stmt.expr)
for stmt in stmts]) | {stmt.var
for stmt in stmts}
dependent = allids.difference(defined, known)
used_known = allids.intersection(known) - STANDARD_IDENTIFIERS
return defined, used_known, dependent