def test_double_pass_donor_cell(n_iters): courant = .5 options = Options(n_iters=n_iters, DPDC=True, nonoscillatory=True) state = np.array([0, 1, 0], dtype=options.dtype) boundary_conditions = (Periodic(),) mpdata = Solver( stepper=Stepper(options=options, n_dims=state.ndim, non_unit_g_factor=False), advectee=ScalarField( state, halo=options.n_halo, boundary_conditions=boundary_conditions ), advector=VectorField( (np.full(state.shape[0] + 1, courant, dtype=options.dtype),), halo=options.n_halo, boundary_conditions=boundary_conditions ) ) steps = 1 conserved = np.sum(mpdata.advectee.get()) mpdata.advance(steps) assert np.sum(mpdata.advectee.get()) == conserved
def __init__(self, nz, dt, advector_of_t, advectee_of_zZ_at_t0, g_factor_of_zZ, mpdata_settings): self.t = 0 self.dt = dt self.advector_of_t = advector_of_t grid = (nz, ) options = Options(n_iters=mpdata_settings['n_iters'], infinite_gauge=mpdata_settings['iga'], flux_corrected_transport=mpdata_settings['fct'], third_order_terms=mpdata_settings['tot']) stepper = Stepper(options=options, grid=grid, non_unit_g_factor=True) bcs = (ExtrapolatedBoundaryCondition(), ) g_factor = ScalarField(data=g_factor_of_zZ( arakawa_c.z_scalar_coord(grid)), halo=options.n_halo, boundary_conditions=bcs) advector = VectorField(data=(np.full(nz + 1, advector_of_t(0)), ), halo=options.n_halo, boundary_conditions=bcs) self.advectee = ScalarField(data=advectee_of_zZ_at_t0( arakawa_c.z_scalar_coord(grid)), halo=options.n_halo, boundary_conditions=bcs) self.solver = Solver(stepper=stepper, advectee=self.advectee, advector=advector, g_factor=g_factor)
def test_single_timestep(options): # Arrange stepper = Stepper(options=options, grid=GRID, non_unit_g_factor=True) advector = nondivergent_vector_field_2d(GRID, SIZE, TIMESTEP, stream_function, options.n_halo) g_factor = ScalarField(RHOD.astype(dtype=options.dtype), halo=options.n_halo, boundary_conditions=(Periodic(), Periodic())) mpdatas = {} for key, value in VALUES.items(): advectee = ScalarField(np.full(GRID, value, dtype=options.dtype), halo=options.n_halo, boundary_conditions=(Periodic(), Periodic())) mpdatas[key] = Solver(stepper=stepper, advectee=advectee, advector=advector, g_factor=g_factor) # Act for mpdata in mpdatas.values(): mpdata.advance(n_steps=1) # Assert for value in mpdatas.values(): assert np.isfinite(value.advectee.get()).all()
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_diffusion_only_2d(data0=np.array([[0, 0, 0], [0, 1., 0], [0, 0, 0]]), mu_coeff=(.1, .1), n_steps=1): # Arrange options = Options(non_zero_mu_coeff=True) boundary_conditions = tuple([Periodic()] * 2) advectee = ScalarField(data0, options.n_halo, boundary_conditions) 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=boundary_conditions) solver = Solver(stepper=Stepper(options=options, grid=data0.shape), advector=advector, advectee=advectee) # Act solver.advance(n_steps=n_steps, mu_coeff=mu_coeff) # 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 __init__(self, *, fields, n_iters=2, infinite_gauge=True, flux_corrected_transport=True, third_order_terms=False): self.grid = fields.g_factor.shape self.asynchronous = False self.thread: (Thread, None) = None options = Options( n_iters=n_iters, infinite_gauge=infinite_gauge, flux_corrected_transport=flux_corrected_transport, third_order_terms=third_order_terms ) disable_threads_if_needed = {} if not conf.JIT_FLAGS['parallel']: disable_threads_if_needed['n_threads'] = 1 stepper = Stepper(options=options, grid=self.grid, non_unit_g_factor=True, **disable_threads_if_needed) # CFL condition for d in range(len(fields.advector)): np.testing.assert_array_less(np.abs(fields.advector[d]), 1) self.advector = fields.advector advector_impl = VectorField(fields.advector, halo=options.n_halo, boundary_conditions=(PeriodicBoundaryCondition(), PeriodicBoundaryCondition())) self.g_factor = fields.g_factor g_factor_impl = ScalarField(fields.g_factor.astype(dtype=options.dtype), halo=options.n_halo, boundary_conditions=(PeriodicBoundaryCondition(), PeriodicBoundaryCondition())) self.mpdatas = {} for k, v in fields.advectees.items(): advectee = ScalarField(np.full(self.grid, v, dtype=options.dtype), halo=options.n_halo, boundary_conditions=(PeriodicBoundaryCondition(), PeriodicBoundaryCondition())) self.mpdatas[k] = Solver(stepper=stepper, advectee=advectee, advector=advector_impl, g_factor=g_factor_impl)
def test_upwind(shape, ij0, out, courant_number): value = 44 scalar_field_init = np.zeros(shape) scalar_field_init[ij0] = value vector_field_init = ( np.full((shape[0] + 1, shape[1]), courant_number[0]), np.full((shape[0], shape[1] + 1), courant_number[1]) ) options = Options(n_iters=1) bcs = (Periodic(), Periodic()) advectee = ScalarField(scalar_field_init, halo=options.n_halo, boundary_conditions=bcs) advector = VectorField(vector_field_init, halo=options.n_halo, boundary_conditions=bcs) mpdata = Solver( stepper=Stepper(options=options, grid=shape, n_threads=1), advector=advector, advectee=advectee ) mpdata.advance(n_steps=1) np.testing.assert_array_equal( mpdata.advectee.get(), out )
def test_timing_2d(benchmark, options, grid_static_str, num_threads, plot=False): if grid_static_str == "static": grid_static = True elif grid_static_str == "dynamic": grid_static = False else: raise ValueError() numba.set_num_threads(num_threads) settings = Settings(n_rotations=6) _, __, psi = from_pdf_2d(settings.pdf, xrange=settings.xrange, yrange=settings.yrange, gridsize=settings.grid) advectee = ScalarField(data=psi.astype(dtype=options.dtype), halo=options.n_halo, boundary_conditions=(Periodic(), Periodic())) advector = VectorField(data=(np.full( (advectee.grid[0] + 1, advectee.grid[1]), COURANT[0], dtype=options.dtype), np.full( (advectee.grid[0], advectee.grid[1] + 1), COURANT[1], dtype=options.dtype)), halo=options.n_halo, boundary_conditions=(Periodic(), Periodic())) if grid_static: stepper = Stepper(options=options, grid=psi.shape) else: stepper = Stepper(options=options, n_dims=2) solver = Solver(stepper=stepper, advectee=advectee, advector=advector) def set_z(): solver.advectee.get()[:] = psi benchmark.pedantic(solver.advance, (settings.n_steps, ), setup=set_z, warmup_rounds=1, rounds=3) if options.n_iters == 1 or options.nonoscillatory: np.testing.assert_almost_equal(np.amin(solver.advectee.get()), H_0) assert np.amax(solver.advectee.get()) < 10 * H if plot: pyplot.imshow(solver.advectee.get()) pyplot.colorbar() pyplot.show()
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_shared_advector(): n_x = 100 arr = np.zeros(n_x) opt1 = Options(n_iters=2, DPDC=True) opt2 = Options(n_iters=2) b_c = (Periodic(),) halo = opt1.n_halo assert opt2.n_halo == halo advector = VectorField(data=(np.zeros(n_x + 1),), halo=halo, boundary_conditions=b_c) _ = Solver( stepper=Stepper(options=opt1, grid=(n_x,)), advectee=ScalarField(data=arr, halo=halo, boundary_conditions=b_c), advector=advector ) solver = Solver( stepper=Stepper(options=opt2, grid=(n_x,)), advectee=ScalarField(data=arr, halo=halo, boundary_conditions=b_c), advector=advector ) solver.advance(1)
def __init__( self, nz, dt, advector_of_t, advectee_of_zZ_at_t0, g_factor_of_zZ, mpdata_settings, ): self.__t = 0 self.dt = dt self.advector_of_t = advector_of_t grid = (nz, ) options = Options( n_iters=mpdata_settings["n_iters"], infinite_gauge=mpdata_settings["iga"], nonoscillatory=mpdata_settings["fct"], third_order_terms=mpdata_settings["tot"], ) stepper = Stepper(options=options, grid=grid, non_unit_g_factor=True) bcs = (Extrapolated(), ) zZ_scalar = arakawa_c.z_scalar_coord(grid) / nz g_factor = ScalarField( data=g_factor_of_zZ(zZ_scalar), halo=options.n_halo, boundary_conditions=bcs, ) advector = VectorField( data=(np.full(nz + 1, advector_of_t(0)), ), halo=options.n_halo, boundary_conditions=bcs, ) self.advectee = ScalarField( data=advectee_of_zZ_at_t0(zZ_scalar), halo=options.n_halo, boundary_conditions=bcs, ) self.solver = Solver( stepper=stepper, advectee=self.advectee, advector=advector, g_factor=g_factor, )
def test_upwind_1d(): state = np.array([0, 1, 0]) courant = 1 options = Options(n_iters=1) mpdata = Solver( stepper=Stepper(options=options, n_dims=len(state.shape), non_unit_g_factor=False), advectee=ScalarField( state.astype(options.dtype), halo=options.n_halo, boundary_conditions=(Periodic(),) ), advector=VectorField( (np.full(state.shape[0] + 1, courant, dtype=options.dtype),), halo=options.n_halo, boundary_conditions=(Periodic(),) ) ) n_steps = 5 conserved = np.sum(mpdata.advectee.get()) mpdata.advance(n_steps) assert np.sum(mpdata.advectee.get()) == conserved
"output": case_data[8] } # Arrange data = case["input"].reshape((case["nx"], case["ny"])) courant = [case["Cx"], case["Cy"]] options = Options(n_iters=case["ni"], dimensionally_split=case["dimsplit"]) grid = data.shape advector_data = [ np.full((grid[0] + 1, grid[1]), courant[0], dtype=options.dtype), np.full((grid[0], grid[1] + 1), courant[1], dtype=options.dtype) ] bcs = (Periodic(), Periodic()) advector = VectorField(advector_data, halo=options.n_halo, boundary_conditions=bcs) advectee = ScalarField(data=data.astype(dtype=options.dtype), halo=options.n_halo, boundary_conditions=bcs) stepper = Stepper(options=options, grid=grid, non_unit_g_factor=False) mpdata = Solver(stepper=stepper, advectee=advectee, advector=advector) sut = mpdata # Act sut.advance(n_steps=case["nt"]) # Assert np.testing.assert_almost_equal(sut.advectee.get(), case["output"].reshape( case["nx"], case["ny"]), decimal=4)
def __init__(self, *, advectees, stream_function, rhod_of_zZ, dt, grid, size, displacement, n_iters=2, infinite_gauge=True, nonoscillatory=True, third_order_terms=False): self.grid = grid self.size = size self.dt = dt self.stream_function = stream_function self.stream_function_time_dependent = ( "t" in inspect.signature(stream_function).parameters) self.asynchronous = False self.thread: (Thread, None) = None self.displacement = displacement self.t = 0 options = Options( n_iters=n_iters, infinite_gauge=infinite_gauge, nonoscillatory=nonoscillatory, third_order_terms=third_order_terms, ) disable_threads_if_needed = {} if not conf.JIT_FLAGS["parallel"]: disable_threads_if_needed["n_threads"] = 1 stepper = Stepper(options=options, grid=self.grid, non_unit_g_factor=True, **disable_threads_if_needed) advector_impl = VectorField( ( np.full((grid[0] + 1, grid[1]), np.nan), np.full((grid[0], grid[1] + 1), np.nan), ), halo=options.n_halo, boundary_conditions=(Periodic(), Periodic()), ) g_factor = make_rhod(self.grid, rhod_of_zZ) g_factor_impl = ScalarField( g_factor.astype(dtype=options.dtype), halo=options.n_halo, boundary_conditions=(Periodic(), Periodic()), ) self.g_factor_vec = ( rhod_of_zZ(zZ=x_vec_coord(self.grid)[-1]), rhod_of_zZ(zZ=z_vec_coord(self.grid)[-1]), ) self.mpdatas = {} for k, v in advectees.items(): advectee_impl = ScalarField( np.asarray(v, dtype=options.dtype), halo=options.n_halo, boundary_conditions=(Periodic(), Periodic()), ) self.mpdatas[k] = Solver( stepper=stepper, advectee=advectee_impl, advector=advector_impl, g_factor=g_factor_impl, )