def test_symbol_cache_aliasing_reverse(): """Test to assert that removing he original u[x, y] instance does not impede our alisaing cache or leaks memory. """ # Ensure a clean cache to start with clear_cache() # FIXME: Currently not working, presumably due to our # failure to cache new instances? # assert(len(_SymbolCache) == 0) # Create first instance of u and fill its data grid = Grid(shape=(3, 4)) u = Function(name='u', grid=grid) u.data[:] = 6. u_ref = weakref.ref(u.data) # Create derivative and delete orignal u[x, y] dx = u.dx del u clear_cache() # We still have a references to u # FIXME: Unreliable cache sizes # assert len(_SymbolCache) == 1 # Ensure u[x + h, y] still holds valid data assert np.allclose(dx.args[0].args[1].data, 6.) del dx clear_cache() # FIXME: Unreliable cache sizes # assert len(_SymbolCache) == 0 # We still have a reference to u_h assert u_ref() is None
def unit_box(name='a', shape=(11, 11), grid=None): """Create a field with value 0. to 1. in each dimension""" grid = grid or Grid(shape=shape) a = Function(name=name, grid=grid) dims = tuple([np.linspace(0., 1., d) for d in shape]) a.data[:] = np.meshgrid(*dims)[1] return a
def test_arithmetic(self): """Test arithmetic operations involving Data objects.""" grid = Grid(shape=(16, 16, 16)) u = Function(name='yu3D', grid=grid, space_order=0) u.data[:] = 1 # Simple arithmetic assert np.all(u.data == 1) assert np.all(u.data + 2. == 3.) assert np.all(u.data - 2. == -1.) assert np.all(u.data * 2. == 2.) assert np.all(u.data / 2. == 0.5) assert np.all(u.data % 2 == 1.) # Increments and partial increments u.data[:] += 2. assert np.all(u.data == 3.) u.data[9, :, :] += 1. assert all(np.all(u.data[i, :, :] == 3.) for i in range(9)) assert np.all(u.data[9, :, :] == 4.) # Right operations __rOP__ u.data[:] = 1. arr = np.ndarray(shape=(16, 16, 16), dtype=np.float32) arr.fill(2.) assert np.all(arr - u.data == 1.)
def test_slicing(self): grid = Grid(shape=(4, 4)) x, y = grid.dimensions glb_pos_map = grid.distributor.glb_pos_map myrank = grid.distributor.myrank u = Function(name='u', grid=grid, space_order=0) u.data[:] = myrank # `u.data` is a view of the global data array restricted, on each rank, # to the local rank domain, so it must be == myrank assert np.all(u.data == myrank) assert np.all(u.data._local == myrank) if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(u.data[:2, :2] == myrank) assert u.data[:2, 2:].size == u.data[2:, :2].size == u.data[2:, 2:].size == 0 elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert np.all(u.data[:2, 2:] == myrank) assert u.data[:2, :2].size == u.data[2:, :2].size == u.data[2:, 2:].size == 0 elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(u.data[2:, :2] == myrank) assert u.data[:2, 2:].size == u.data[:2, :2].size == u.data[2:, 2:].size == 0 else: assert np.all(u.data[2:, 2:] == myrank) assert u.data[:2, 2:].size == u.data[2:, :2].size == u.data[:2, :2].size == 0
def test_indexing(self): grid = Grid(shape=(4, 4)) x, y = grid.dimensions glb_pos_map = grid.distributor.glb_pos_map myrank = grid.distributor.myrank u = Function(name='u', grid=grid, space_order=0) u.data[:] = myrank if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert u.data[0, 0] == myrank assert u.data[2, 2] is None assert u.data[2].size == 0 assert u.data[:, 2].size == 0 elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert u.data[0, 0] is None assert u.data[2, 2] is None assert u.data[2].size == 0 assert np.all(u.data[:, 2] == [myrank, myrank]) elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert u.data[0, 0] is None assert u.data[2, 2] is None assert np.all(u.data[2] == [myrank, myrank]) assert u.data[:, 2].size == 0 else: assert u.data[0, 0] is None assert u.data[2, 2] == myrank assert np.all(u.data[2] == [myrank, myrank]) assert np.all(u.data[:, 2] == [myrank, myrank])
def test_partial_offloading(self): """ Check that :class:`Function` objects not using any :class:`SpaceDimension` are computed in Devito-land, rather than via YASK. """ shape = (4, 4, 4) grid = Grid(shape=shape) dx = Dimension(name='dx') dy = Dimension(name='dy') dz = Dimension(name='dz') x, y, z = grid.dimensions u = TimeFunction(name='yu4D', grid=grid, space_order=0) f = Function(name='f', dimensions=(dx, dy, dz), shape=shape) u.data_with_halo[:] = 0. f.data[:] = 0. eqns = [Eq(u.forward, u + 1.), Eq(f, u[1, dx, dy, dz] + 1.)] op = Operator(eqns) op(time=0) assert np.all(u.data[0] == 0.) assert np.all(u.data[1] == 1.) assert np.all(f.data == 2.)
def a(shape=(11, 11)): grid = Grid(shape=shape) a = Function(name='a', grid=grid) xarr = np.linspace(0., 1., shape[0]) yarr = np.linspace(0., 1., shape[1]) a.data[:] = np.meshgrid(xarr, yarr)[1] return a
def test_localviews(self): grid = Grid(shape=(4, 4)) x, y = grid.dimensions glb_pos_map = grid.distributor.glb_pos_map myrank = grid.distributor.myrank u = Function(name='u', grid=grid) u.data[:] = grid.distributor.myrank assert u.data_ro_domain._local[0, 0] == grid.distributor.myrank assert u.data_ro_domain._local[1, 1] == grid.distributor.myrank assert u.data_ro_domain._local[-1, -1] == grid.distributor.myrank assert u.data_ro_with_halo._local[1, 1] == grid.distributor.myrank if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(u.data_ro_with_halo._local[1:, 1:] == myrank) assert np.all(u.data_ro_with_halo._local[0] == 0.) elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert np.all(u.data_ro_with_halo._local[1:3, :2] == myrank) assert np.all(u.data_ro_with_halo._local[0] == 0.) elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(u.data_ro_with_halo._local[:2, 1:3] == myrank) assert np.all(u.data_ro_with_halo._local[2] == 0.) else: assert np.all(u.data_ro_with_halo._local[:2, :2] == myrank) assert np.all(u.data_ro_with_halo._local[2] == 0.)
def test_broadcasting(self): """ Test Data broadcasting, expected to behave as NumPy broadcasting. Notes ----- Refer to https://docs.scipy.org/doc/numpy-1.15.0/user/basics.broadcasting.html for more info about NumPy broadcasting rules. """ grid = Grid(shape=(4, 4, 4)) u = Function(name='yu3D', grid=grid) u.data[:] = 2. # Assign from array with lower-dimensional shape v = np.ones(shape=(4, 4), dtype=u.dtype) u.data[:] = v assert np.all(u.data == 1.) # Assign from array with higher-dimensional shape causes a ValueError exception v = np.zeros(shape=(4, 4, 4, 4), dtype=u.dtype) try: u.data[:] = v except ValueError: assert True except: assert False # Assign from array having shape with some 1-valued entries v = np.zeros(shape=(4, 1, 4), dtype=u.dtype) u.data[:] = v assert np.all(u.data == 0.)
def test_fd_space_staggered(self, space_order, stagger): """ This test compares the discrete finite-difference scheme against polynomials For a given order p, the finite difference scheme should be exact for polynomials of order p :param derivative: name of the derivative to be tested :param space_order: space order of the finite difference stencil """ clear_cache() # dummy axis dimension nx = 100 xx = np.linspace(-1, 1, nx) dx = xx[1] - xx[0] # Symbolic data grid = Grid(shape=(nx,), dtype=np.float32) x = grid.dimensions[0] # Location of the staggered function if stagger == left: off = -.5 side = -x xx2 = xx - off * dx elif stagger == right: off = .5 side = x xx2 = xx[:-1] - off * dx else: off = 0 side = NODE xx2 = xx u = Function(name="u", grid=grid, space_order=space_order, staggered=(side,)) du = Function(name="du", grid=grid, space_order=space_order) # Define polynomial with exact fd coeffs = np.ones((space_order-1,), dtype=np.float32) polynome = sum([coeffs[i]*x**i for i in range(0, space_order-1)]) polyvalues = np.array([polynome.subs(x, xi) for xi in xx2], np.float32) # Fill original data with the polynomial values u.data[:] = polyvalues # True derivative of the polynome Dpolynome = diff(polynome) Dpolyvalues = np.array([Dpolynome.subs(x, xi) for xi in xx], np.float32) # FD derivative, symbolic u_deriv = generic_derivative(u, deriv_order=1, fd_order=space_order, dim=x, stagger=stagger) # Compute numerical FD stencil = Eq(du, u_deriv) op = Operator(stencil, subs={x.spacing: dx}) op.apply() # Check exactness of the numerical derivative except inside space_brd space_border = space_order error = abs(du.data[space_border:-space_border] - Dpolyvalues[space_border:-space_border]) assert np.isclose(np.mean(error), 0., atol=1e-3)
def test_trivial_insertion(self): grid = Grid(shape=(4, 4)) u = Function(name='u', grid=grid, space_order=0) v = Function(name='v', grid=grid, space_order=1) u.data[:] = 1. assert np.all(u.data == 1.) assert np.all(u.data._local == 1.) v.data_with_halo[:] = 1. assert v.data_with_halo[:].shape == (3, 3) assert np.all(v.data_with_halo == 1.) assert np.all(v.data_with_halo[:] == 1.) assert np.all(v.data_with_halo._local == 1.)
def test_function(): grid = Grid(shape=(3, 3, 3)) f = Function(name='f', grid=grid) f.data[0] = 1. pkl_f = pickle.dumps(f) new_f = pickle.loads(pkl_f) # .data is initialized, so it should have been pickled too assert np.all(f.data[0] == 1.) assert np.all(new_f.data[0] == 1.) assert f.space_order == new_f.space_order assert f.dtype == new_f.dtype assert f.shape == new_f.shape
def test_inject_from_field(shape, coords, result, npoints=19): """Test point injection from a second field along a line through the middle of the grid. """ a = unit_box(shape=shape) a.data[:] = 0. b = Function(name='b', grid=a.grid) b.data[:] = 1. p = points(a.grid, ranges=coords, npoints=npoints) expr = p.inject(field=a, expr=b) Operator(expr)(a=a, b=b) indices = [slice(4, 6, 1) for _ in coords] indices[0] = slice(1, -1, 1) assert np.allclose(a.data[indices], result, rtol=1.e-5)
def test_halo_indexing(self): """Test data packing/unpacking in presence of a halo region.""" domain_shape = (16, 16, 16) grid = Grid(shape=domain_shape) u = Function(name='yu3D', grid=grid, space_order=2) assert u.shape == u.data.shape == domain_shape assert u._shape_with_inhalo == u.data_with_halo.shape == (20, 20, 20) assert u.shape_with_halo == u._shape_with_inhalo # W/o MPI, these two coincide # Test simple insertion and extraction u.data_with_halo[0, 0, 0] = 1. u.data[0, 0, 0] = 2. assert u.data_with_halo[0, 0, 0] == 1. assert u.data[0, 0, 0] == 2. assert u.data_with_halo[2, 2, 2] == 2. # Test negative indices u.data_with_halo[-1, -1, -1] = 3. assert u.data[-1, -1, -1] == 0. assert u.data_with_halo[-1, -1, -1] == 3.
def test_symbol_cache_aliasing(): """Test to assert that our aliasing cache isn't defeated by sympys non-aliasing symbol cache. For further explanation consider the symbol u[x, y] and it's first derivative in x, which includes the symbols u[x, y] and u[x + h, y]. The two functions are aliased in devito's caching mechanism to allow multiple stencil indices pointing at the same data object u, but SymPy treats these two instances as separate functions and thus is allowed to delete one or the other when the cache is cleared. The test below asserts that u[x + h, y] is deleted, the data on u is still intact through our own caching mechanism.""" # Ensure a clean cache to start with clear_cache() # FIXME: Currently not working, presumably due to our # failure to cache new instances? # assert(len(_SymbolCache) == 0) # Create first instance of u and fill its data grid = Grid(shape=(3, 4)) u = Function(name='u', grid=grid) u.data[:] = 6. u_ref = weakref.ref(u.data) # Create u[x + h, y] and delete it again dx = u.dx # Contains two u symbols: u[x, y] and u[x + h, y] del dx clear_cache() # FIXME: Unreliable cache sizes # assert len(_SymbolCache) == 1 # We still have a reference to u assert np.allclose(u.data, 6.) # u.data is alive and well # Remove the final instance and ensure u.data got deallocated del u clear_cache() assert u_ref() is None
def test_from_replicated_to_distributed(self): shape = (4, 4) grid = Grid(shape=shape) x, y = grid.dimensions glb_pos_map = grid.distributor.glb_pos_map u = Function(name='u', grid=grid, space_order=0) # distributed v = Function(name='v', grid=grid, space_order=0) # distributed a = np.arange(16).reshape(shape) # replicated # Full array u.data[:] = a if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(u.data == [[0, 1], [4, 5]]) elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert np.all(u.data == [[2, 3], [6, 7]]) elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(u.data == [[8, 9], [12, 13]]) else: assert np.all(u.data == [[10, 11], [14, 15]]) # Subsection (all ranks touched) u.data[:] = 0 u.data[1:3, 1:3] = a[1:3, 1:3] # Same as above but with negative indices v.data[:] = 0 v.data[1:-1, 1:-1] = a[1:-1, 1:-1] if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(u.data == [[0, 0], [0, 5]]) assert np.all(v.data == [[0, 0], [0, 5]]) elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert np.all(u.data == [[0, 0], [6, 0]]) assert np.all(v.data == [[0, 0], [6, 0]]) elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(u.data == [[0, 9], [0, 0]]) assert np.all(v.data == [[0, 9], [0, 0]]) else: assert np.all(u.data == [[10, 0], [0, 0]]) assert np.all(v.data == [[10, 0], [0, 0]]) # The assigned data must have same shape as the one of the distributed array, # otherwise an exception is expected try: u.data[1:3, 1:3] = a[1:2, 1:2] except ValueError: assert True except: assert False
def _new_operator1(shape, blockshape=None, dle=None): blockshape = as_tuple(blockshape) grid = Grid(shape=shape, dtype=np.int32) infield = Function(name='infield', grid=grid) infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape) outfield = Function(name='outfield', grid=grid) stencil = Eq(outfield.indexify(), outfield.indexify() + infield.indexify()*3.0) op = Operator(stencil, dle=dle) blocksizes = get_blocksizes(op, dle, grid, blockshape) op(infield=infield, outfield=outfield, **blocksizes) return outfield, op
def _gen_phys_param(self, field, name, space_order, is_param=False, default_value=0, init_empty=False): if field is None and not init_empty: return default_value if isinstance(field, np.ndarray) or init_empty: function = Function(name=name, grid=self.grid, space_order=space_order, parameter=is_param) if not init_empty: if name is 'rho': initialize_function(function, 1 / field, self.nbpml) else: initialize_function(function, field, self.nbpml) else: function = Constant(name=name, value=field) self._physical_parameters.append(name) return function
def test_function_wo(self): grid = Grid(shape=(3, 3, 3)) i = Dimension(name='i') f = Function(name='f', shape=(1,), dimensions=(i,), grid=grid) u = TimeFunction(name='u', grid=grid) eqns = [Eq(u.forward, u + 1), Eq(f[0], u[0, 0, 0, 0])] op = Operator(eqns, opt='noop', language='openmp') assert len(op.body.maps) == 1 assert op.body.maps[0].pragmas[0].value ==\ ('omp target enter data map(to: u[0:u_vec->size[0]]' '[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])') assert len(op.body.unmaps) == 2 assert op.body.unmaps[0].pragmas[0].value ==\ ('omp target update from(u[0:u_vec->size[0]]' '[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])') assert op.body.unmaps[1].pragmas[0].value ==\ ('omp target exit data map(release: u[0:u_vec->size[0]]' '[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]]) if(devicerm)')
def test_array_rw(self): grid = Grid(shape=(3, 3, 3)) f = Function(name='f', grid=grid) u = TimeFunction(name='u', grid=grid, space_order=2) eqn = Eq(u.forward, u*cos(f*2)) op = Operator(eqn, language='openmp') assert len(op.body.allocs) == 1 assert str(op.body.allocs[0]) ==\ ('float *r0_vec = (float*) ' 'omp_target_alloc(sizeof(float[x_size*y_size*z_size]), ' 'omp_get_default_device());') assert len(op.body.maps) == 2 assert all('r0' not in str(i) for i in op.body.maps) assert len(op.body.frees) == 1 assert str(op.body.frees[0]) ==\ 'omp_target_free(r0_vec, omp_get_default_device());' assert len(op.body.unmaps) == 3 assert all('r0' not in str(i) for i in op.body.unmaps)
def test_if_parallel(self): a = np.arange(36).reshape((6, 6)) grid = Grid(shape=(18, 18)) x, y = grid.dimensions glb_pos_map = grid.distributor.glb_pos_map f = Function(name='f', grid=grid, halo=((3, 3), (3, 3)), dtype=np.int32) initialize_function(f, a, 6, mode='reflect') if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(a[::-1, 0:3] - np.array(f.data[0:6, 6:9]) == 0) assert np.all(a[0:3, ::-1] - np.array(f.data[6:9, 0:6]) == 0) elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert np.all(a[::-1, 3:6] - np.array(f.data[0:6, 9:12]) == 0) assert np.all(a[0:3, ::-1] - np.array(f.data[6:9, 12:18]) == 0) elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert np.all(a[::-1, 0:3] - np.array(f.data[12:18, 6:9]) == 0) assert np.all(a[3:6, ::-1] - np.array(f.data[9:12, 0:6]) == 0) else: assert np.all(a[::-1, 3:6] - np.array(f.data[12:18, 9:12]) == 0) assert np.all(a[3:6, ::-1] - np.array(f.data[9:12, 12:18]) == 0)
def test_save_w_nonaffine_time(self): factor = 4 grid = Grid(shape=(11, 11)) x, y = grid.dimensions t = grid.stepping_dim time = grid.time_dim time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor) f = Function(name='f', grid=grid, dtype=np.int32) u = TimeFunction(name='u', grid=grid) usave = TimeFunction(name='usave', grid=grid, save=2, time_dim=time_subsampled) save_shift = Constant(name='save_shift', dtype=np.int32) eqns = [Eq(u.forward, u[t, f[x, x], f[y, y]] + 1.), Eq(usave.subs(time_subsampled, time_subsampled - save_shift), u)] op = Operator(eqns, opt=('buffering', 'tasking', 'orchestrate')) # We just check the generated code here assert len([i for i in FindSymbols().visit(op) if isinstance(i, Lock)]) == 1 assert len(op._func_table) == 2
def test_misc_dims(self): """ Tests grid-independent :class:`Function`s, which require YASK's "misc" dimensions. """ dx = Dimension(name='dx') grid = Grid(shape=(10, 10)) x, y = grid.dimensions time = grid.time_dim u = TimeFunction(name='u', grid=grid, time_order=1, space_order=4, save=4) c = Function(name='c', dimensions=(x, dx), shape=(10, 5)) step = Eq(u.forward, ( u[time, x-2, y] * c[x, 0] + u[time, x-1, y] * c[x, 1] + u[time, x, y] * c[x, 2] + u[time, x+1, y] * c[x, 3] + u[time, x+2, y] * c[x, 4])) for i in range(10): c.data[i, 0] = 1.0+i c.data[i, 1] = 1.0+i c.data[i, 2] = 3.0+i c.data[i, 3] = 6.0+i c.data[i, 4] = 5.0+i u.data[:] = 0.0 u.data[0, 2, :] = 2.0 op = Operator(step) assert 'run_solution' in str(op) op(time_m=0, time_M=0) assert(np.all(u.data[1, 0, :] == 10.0)) assert(np.all(u.data[1, 1, :] == 14.0)) assert(np.all(u.data[1, 2, :] == 10.0)) assert(np.all(u.data[1, 3, :] == 8.0)) assert(np.all(u.data[1, 4, :] == 10.0)) assert(np.all(u.data[1, 5:10, :] == 0.0))
def test_arrays_defined_over_subdims(self): """ Check code generation when an Array uses a SubDimension. """ grid = Grid(shape=(3, )) x, = grid.dimensions xi, = grid.interior.dimensions f = Function(name='f', grid=grid) a = Array(name='a', dimensions=(xi, ), dtype=grid.dtype) op = Operator( [Eq(a[xi], 1), Eq(f, f + a[xi + 1], subdomain=grid.interior)], dle=('advanced', { 'openmp': False })) assert len(op.parameters) == 6 # neither `x_size` nor `xi_size` are expected here assert not any(i.name in ('x_size', 'xi_size') for i in op.parameters) # Try running it -- regardless of what it will produce, this should run # ie, this checks this error isn't raised: # "ValueError: No value found for parameter xi_size" op()
def test_default_functions(self): """ Test the default argument derivation for functions. """ grid = Grid(shape=(5, 6, 7)) f = TimeFunction(name='f', grid=grid) g = Function(name='g', grid=grid) op = Operator(Eq(g, g + f)) expected = { 'x_size': 5, 'x_s': 0, 'x_e': 5, 'y_size': 6, 'y_s': 0, 'y_e': 6, 'z_size': 7, 'z_s': 0, 'z_e': 7, 'f': f.data, 'g': g.data, } self.verify_arguments(op.arguments(), expected)
def test_operator_leakage_sparse(self): """ Test to ensure that Operator creation does not cause memory leaks for SparseTimeFunctions. """ grid = Grid(shape=(5, 6)) a = Function(name='a', grid=grid) s = SparseTimeFunction(name='s', grid=grid, npoint=1, nt=1) w_a = weakref.ref(a) w_s = weakref.ref(s) # Create operator and delete everything again op = Operator(s.interpolate(a)) w_op = weakref.ref(op) del op del s del a clear_cache() # Test whether things are still hanging around assert w_a() is None assert w_s() is None assert w_op() is None
def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=1000.0, space_order=4, kernel='OT2', nbpml=40, full_run=False, autotune=False, constant=False, checkpointing=False, **kwargs): solver = acoustic_setup(shape=shape, spacing=spacing, nbpml=nbpml, tn=tn, space_order=space_order, kernel=kernel, constant=constant, **kwargs) # Smooth velocity initial_vp = Function(name='v0', grid=solver.model.grid, space_order=space_order) smooth(initial_vp, solver.model.m) dm = np.float32(initial_vp.data**2 - solver.model.m.data) info("Applying Forward") # Whether or not we save the whole time history. We only need the full wavefield # with 'save=True' if we compute the gradient without checkpointing, if we use # checkpointing, PyRevolve will take care of the time history save = full_run and not checkpointing # Define receiver geometry (spread across x, just below surface) rec, u, summary = solver.forward(save=save, autotune=autotune) if constant: # With a new m as Constant m0 = Constant(name="m", value=.25, dtype=np.float32) solver.forward(save=save, m=m0) # With a new m as a scalar value solver.forward(save=save, m=.25) if not full_run: return summary.gflopss, summary.oi, summary.timings, [rec, u.data] info("Applying Adjoint") solver.adjoint(rec, autotune=autotune) info("Applying Born") solver.born(dm, autotune=autotune) info("Applying Gradient") solver.gradient(rec, u, autotune=autotune, checkpointing=checkpointing)
def fwi_gradient(m_in): # Create symbols to hold the gradient and residual grad = Function(name="grad", grid=model.grid) residual = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis, coordinates=geometry.rec_positions) objective = 0. # Creat forward wavefield to reuse to avoid memory overload u0 = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=4, save=geometry.nt) for i in range(nshots): # Important: We force previous wavefields to be destroyed, # so that we may reuse the memory. clear_cache() # Update source location geometry.src_positions[0, :] = source_locations[i, :] # Generate synthetic data from true model true_d, _, _ = solver.forward(m=model.m) # Compute smooth data and full forward wavefield u0 u0.data.fill(0.) smooth_d, _, _ = solver.forward(m=m_in, save=True, u=u0) # Compute gradient from data residual and update objective function residual.data[:] = smooth_d.data[:] - true_d.data[:] objective += .5 * np.linalg.norm(residual.data.flatten())**2 solver.gradient(rec=residual, u=u0, m=m_in, grad=grad) return objective, grad.data
def gradient(self, rec, u, v=None, grad=None, m=None, **kwargs): """ Gradient modelling function for computing the adjoint of the Linearized Born modelling function, ie. the action of the Jacobian adjoint on an input data. :param recin: Receiver data as a numpy array :param u: Symbol for full wavefield `u` (created with save=True) :param v: (Optional) Symbol to store the computed wavefield :param grad: (Optional) Symbol to store the gradient field :returns: Gradient field and performance summary """ # Gradient symbol if grad is None: grad = Function(name='grad', grid=self.model.grid) # Create the forward wavefield if v is None: v = TimeFunction(name='v', grid=self.model.grid, time_order=2, space_order=self.space_order) # Pick m from model unless explicitly provided if m is None: m = m or self.model.m summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, m=m, dt=kwargs.pop('dt', self.dt), **kwargs) return grad, summary
def test_mpi_fullmode_objects(): grid = Grid(shape=(4, 4, 4)) x, y, _ = grid.dimensions # Message f = Function(name='f', grid=grid) obj = MPIMsgEnriched('msg', f, [Halo(x, LEFT)]) pkl_obj = pickle.dumps(obj) new_obj = pickle.loads(pkl_obj) assert obj.name == new_obj.name assert obj.function.name == new_obj.function.name assert all(obj.function.dimensions[i].name == new_obj.function.dimensions[i].name for i in range(grid.dim)) assert new_obj.function.dimensions[0] is new_obj.halos[0].dim # Region x_m, x_M = x.symbolic_min, x.symbolic_max y_m, y_M = y.symbolic_min, y.symbolic_max obj = MPIRegion('reg', 1, [y, x], [(((x, OWNED, LEFT),), {x: (x_m, Min(x_M, x_m))}), (((y, OWNED, LEFT),), {y: (y_m, Min(y_M, y_m))})]) pkl_obj = pickle.dumps(obj) new_obj = pickle.loads(pkl_obj) assert obj.prefix == new_obj.prefix assert obj.key == new_obj.key assert obj.name == new_obj.name assert len(new_obj.arguments) == 2 assert all(d0.name == d1.name for d0, d1 in zip(obj.arguments, new_obj.arguments)) assert all(new_obj.arguments[i] is new_obj.owned[i][0][0][0] # `x` and `y` for i in range(2)) assert new_obj.owned[0][0][0][1] is new_obj.owned[1][0][0][1] # `OWNED` assert new_obj.owned[0][0][0][2] is new_obj.owned[1][0][0][2] # `LEFT` for n, i in enumerate(new_obj.owned): d, v = list(i[1].items())[0] assert d is new_obj.arguments[n] assert v[0] is d.symbolic_min assert v[1] == Min(d.symbolic_max, d.symbolic_min)
def test_multiple_subnests_v0(self): grid = Grid(shape=(3, 3, 3)) x, y, z = grid.dimensions t = grid.stepping_dim f = Function(name='f', grid=grid) u = TimeFunction(name='u', grid=grid, space_order=3) eqn = Eq( u.forward, _R( _R(u[t, x, y, z] + u[t, x + 1, y + 1, z + 1]) * 3. * f + _R(u[t, x + 2, y + 2, z + 2] + u[t, x + 3, y + 3, z + 3]) * 3. * f) + 1.) op = Operator(eqn, opt=('advanced', { 'openmp': True, 'cire-mingain': 0, 'par-nested': 0, 'par-collapse-ncores': 1, 'par-dynamic-work': 0 })) bns, _ = assert_blocking(op, {'x0_blk0'}) trees = retrieve_iteration_tree(bns['x0_blk0']) assert len(trees) == 2 assert trees[0][0] is trees[1][0] assert trees[0][0].pragmas[0].value ==\ 'omp for collapse(2) schedule(dynamic,1)' assert trees[0][2].pragmas[0].value == ('omp parallel for collapse(2) ' 'schedule(dynamic,1) ' 'num_threads(nthreads_nested)') assert trees[1][2].pragmas[0].value == ('omp parallel for collapse(2) ' 'schedule(dynamic,1) ' 'num_threads(nthreads_nested)')
def run(shape=(50, 50, 50), spacing=(10.0, 10.0, 10.0), tn=1000.0, space_order=4, nbl=40, full_run=False, autotune=False, **kwargs): solver = acoustic_ssa_setup(shape=shape, spacing=spacing, nbl=nbl, tn=tn, space_order=space_order, **kwargs) info("Applying Forward") # Define receiver geometry (spread across x, just below surface) rec, u, summary = solver.forward(save=full_run, autotune=autotune) if not full_run: return summary.gflopss, summary.oi, summary.timings, [rec, u.data] # Smooth velocity initial_vp = Function(name='v0', grid=solver.model.grid, space_order=space_order) smooth(initial_vp, solver.model.vp) dm = solver.model.vp - initial_vp info("Applying Adjoint") solver.adjoint(rec, autotune=autotune) info("Applying Born") solver.jacobian(dm, autotune=autotune) info("Applying Gradient") solver.jacobian_adjoint(rec, u, autotune=autotune) return summary.gflopss, summary.oi, summary.timings, [rec, u.data]
def test_argument_derivation_order(self, nt=100): """ Ensure the precedence order of arguments is respected Defaults < (overriden by) Tensor Arguments < Dimensions < Scalar Arguments """ i, j, k = dimify('i j k') shape = (10, 10, 10) grid = Grid(shape=shape, dimensions=(i, j, k)) a = Function(name='a', grid=grid).indexed b_function = TimeFunction(name='b', grid=grid, save=True, time_dim=nt) b = b_function.indexed time = b_function.indices[0] b1 = TimeFunction(name='b1', grid=grid, save=True, time_dim=nt + 1).indexed eqn = Eq(b[time, i, j, k], a[i, j, k]) op = Operator(eqn) # Simple case, same as that tested above. # Repeated here for clarity of further tests. op_arguments, _ = op.arguments() assert (op_arguments[time.start_name] == 0) assert (op_arguments[time.end_name] == nt) # Providing a tensor argument should infer the dimension size from its shape op_arguments, _ = op.arguments(b=b1) assert (op_arguments[time.start_name] == 0) assert (op_arguments[time.end_name] == nt + 1) # Providing a dimension size explicitly should override the automatically inferred op_arguments, _ = op.arguments(b=b1, time=nt - 1) assert (op_arguments[time.start_name] == 0) assert (op_arguments[time.end_name] == nt - 1) # Providing a scalar argument explicitly should override the automatically\ # inferred op_arguments, _ = op.arguments(b=b1, time=nt - 1, time_e=nt - 2) assert (op_arguments[time.start_name] == 0) assert (op_arguments[time.end_name] == nt - 2)
def test_expr_like_lowering(self): """ Test the lowering of an expr-like ConditionalDimension's condition. This test makes an Operator that should indexify and lower the condition passed in the Conditional Dimension """ grid = Grid(shape=(3, 3)) g1 = Function(name='g1', grid=grid) g2 = Function(name='g2', grid=grid) g1.data[:] = 0.49 g2.data[:] = 0.49 x, y = grid.dimensions ci = ConditionalDimension(name='ci', parent=y, condition=Le((g1 + g2), 1.01 * (g1 + g2))) f = Function(name='f', shape=grid.shape, dimensions=(x, ci)) Operator(Eq(f, g1 + g2)).apply() assert np.all(f.data[:] == g1.data[:] + g2.data[:])
def test_default_rules_deriv_offset(self, so, offset): """ Test that default_rules generates correct derivatives when derivatives are evaluated at offset x0. """ grid = Grid(shape=(11, ), extent=(10., )) x = grid.dimensions[0] h_x = x.spacing f_std = Function(name='f', grid=grid, space_order=so) f_sym = Function(name='f', grid=grid, space_order=so, coefficients='symbolic') g = Function(name='g', grid=grid, space_order=so) eval_std = str( Eq(g, f_std.dx(x0=x + offset * h_x / 2)).evaluate.evalf(_PRECISION)) eval_sym = str( Eq(g, f_sym.dx(x0=x + offset * h_x / 2)).evaluate.evalf(_PRECISION)) assert eval_std == eval_sym
def __init__(self, origin, spacing, shape, space_order, nbl=20, dtype=np.float32, subdomains=(), damp_mask=False): self.shape = shape self.nbl = int(nbl) self.origin = tuple([dtype(o) for o in origin]) # Origin of the computational domain with boundary to inject/interpolate # at the correct index origin_pml = tuple( [dtype(o - s * nbl) for o, s in zip(origin, spacing)]) phydomain = PhysicalDomain(self.nbl) subdomains = subdomains + (phydomain, ) shape_pml = np.array(shape) + 2 * self.nbl # Physical extent is calculated per cell, so shape - 1 extent = tuple(np.array(spacing) * (shape_pml - 1)) self.grid = Grid(extent=extent, shape=shape_pml, origin=origin_pml, dtype=dtype, subdomains=subdomains) if self.nbl != 0: # Create dampening field as symbol `damp` self.damp = Function(name="damp", grid=self.grid) initialize_damp(self.damp, self.nbl, self.spacing, mask=damp_mask) self._physical_parameters = ['damp'] else: self.damp = 1 if damp_mask else 0 self._physical_parameters = []
def test_operator_leakage_function(): """ Test to ensure that Operator creation does not cause memory leaks for (Time)Functions. """ grid = Grid(shape=(5, 6)) f = Function(name='f', grid=grid) g = TimeFunction(name='g', grid=grid) # Take weakrefs to test whether symbols are dead or alive w_f = weakref.ref(f) w_g = weakref.ref(g) # Create operator and delete everything again op = Operator(Eq(f, 2 * g)) w_op = weakref.ref(op) del op del f del g clear_cache() # Test whether things are still hanging around assert w_f() is None assert w_g() is None assert w_op() is None
def GradientOperator(model, source, receiver, space_order=4, save=True, kernel='OT2', **kwargs): """ Constructor method for the gradient operator in an acoustic media :param model: :class:`Model` object containing the physical parameters :param source: :class:`PointData` object containing the source geometry :param receiver: :class:`PointData` object containing the acquisition geometry :param time_order: Time discretization order :param space_order: Space discretization order """ m, damp = model.m, model.damp # Gradient symbol and wavefield symbols grad = Function(name='grad', grid=model.grid) u = TimeFunction(name='u', grid=model.grid, save=source.nt if save else None, time_order=2, space_order=space_order) v = TimeFunction(name='v', grid=model.grid, save=None, time_order=2, space_order=space_order) rec = Receiver(name='rec', grid=model.grid, time_range=receiver.time_range, npoint=receiver.npoint) s = model.grid.stepping_dim.spacing eqn = iso_stencil(v, m, s, damp, kernel, forward=False) if kernel == 'OT2': gradient_update = Inc(grad, - u.dt2 * v) elif kernel == 'OT4': gradient_update = Inc(grad, - (u.dt2 + s**2 / 12.0 * u.laplace2(m**(-2))) * v) # Add expression for receiver injection receivers = rec.inject(field=v.backward, expr=rec * s**2 / m, offset=model.nbpml) # Substitute spacing terms to reduce flops return Operator(eqn + receivers + [gradient_update], subs=model.spacing_map, name='Gradient', **kwargs)
def test_no_index(self): """Test behaviour when the ConditionalDimension is used as a symbol in an expression.""" nt = 19 grid = Grid(shape=(11, 11)) time = grid.time_dim u = TimeFunction(name='u', grid=grid) assert(grid.stepping_dim in u.indices) v = Function(name='v', grid=grid) factor = 4 time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor) eqns = [Eq(u.forward, u + 1), Eq(v, v + u*u*time_subsampled)] op = Operator(eqns) op.apply(t_M=nt-2) assert np.all(np.allclose(u.data[(nt-1) % 3], nt-1)) # expected result is 1024 # v = u[0]**2 * 0 + u[4]**2 * 1 + u[8]**2 * 2 + u[12]**2 * 3 + u[16]**2 * 4 # with u[t] = t # v = 16 * 1 + 64 * 2 + 144 * 3 + 256 * 4 = 1600 assert np.all(np.allclose(v.data, 1600))
def test_function_wo(self): grid = Grid(shape=(3, 3, 3)) i = Dimension(name='i') f = Function(name='f', shape=(1, ), dimensions=(i, ), grid=grid) u = TimeFunction(name='u', grid=grid) eqns = [Eq(u.forward, u + 1), Eq(f[0], u[0, 0, 0, 0])] op = Operator(eqns, opt='noop') assert len(op.body[1].header) == 2 assert len(op.body[1].footer) == 2 assert op.body[1].header[0].value ==\ ('omp target enter data map(to: u[0:u_vec->size[0]]' '[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])') assert str(op.body[1].header[1]) == '' assert str(op.body[1].footer[0]) == '' assert op.body[1].footer[1].contents[0].value ==\ ('omp target update from(u[0:u_vec->size[0]]' '[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])') assert op.body[1].footer[1].contents[1].value ==\ ('omp target exit data map(release: u[0:u_vec->size[0]]' '[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
def test_sdf_distance(self, space_order): """Check the mean error in SDF values""" # Error threshold thres = 0.11 # 0.11 grid increments # Sphere of radius 40, center (50, 50, 50) sphere = 'tests/trial_surfaces/sphere.ply' # Grid configuration extent = (100., 100., 100.) shape = (101, 101, 101) origin = (0., 0., 0.) grid = Grid(shape=shape, extent=extent, origin=origin) x, y, z = grid.dimensions h_x, h_y, h_z = grid.spacing f = Function(name='f', grid=grid, space_order=space_order) # Create signed distance function sig = SignedDistanceFunction(f, sphere) pos_x, pos_y, pos_z = np.where(sig.sdf.data != -space_order//2-1) def true_dist(pos_1, pos_2, pos_3): return 40 - np.sqrt((h_x*pos_1 - 50)**2 + (h_y*pos_2 - 50)**2 + (h_z*pos_3 - 50)**2) err = np.absolute(sig.sdf.data[pos_x, pos_y, pos_z] - true_dist(pos_x, pos_y, pos_z)) avg = np.mean(err) if avg > thres: err_message = "Mean error in measured distances is {:.6}" raise ValueError(err_message.format(avg))
def test_precomputed_interpolation(): """ Test interpolation with PrecomputedSparseFunction which accepts precomputed values for interpolation coefficients """ shape = (101, 101) points = [(.05, .9), (.01, .8), (0.07, 0.84)] origin = (0, 0) grid = Grid(shape=shape, origin=origin) r = 2 # Constant for linear interpolation # because we interpolate across 2 neighbouring points in each dimension def init(data): for i in range(data.shape[0]): for j in range(data.shape[1]): data[i, j] = sin(grid.spacing[0] * i) + sin(grid.spacing[1] * j) return data m = Function(name='m', grid=grid, initializer=init, space_order=0) gridpoints, interpolation_coeffs = precompute_linear_interpolation( points, grid, origin) sf = PrecomputedSparseFunction(name='s', grid=grid, r=r, npoint=len(points), gridpoints=gridpoints, interpolation_coeffs=interpolation_coeffs) eqn = sf.interpolate(m) op = Operator(eqn) op() expected_values = [sin(point[0]) + sin(point[1]) for point in points] assert (all(np.isclose(sf.data, expected_values, rtol=1e-6)))
def test_array_rw(self): grid = Grid(shape=(3, 3, 3)) f = Function(name='f', grid=grid) u = TimeFunction(name='u', grid=grid, space_order=2) eqn = Eq(u.forward, u * cos(f * 2)) op = Operator(eqn) assert len(op.body[1].header) == 7 assert str(op.body[1].header[0]) == 'float (*r1)[y_size][z_size];' assert op.body[1].header[1].text ==\ 'posix_memalign((void**)&r1, 64, sizeof(float[x_size][y_size][z_size]))' assert op.body[1].header[2].value ==\ ('omp target enter data map(alloc: r1[0:x_size][0:y_size][0:z_size])' '') assert len(op.body[1].footer) == 6 assert str(op.body[1].footer[0]) == '' assert op.body[1].footer[1].value ==\ ('omp target exit data map(delete: r1[0:x_size][0:y_size][0:z_size])' ' if((x_size != 0) && (y_size != 0) && (z_size != 0))') assert op.body[1].footer[2].text == 'free(r1)'
def test_override_function_size(self): """ Test runtime size overrides for :class:`Function` dimensions. Note: The current behaviour for size-only arguments seems ambiguous (eg. op(x=3, y=4), as it sets `dim_size` as well as `dim_end`. Since `dim_size` is used for the cast, we can get garbage results if it does not agree with the shape of the provided data. This should error out, or potentially we could set the corresponding size, while aliasing `dim` to `dim_e`? The same should be tested for :class:`TimeFunction` once fixed. """ grid = Grid(shape=(5, 6, 7)) g = Function(name='g', grid=grid) op = Operator(Eq(g, 1.)) args = {'x': 3, 'y': 4, 'z': 5} arguments = op.arguments(**args) expected = { 'x_size': 5, 'x_s': 0, 'x_e': 3, 'y_size': 6, 'y_s': 0, 'y_e': 4, 'z_size': 7, 'z_s': 0, 'z_e': 5, 'g': g.data } self.verify_arguments(arguments, expected) # Verify execution op(**args) assert (g.data[3:, 4:, 5:] == 0.).all() assert (g.data[:3, :4, :5] == 1.).all()
def test_tasking_fused(self): nt = 10 bundle0 = Bundle() grid = Grid(shape=(10, 10, 10), subdomains=bundle0) tmp = Function(name='tmp', grid=grid) u = TimeFunction(name='u', grid=grid, save=nt) v = TimeFunction(name='v', grid=grid, save=nt) w = TimeFunction(name='w', grid=grid) eqns = [Eq(w.forward, w + 1), Eq(tmp, w.forward), Eq(u.forward, tmp, subdomain=bundle0), Eq(v.forward, tmp, subdomain=bundle0)] op = Operator(eqns, opt=('tasking', 'fuse', 'orchestrate')) # Check generated code assert len(retrieve_iteration_tree(op)) == 4 locks = [i for i in FindSymbols().visit(op) if isinstance(i, Lock)] assert len(locks) == 1 # Only 1 because it's only `tmp` that needs protection assert len(op._func_table) == 2 exprs = FindNodes(Expression).visit(op._func_table['copy_device_to_host0'].root) assert len(exprs) == 19 assert str(exprs[12]) == 'int id = sdata0->id;' assert str(exprs[13]) == 'const int time = sdata0->time;' assert str(exprs[14]) == 'lock0[0] = 1;' assert exprs[15].write is u assert exprs[16].write is v assert str(exprs[17]) == 'lock0[0] = 2;' assert str(exprs[18]) == 'sdata0->flag = 1;' op.apply(time_M=nt-2) assert np.all(u.data[nt-1] == 9) assert np.all(v.data[nt-1] == 9)
def _initialize_bcs(self, bcs="damp"): # Create dampening field as symbol `damp` if self.nbl == 0: self.damp = 1 if bcs == "mask" else 0 return # First initialization init = self.damp is None # Get current Function if alread yinitialized self.damp = self.damp or Function(name="damp", grid=self.grid) if callable(bcs): bcs(self.damp, self.nbl) else: re_init = ((bcs == "mask" and mmin(self.damp) == 0) or (bcs == "damp" and mmax(self.damp) == 1)) if init or re_init: if re_init and not init: bcs_o = "damp" if bcs == "mask" else "mask" warning("Re-initializing damp profile from %s to %s" % (bcs_o, bcs)) warning("Model has to be created with `bcs=\"%s\"`" "for this WaveSolver" % bcs) initialize_damp(self.damp, self.padsizes, self.spacing, abc_type=bcs, fs=self.fs) self._physical_parameters.update(['damp'])
def test_multiple_subnests(self): grid = Grid(shape=(3, 3, 3)) x, y, z = grid.dimensions t = grid.stepping_dim f = Function(name='f', grid=grid) u = TimeFunction(name='u', grid=grid) eqn = Eq(u.forward, ((u[t, x, y, z] + u[t, x+1, y+1, z+1])*3*f + (u[t, x+2, y+2, z+2] + u[t, x+3, y+3, z+3])*3*f + 1)) op = Operator(eqn, dse='aggressive', dle=('advanced', {'openmp': True})) trees = retrieve_iteration_tree(op._func_table['bf0'].root) assert len(trees) == 2 assert trees[0][0] is trees[1][0] assert trees[0][0].pragmas[0].value ==\ 'omp for collapse(1) schedule(dynamic,1)' assert trees[0][2].pragmas[0].value == ('omp parallel for collapse(1) ' 'schedule(dynamic,1) ' 'num_threads(nthreads_nested)') assert trees[1][2].pragmas[0].value == ('omp parallel for collapse(1) ' 'schedule(dynamic,1) ' 'num_threads(nthreads_nested)')
def test_misc_data(self): """ Test data insertion/indexing for Functions with mixed distributed/replicated Dimensions. """ dx = Dimension(name='dx') grid = Grid(shape=(4, 4)) x, y = grid.dimensions glb_pos_map = grid.distributor.glb_pos_map # Note: `grid` must be passed to `c` since `x` is a distributed dimension, # and `grid` carries the `x` decomposition c = Function(name='c', grid=grid, dimensions=(x, dx), shape=(4, 5)) # Data insertion for i in range(4): c.data[i, 0] = 1.0+i c.data[i, 1] = 1.0+i c.data[i, 2] = 3.0+i c.data[i, 3] = 6.0+i c.data[i, 4] = 5.0+i # Data indexing if LEFT in glb_pos_map[x]: assert(np.all(c.data[0] == [1., 1., 3., 6., 5.])) assert(np.all(c.data[1] == [2., 2., 4., 7., 6.])) else: assert(np.all(c.data[2] == [3., 3., 5., 8., 7.])) assert(np.all(c.data[3] == [4., 4., 6., 9., 8.])) # Same as before, but with negative indices and non-trivial slices if LEFT in glb_pos_map[x]: assert(np.all(c.data[0:-3] == [1., 1., 3., 6., 5.])) assert(np.all(c.data[-3:-2] == [2., 2., 4., 7., 6.])) else: assert(np.all(c.data[-2:-1] == [3., 3., 5., 8., 7.])) assert(np.all(c.data[-1] == [4., 4., 6., 9., 8.]))
def test_at_is_actually_working(shape, expected): """ Check that autotuning is actually running when switched on, in both 2D and 3D operators. """ grid = Grid(shape=shape) infield = Function(name='infield', grid=grid) infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape) outfield = Function(name='outfield', grid=grid) stencil = Eq(outfield.indexify(), outfield.indexify() + infield.indexify()*3.0) op = Operator(stencil, dle=('blocking', {'openmp': False, 'blockinner': True, 'blockalways': True})) # Run with whatever `configuration` says (by default, basic+preemptive) op(infield=infield, outfield=outfield, autotune=True) assert op._state['autotuning'][-1]['runs'] == 4 assert op._state['autotuning'][-1]['tpr'] == 1 # Now try `aggressive` autotuning configuration['autotuning'] = 'aggressive' op(infield=infield, outfield=outfield, autotune=True) assert op._state['autotuning'][-1]['runs'] == expected assert op._state['autotuning'][-1]['tpr'] == 1 configuration['autotuning'] = configuration._defaults['autotuning'] # Try again, but using the Operator API directly op(infield=infield, outfield=outfield, autotune='aggressive') assert op._state['autotuning'][-1]['runs'] == expected assert op._state['autotuning'][-1]['tpr'] == 1 # Similar to above op(infield=infield, outfield=outfield, autotune=('aggressive', 'preemptive')) assert op._state['autotuning'][-1]['runs'] == expected assert op._state['autotuning'][-1]['tpr'] == 1
def test_simple_indexing(self): """Test data packing/unpacking via basic indexing.""" grid = Grid(shape=(16, 16, 16)) u = Function(name='yu3D', grid=grid, space_order=0) # Test simple insertion and extraction u.data[0, 1, 1] = 1. assert u.data[0, 0, 0] == 0. assert u.data[0, 1, 1] == 1. assert np.all(u.data == u.data[:, :, :]) assert 1. in u.data[0] assert 1. in u.data[0, 1] # Test negative indices assert u.data[0, -15, -15] == 1. u.data[6, 0, 0] = 1. assert u.data[-10, :, :].sum() == 1. # Test setting whole array to given value u.data[:] = 3. assert np.all(u.data == 3.) # Test insertion of single value into block u.data[5, :, 5] = 5. assert np.all(u.data[5, :, 5] == 5.) # Test extraction of block with negative indices sliced = u.data[-11, :, -11] assert sliced.shape == (16,) assert np.all(sliced == 5.) # Test insertion of block into block block = np.ndarray(shape=(1, 16, 1), dtype=np.float32) block.fill(4.) u.data[4:5, :, 4:5] = block assert np.all(u.data[4, :, 4] == block)
def test_timesteps_per_at_run(): """ Check that each autotuning run (ie with a given block shape) takes ``autotuning.core.options['squeezer']`` timesteps, for an operator performing the increment ``a[t + timeorder, ...] = f(a[t, ...], ...)``. """ shape = (30, 30, 30) grid = Grid(shape=shape) x, y, z = grid.dimensions t = grid.stepping_dim # Function infield = Function(name='infield', grid=grid) infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape) outfield = Function(name='outfield', grid=grid) stencil = Eq(outfield.indexify(), outfield.indexify() + infield.indexify()*3.0) op = Operator(stencil, dle=('blocking', {'openmp': False, 'blockalways': True})) op(infield=infield, outfield=outfield, autotune=True) assert op._state['autotuning'][-1]['runs'] == 4 assert op._state['autotuning'][-1]['tpr'] == 1 # TimeFunction with increasing time order; increasing the time order # shouldn't affect how many iterations the autotuner is gonna run for to in [1, 2, 4]: infield = TimeFunction(name='infield', grid=grid, time_order=to) infield.data[:] = np.arange(reduce(mul, infield.shape), dtype=np.int32).reshape(infield.shape) outfield = TimeFunction(name='outfield', grid=grid, time_order=to) stencil = Eq(outfield[t + to, x, y, z], outfield.indexify() + infield.indexify()*3.0) op = Operator(stencil, dle=('blocking', {'openmp': False, 'blockalways': True})) op(infield=infield, outfield=outfield, time=20, autotune=True) assert op._state['autotuning'][-1]['runs'] == 4 assert op._state['autotuning'][-1]['tpr'] == options['squeezer'] + 1
def test_index_alignment(const): """ A much simpler test meant to ensure that the forward and reverse indices are correctly aligned (i.e. u * v , where u is the forward field and v the reverse field corresponds to the correct timesteps in u and v). The forward operator does u = u + 1 which means that the field a will be equal to nt (0 -> 1 -> 2 -> 3), the number of timesteps this operator is run for. The field at the last time step of the forward is used to initialise the field v for the reverse pass. The reverse operator does v = v - 1, which means that if the reverse operator is run for the same number of timesteps as the forward operator, v should be 0 at the last time step (3 -> 2 -> 1 -> 0). There is also a grad = grad + u * v accumulator in the reverse operator. If the alignment is correct, u and v should have the same value at every time step: 0 -> 1 -> 2 -> 3 u 0 <- 1 <- 2 <- 3 v and hence grad = 0*0 + 1*1 + 2*2 + 3*3 = sum(n^2) where n -> [0, nt] If the test fails, the resulting number can tell you how the fields are misaligned """ n = 4 grid = Grid(shape=(2, 2)) order_of_eqn = 1 modulo_factor = order_of_eqn + 1 nt = n - order_of_eqn u = TimeFunction(name='u', grid=grid, save=n) # Increment one in the forward pass 0 -> 1 -> 2 -> 3 fwd_op = Operator(Eq(u.forward, u + 1.*const)) # Invocation 1 fwd_op(time=nt-1, constant=1) last_time_step_v = nt % modulo_factor # Last time step should be equal to the number of timesteps we ran assert(np.allclose(u.data[nt, :, :], nt)) v = TimeFunction(name='v', grid=grid, save=None) v.data[last_time_step_v, :, :] = u.data[nt, :, :] # Decrement one in the reverse pass 3 -> 2 -> 1 -> 0 adj_eqn = Eq(v, v.forward - 1.*const) adj_op = Operator(adj_eqn) # Invocation 2 adj_op(time=nt-1, constant=1) # Last time step should be back to 0 assert(np.allclose(v.data[0, :, :], 0)) # Reset v to run the backward again v.data[last_time_step_v, :, :] = u.data[nt, :, :] prod = Function(name="prod", grid=grid) # Multiply u and v and add them # = 3*3 + 2*2 + 1*1 + 0*0 prod_eqn = Eq(prod, prod + u * v) comb_op = Operator([adj_eqn, prod_eqn]) # Invocation 3 comb_op(time=nt-1, constant=1) final_value = sum([n**2 for n in range(nt)]) # Final value should be sum of squares of first nt natural numbers assert(np.allclose(prod.data, final_value)) # Now reset to repeat all the above tests with checkpointing prod.data[:] = 0 v.data[last_time_step_v, :, :] = u.data[nt, :, :] # Checkpointed version doesn't require to save u u_nosave = TimeFunction(name='u_n', grid=grid) # change equations to use new symbols fwd_eqn_2 = Eq(u_nosave.forward, u_nosave + 1.*const) fwd_op_2 = Operator(fwd_eqn_2) cp = DevitoCheckpoint([u_nosave]) wrap_fw = CheckpointOperator(fwd_op_2, constant=1) prod_eqn_2 = Eq(prod, prod + u_nosave * v) comb_op_2 = Operator([adj_eqn, prod_eqn_2]) wrap_rev = CheckpointOperator(comb_op_2, constant=1) wrp = Revolver(cp, wrap_fw, wrap_rev, None, nt) # Invocation 4 wrp.apply_forward() assert(np.allclose(u_nosave.data[last_time_step_v, :, :], nt)) # Invocation 5 wrp.apply_reverse() assert(np.allclose(v.data[0, :, :], 0)) assert(np.allclose(prod.data, final_value))
def test_indexing_in_views(self): grid = Grid(shape=(4, 4)) x, y = grid.dimensions glb_pos_map = grid.distributor.glb_pos_map myrank = grid.distributor.myrank u = Function(name='u', grid=grid, space_order=0) u.data[:] = myrank # Note that the `1`s are global indices view = u.data[1:, 1:] assert np.all(view[:] == myrank) if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert view.shape == (1, 1) assert np.all(view == 0.) assert view[0, 0] == 0. assert view[1, 1] is None assert view[1].shape == (0, 1) elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert view.shape == (1, 2) assert np.all(view == 1.) assert view[0, 0] is None assert view[1, 1] is None assert view[1].shape == (0, 2) elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert view.shape == (2, 1) assert np.all(view == 2.) assert view[0, 0] is None assert view[1, 1] is None assert view[1].shape == (1,) assert np.all(view[1] == 2.) else: assert view.shape == (2, 2) assert np.all(view == 3.) assert view[0, 0] is None assert view[1, 1] == 3. assert view[1].shape == (2,) assert np.all(view[1] == 3.) # Now we further slice into `view` view2 = view[1:, 1:] assert np.all(view2[:] == myrank) if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert view2.shape == (0, 0) elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert view2.shape == (0, 2) elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]: assert view2.shape == (2, 0) else: assert view2.shape == (2, 2) # Now a change in `view2` by the only rank that "sees" it should affect # both `view` and `u.data` view2[:] += 1 if RIGHT in glb_pos_map[x] and RIGHT in glb_pos_map[y]: assert np.all(u.data[:] == myrank + 1) assert np.all(view[:] == myrank + 1) assert np.all(view2[:] == myrank + 1) else: assert np.all(view[:] == myrank) assert np.all(view2[:] == myrank) assert view2.size == 0