def make_condensational_growth_solver(nr, r_min, r_max, GC_max, grid_layout, psi_coord, pdf_of_r, drdt_of_r, opts: Options): # psi = psi(p) dp_dr = psi_coord.dx_dr dx_dr = grid_layout.dx_dr xh, dx = np.linspace( grid_layout.x(r_min), grid_layout.x(r_max), nr + 1, retstep=True ) rh = grid_layout.r(xh) x = np.linspace( xh[0] + dx / 2, xh[-1] - dx / 2, nr ) r = grid_layout.r(x) def pdf_of_r_over_psi(r): return pdf_of_r(r) / psi_coord.dx_dr(r) psi = discretised_analytical_solution(rh, pdf_of_r_over_psi, midpoint_value=True, r=r) dp_dt = drdt_of_r(rh) * dp_dr(rh) G = dp_dr(r) / dx_dr(r) # C = dr_dt * dt / dr # GC = dp_dr / dx_dr * dr_dt * dt / dr = # \ \_____ / _..____/ # \_____.._____/ \_ dt/dx # | # dp_dt dt = GC_max * dx / np.amax(dp_dt) GCh = dp_dt * dt / dx # CFL condition np.testing.assert_array_less(np.abs(GCh), 1) g_factor = ScalarField(G.astype(dtype=opts.dtype), halo=opts.n_halo, boundary_conditions=(ExtrapolatedBoundaryCondition(),)) state = ScalarField(psi.astype(dtype=opts.dtype), halo=opts.n_halo, boundary_conditions=(ConstantBoundaryCondition(0),)) GC_field = VectorField([GCh.astype(dtype=opts.dtype)], halo=opts.n_halo, boundary_conditions=(ConstantBoundaryCondition(0),)) stepper = Stepper( options=opts, n_dims=1, non_unit_g_factor=True ) return ( Solver(stepper=stepper, g_factor=g_factor, advectee=state, advector=GC_field), r, rh, dx, dt, g_factor )
def test_timing_2d(benchmark, options, dtype, grid_static_str, concurrency_str, plot=False): if grid_static_str == "static": grid_static = True elif grid_static_str == "dynamic": grid_static = False else: raise ValueError() if concurrency_str == "serial": numba.set_num_threads(1) else: numba.set_num_threads(numba.config.NUMBA_NUM_THREADS) setup = Setup(n_rotations=6) _, __, z = from_pdf_2d(setup.pdf, xrange=setup.xrange, yrange=setup.yrange, gridsize=setup.grid) C = (-.5, .25) grid = z.shape advector_data = [ np.full((grid[0] + 1, grid[1]), C[0], dtype=options.dtype), np.full((grid[0], grid[1] + 1), C[1], dtype=options.dtype) ] advector = VectorField(advector_data, halo=options.n_halo, boundary_conditions=(PeriodicBoundaryCondition(), PeriodicBoundaryCondition())) advectee = ScalarField(data=z.astype(dtype=options.dtype), halo=options.n_halo, boundary_conditions=(PeriodicBoundaryCondition(), PeriodicBoundaryCondition())) if grid_static: stepper = Stepper(options=options, grid=grid) else: stepper = Stepper(options=options, n_dims=2) solver = Solver(stepper=stepper, advectee=advectee, advector=advector) def set_z(): solver.advectee.get()[:] = z benchmark.pedantic(solver.advance, (setup.nt, ), setup=set_z, warmup_rounds=1, rounds=3) if options.n_iters == 1 or options.flux_corrected_transport: np.testing.assert_almost_equal(np.amin(solver.advectee.get()), h0) assert np.amax(solver.advectee.get()) < 10 * h if plot: pyplot.imshow(solver.advectee.get()) pyplot.colorbar() pyplot.show()
def test_make_upwind(self): # Arrange psi_data = np.array((0, 1, 0)) flux_data = np.array((0, 0, 1, 0)) options = Options() halo = options.n_halo traversals = Traversals(grid=psi_data.shape, halo=halo, jit_flags={}, n_threads=1) upwind = make_upwind(options=options, non_unit_g_factor=False, traversals=traversals) bc = [PeriodicBoundaryCondition()] psi = ScalarField(psi_data, halo, bc) psi_impl = psi.impl flux_impl = VectorField((flux_data, ), halo, bc).impl null_impl = ScalarField.make_null(len(psi_data.shape)).impl # Act upwind(psi_impl[0], *flux_impl, *null_impl) # Assert np.testing.assert_array_equal(psi.get(), np.roll(psi_data, 1))
def test_apply_scalar(n_threads, halo, grid, loop): n_dims = len(grid) if n_dims == 1 and n_threads > 1: return # arrange traversals = Traversals(grid, halo, jit_flags, n_threads) sut = traversals.apply_scalar(loop=loop) scl_null_arg_impl = ScalarField.make_null(n_dims).impl vec_null_arg_impl = VectorField.make_null(n_dims).impl out = ScalarField(np.zeros(grid), halo, [ConstantBoundaryCondition(np.nan)] * n_dims) # act sut(_cell_id_scalar, _cell_id_scalar, *out.impl[0], *vec_null_arg_impl[0], *vec_null_arg_impl[1], *scl_null_arg_impl[0], *scl_null_arg_impl[1], *scl_null_arg_impl[0], *scl_null_arg_impl[1], *scl_null_arg_impl[0], *scl_null_arg_impl[1]) # assert data = out.get() assert data.shape == grid focus = (-halo, -halo) for i in range(halo, halo + grid[0]): for j in (-1, ) if n_dims == 1 else range(halo, halo + grid[1]): value = indexers[n_dims].at[0](focus, data, i, j) assert value == (n_dims if loop else 1) * cell_id(i, j) assert scl_null_arg_impl[0][0][meta_halo_valid] assert vec_null_arg_impl[0][0][meta_halo_valid] assert not out.impl[0][0][meta_halo_valid]
def test_diffusion_only_2d( data0=np.array([[0, 0, 0], [0, 1., 0], [0, 0, 0]]), mu=(.1, .1), nt=1): # Arrange options = Options(non_zero_mu_coeff=True) bc = [PeriodicBoundaryCondition()] * 2 advectee = ScalarField(data0, options.n_halo, bc) advector = VectorField(data=(np.zeros( (data0.shape[0] + 1, data0.shape[1])), np.zeros( (data0.shape[0], data0.shape[1] + 1))), halo=options.n_halo, boundary_conditions=bc) solver = Solver(stepper=Stepper(options=options, grid=data0.shape), advector=advector, advectee=advectee) # Act solver.advance(nt=nt, mu_coeff=mu) # Assert data1 = solver.advectee.get() np.testing.assert_almost_equal(actual=np.sum(data1), desired=np.sum(data0)) assert np.amax(data0) > np.amax(data1) assert np.amin(data1) >= 0 assert np.count_nonzero(data1) == 5
def test_scalar_1d(self, data, halo, side): # arrange field = ScalarField(data, halo, (PeriodicBoundaryCondition(),)) meta_and_data, fill_halos = field.impl traversals = Traversals(grid=data.shape, halo=halo, jit_flags={}) sut, _ = traversals.make_boundary_conditions() # act sut(*meta_and_data, *fill_halos) # assert if side == LEFT: np.testing.assert_array_equal(field.data[:halo], data[-halo:]) elif side == RIGHT: np.testing.assert_array_equal(field.data[-halo:], data[:halo]) else: raise ValueError()
def test_scalar_2d(self, halo, n_threads): # arrange data = np.array([[1, 6], [2, 7], [3, 8], [4, 9]]) bc = (PeriodicBoundaryCondition(), PolarBoundaryCondition(data.shape, 0, 1)) field = ScalarField(data, halo, bc) meta_and_data, fill_halos = field.impl traversals = Traversals(grid=data.shape, halo=halo, jit_flags={}, n_threads=n_threads) sut = traversals._fill_halos_scalar # act for thread_id in numba.prange(n_threads): sut(thread_id, *meta_and_data, *fill_halos) # assert np.testing.assert_array_equal( field.data[halo:-halo, :halo], np.roll(field.get()[:, :halo], data.shape[0] // 2, axis=0)) np.testing.assert_array_equal( field.data[halo:-halo, -halo:], np.roll(field.get()[:, -halo:], data.shape[0] // 2, axis=0))
def test_1d_contiguous(): grid = (44, ) data = np.empty(grid) bc = (PeriodicBoundaryCondition(), ) sut = ScalarField(data, halo=1, boundary_conditions=bc) assert sut.get().data.contiguous
def test_2d_second_dim_contiguous(): grid = (44, 44) data = np.empty(grid) bc = (PeriodicBoundaryCondition(), PeriodicBoundaryCondition()) sut = ScalarField(data, halo=1, boundary_conditions=bc) assert sut.get()[0, :].data.contiguous