def test_that_code_block_is_callable(self):
code = '''
x = 1
d_x[d_idx] += s_x[s_idx] + x
'''
cb = BasicCodeBlock(code=code)
# The code block should be callable.
res = cb()
self.assertTrue(sum(res.d_x) == 1)
# Should take arguments to update the context.
x = numpy.ones(2, dtype=float)*10
res = cb(s_x=x)
self.assertTrue(sum(res.d_x) == 11)
def test_basic_code_block(self):
code = '''
x = 1
d_x[d_idx] += s_x[s_idx] + x
'''
cb = BasicCodeBlock(code=code)
expect = ['d_idx', 'd_x', 's_idx', 's_x', 'x']
self.assert_seq_equal(cb.symbols, expect)
self.assert_seq_equal(cb.src_arrays, ['s_x'])
self.assert_seq_equal(cb.dest_arrays, ['d_x'])
ctx = cb.context
self.assertEqual(ctx.s_idx, 0)
self.assertEqual(ctx.d_idx, 0)
x = numpy.zeros(2, dtype=float)
self.assertTrue(numpy.alltrue(ctx.d_x == x))
self.assertTrue(numpy.alltrue(ctx.s_x == x))
def test_order_of_precomputed(self):
try:
pre_comp = Group.pre_comp
pre_comp.AIJ = BasicCodeBlock(code=dedent("""
AIJ[0] = XIJ[0]/RIJ
AIJ[1] = XIJ[1]/RIJ
AIJ[2] = XIJ[2]/RIJ
"""),
AIJ=[1.0, 0.0, 0.0])
input = dict(
(x, pre_comp[x]) for x in ['RIJ', 'R2IJ', 'XIJ', 'HIJ', 'AIJ'])
pre = sort_precomputed(input, pre_comp)
self.assertEqual(list(pre.keys()),
['HIJ', 'XIJ', 'R2IJ', 'RIJ', 'AIJ'])
finally:
from pysph.sph.equation import precomputed_symbols
Group.pre_comp = precomputed_symbols()