import json import os import platform import socket from collections import namedtuple import numpy as np import pytest from brian2 import _cache_dirs_and_extensions, clear_cache, prefs from brian2.codegen.codeobject import CodeObject from brian2.codegen.cpp_prefs import compiler_supports_c99, get_compiler_and_args from brian2.codegen.generators.cython_generator import CythonNodeRenderer from brian2.codegen.optimisation import optimise_statements from brian2.codegen.runtime.cython_rt import CythonCodeObject from brian2.codegen.statements import Statement from brian2.codegen.translation import ( analyse_identifiers, get_identifiers_recursively, make_statements, parse_statement, ) from brian2.core.functions import DEFAULT_CONSTANTS, DEFAULT_FUNCTIONS, Function from brian2.core.variables import ArrayVariable, Constant, Subexpression, Variable from brian2.devices.device import auto_target, device from brian2.parsing.rendering import CPPNodeRenderer, NodeRenderer, NumpyNodeRenderer from brian2.parsing.sympytools import str_to_sympy, sympy_to_str from brian2.units import ms, second from brian2.units.fundamentalunits import Unit from brian2.utils.logger import catch_logs FakeGroup = namedtuple("FakeGroup", ["variables"]) @pytest.mark.codegen_independent def test_auto_target(): # very basic test that the "auto" codegen target is useable assert issubclass(auto_target(), CodeObject) @pytest.mark.codegen_independent 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"} @pytest.mark.codegen_independent 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"} @pytest.mark.codegen_independent 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) @pytest.mark.codegen_independent 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) @pytest.mark.codegen_independent 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" @pytest.mark.codegen_independent 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) @pytest.mark.codegen_independent 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" @pytest.mark.codegen_independent 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)" @pytest.mark.codegen_independent 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 should not be simplified to 0! Statement("v1", "=", "rand() - rand()", "", float), Statement("v1", "=", "3*rand() - 3*rand()", "", float), Statement("v1", "=", "3*rand() - ((1+2)*rand())", "", float), # This should not pull out rand()*N Statement("v1", "=", "s1*rand()*N", "", float), Statement("v1", "=", "s2*rand()*N", "", float), # This is not important mathematically, but it would change the numbers # that are generated Statement("v1", "=", "0*rand()*N", "", float), Statement("v1", "=", "0/rand()*N", "", float), ] scalar, vector = optimise_statements([], statements, variables) for vs in vector[:3]: assert ( vs.expr.count("rand()") == 2 ), f"Expression should still contain two rand() calls, but got {str(vs)}" for vs in vector[3:]: assert ( vs.expr.count("rand()") == 1 ), f"Expression should still contain a rand() call, but got {str(vs)}" @pytest.mark.codegen_independent 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", "", float), Statement("v1", "=", "N*v1 - N*v1", "", float), Statement("v1", "=", "v1*N * 0", "", float), Statement("v1", "=", "v1 * 0", "", float), Statement("v1", "=", "v1 * 0.0", "", float), Statement("v1", "=", "0.0 / (v1*N)", "", float), # Should be simplified to 0 Statement("i1", "=", "i1*N * 0", "", int), Statement("i1", "=", "0 * i1", "", int), Statement("i1", "=", "0 * i1*N", "", int), Statement("i1", "=", "i1 * 0", "", int), # Should be simplified to v1*N Statement("v2", "=", "0 + v1*N", "", float), Statement("v2", "=", "v1*N + 0.0", "", float), Statement("v2", "=", "v1*N - 0", "", float), Statement("v2", "=", "v1*N - 0.0", "", float), Statement("v2", "=", "1 * v1*N", "", float), Statement("v2", "=", "1.0 * v1*N", "", float), Statement("v2", "=", "v1*N / 1.0", "", float), Statement("v2", "=", "v1*N / 1", "", 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" @pytest.mark.codegen_independent 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)", "", float), Statement("v1", "=", "exp(N)*v1", "", float), Statement("v1", "=", "exp(0)*v1", "", 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" @pytest.mark.codegen_independent def test_automatic_augmented_assignments(): # We test that statements that could be rewritten as augmented assignments # are correctly rewritten (using sympy to test for symbolic equality) variables = { "x": ArrayVariable("x", owner=None, size=10, device=device), "y": ArrayVariable("y", owner=None, size=10, device=device), "z": ArrayVariable("y", owner=None, size=10, device=device), "b": ArrayVariable("b", owner=None, size=10, dtype=bool, device=device), "clip": DEFAULT_FUNCTIONS["clip"], "inf": DEFAULT_CONSTANTS["inf"], } statements = [ # examples that should be rewritten # Note that using our approach, we will never get -= or /= but always # the equivalent += or *= statements ("x = x + 1.0", "x += 1.0"), ("x = 2.0 * x", "x *= 2.0"), ("x = x - 3.0", "x += -3.0"), ("x = x/2.0", "x *= 0.5"), ("x = y + (x + 1.0)", "x += y + 1.0"), ("x = x + x", "x *= 2"), ("x = x + y + z", "x += y + z"), ("x = x + y + z", "x += y + z"), # examples that should not be rewritten ("x = 1.0/x", "x = 1.0/x"), ("x = 1.0", "x = 1.0"), ("x = 2.0*(x + 1.0)", "x = 2.0*(x + 1.0)"), ("x = clip(x + y, 0.0, inf)", "x = clip(x + y, 0.0, inf)"), ("b = b or False", "b = b or False"), ] for orig, rewritten in statements: scalar, vector = make_statements(orig, variables, np.float32) try: # we augment the assertion error with the original statement assert ( len(scalar) == 0 ), f"Did not expect any scalar statements but got {str(scalar)}" assert ( len(vector) == 1 ), f"Did expect a single statement but got {str(vector)}" statement = vector[0] expected_var, expected_op, expected_expr, _ = parse_statement(rewritten) assert ( expected_var == statement.var ), f"expected write to variable {expected_var}, not to {statement.var}" assert ( expected_op == statement.op ), f"expected operation {expected_op}, not {statement.op}" # Compare the two expressions using sympy to allow for different order etc. sympy_expected = str_to_sympy(expected_expr) sympy_actual = str_to_sympy(statement.expr) assert sympy_expected == sympy_actual, ( f"RHS expressions '{sympy_to_str(sympy_expected)}' and" f" '{sympy_to_str(sympy_actual)}' are not identical" ) except AssertionError as ex: raise AssertionError( f"Transformation for statement '{orig}' gave an unexpected result: {ex}" ) @pytest.mark.codegen_independent @pytest.mark.parametrize( "s", [ "x, y = 3", "x * y", "x = ", "x.a = 3", "x++", "x[0] = 3", "dx/dt = -v / tau", "v == 3*mV", ], ) def test_incorrect_statements(s): with pytest.raises(ValueError): parse_statement(s) def test_clear_cache(): target = prefs.codegen.target if target == "numpy": assert "numpy" not in _cache_dirs_and_extensions with pytest.raises(ValueError): clear_cache("numpy") else: assert target in _cache_dirs_and_extensions cache_dir, _ = _cache_dirs_and_extensions[target] # Create a file that should not be there fname = os.path.join(cache_dir, "some_file.py") open(fname, "w").close() # clear_cache should refuse to clear the directory with pytest.raises(IOError): clear_cache(target) os.remove(fname) @pytest.mark.skipif( platform.system() == "Windows", reason="CC and CXX variables are ignored on Windows.", ) def test_compiler_error(): # In particular on OSX with clang in a conda environment, compilation might fail. # Switching to a system gcc might help in such cases. Make sure that the error # message mentions that. old_CC = os.environ.get("CC", None) old_CXX = os.environ.get("CXX", None) os.environ.update({"CC": "non-existing-compiler", "CXX": "non-existing-compiler++"}) try: with catch_logs() as l: assert not CythonCodeObject.is_available() assert len(l) > 0 # There are additional warnings about compiler flags last_warning = l[-1] assert last_warning[1].endswith(".failed_compile_test") assert "CC" in last_warning[2] and "CXX" in last_warning[2] finally: if old_CC: os.environ["CC"] = old_CC else: del os.environ["CC"] if old_CXX: os.environ["CXX"] = old_CXX else: del os.environ["CXX"] def test_compiler_c99(): # On a user's computer, we do not know whether the compiler actually # has C99 support, so we just check whether the test does not raise an # error # The compiler check previously created spurious '-.o' files (see #1348) if os.path.exists("-.o"): os.remove("-.o") c99_support = compiler_supports_c99() assert not os.path.exists("-.o") # On our Azure test server we know that the compilers support C99 if os.environ.get("AGENT_OS", ""): assert c99_support def test_cpp_flags_support(): from distutils.ccompiler import get_default_compiler from brian2.codegen.cpp_prefs import _compiler_flag_compatibility _compiler_flag_compatibility.clear() # make sure cache is empty compiler = get_default_compiler() if compiler == "msvc": pytest.skip("No flag support check for msvc") old_prefs = prefs["codegen.cpp.extra_compile_args"] # Should always be supported prefs["codegen.cpp.extra_compile_args"] = ["-w"] _, compile_args = get_compiler_and_args() assert compile_args == prefs["codegen.cpp.extra_compile_args"] # Should never be supported and raise a warning prefs["codegen.cpp.extra_compile_args"] = ["-invalidxyz"] with catch_logs() as l: _, compile_args = get_compiler_and_args() assert len(l) == 1 and l[0][0] == "WARNING" assert compile_args == [] prefs["codegen.cpp.extra_compile_args"] = old_prefs @pytest.mark.skipif( platform.system() != "Windows", reason="MSVC flags are only relevant on Windows" ) @pytest.mark.skipif( prefs["codegen.target"] == "numpy", reason="Test only relevant for compiled code" ) def test_msvc_flags(): # Very basic test that flags are stored to disk import brian2.codegen.cpp_prefs as cpp_prefs user_dir = os.path.join(os.path.expanduser("~"), ".brian") flag_file = os.path.join(user_dir, "cpu_flags.txt") assert len(cpp_prefs.msvc_arch_flag) assert os.path.exists(flag_file) with open(flag_file, encoding="utf-8") as f: previously_stored_flags = json.load(f) hostname = socket.gethostname() assert hostname in previously_stored_flags assert len(previously_stored_flags[hostname]) @pytest.mark.codegen_independent @pytest.mark.parametrize( "renderer", [ NodeRenderer(), NumpyNodeRenderer(), CythonNodeRenderer(), CPPNodeRenderer(), ], ) def test_number_rendering(renderer): import ast for number in [0.5, np.float32(0.5), np.float64(0.5)]: # In numpy 2.0, repr(np.float64(0.5)) is 'np.float64(0.5)' node = ast.Constant(value=number) assert renderer.render_node(node) == "0.5" if __name__ == "__main__": test_auto_target() test_analyse_identifiers() test_get_identifiers_recursively() test_write_to_subexpression() test_repeated_subexpressions() test_nested_subexpressions() test_apply_loop_invariant_optimisation() test_apply_loop_invariant_optimisation_integer() test_apply_loop_invariant_optimisation_boolean() test_apply_loop_invariant_optimisation_no_optimisation() test_apply_loop_invariant_optimisation_simplification() test_apply_loop_invariant_optimisation_constant_evaluation() test_automatic_augmented_assignments() test_clear_cache() test_msvc_flags()