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 make_step_impl(options, non_unit_g_factor, grid, n_threads): n_iters = options.n_iters n_dims = len(grid) halo = options.n_halo non_zero_mu_coeff = options.non_zero_mu_coeff flux_corrected_transport = options.flux_corrected_transport traversals = Traversals(grid, halo, options.jit_flags, n_threads=n_threads) upwind = make_upwind(options, non_unit_g_factor, traversals) flux_first_pass = make_flux_first_pass(options, traversals) flux_subsequent = make_flux_subsequent(options, traversals) antidiff = make_antidiff(non_unit_g_factor, options, traversals) laplacian = make_laplacian(non_unit_g_factor, options, traversals) fct_psi_min = make_psi_extremum(min, options, traversals) fct_psi_max = make_psi_extremum(max, options, traversals) fct_beta_down = make_beta(min, non_unit_g_factor, options, traversals) fct_beta_up = make_beta(max, non_unit_g_factor, options, traversals) fct_correction = make_correction(options, traversals) @numba.njit(**options.jit_flags) def axpy(out_meta, out0, out1, a, x_meta, x0, x1, y_meta, y0, y1): out0[:] = a[0] * x0[:] + y0[:] if n_dims > 1: out1[:] = a[1] * x1[:] + y1[:] out_meta[meta_halo_valid] = False @numba.njit(**options.jit_flags) def step(nt, mu_coeff, advectee, advectee_bc, advector, advector_bc, g_factor, g_factor_bc, vectmp_a, vectmp_a_bc, vectmp_b, vectmp_b_bc, vectmp_c, vectmp_c_bc, psi_min, psi_min_bc, psi_max, psi_max_bc, beta_up, beta_up_bc, beta_down, beta_down_bc): vec_bc = advector_bc null_vecfield = advector null_vecfield_bc = vec_bc time = clock() for _ in range(nt): if non_zero_mu_coeff: advector_orig = advector advector = vectmp_c for it in range(n_iters): if it == 0: if flux_corrected_transport: fct_psi_min(psi_min, advectee, advectee_bc, null_vecfield, null_vecfield_bc) fct_psi_max(psi_max, advectee, advectee_bc, null_vecfield, null_vecfield_bc) if non_zero_mu_coeff: laplacian(advector, advectee, advectee_bc, null_vecfield_bc) axpy(*advector, mu_coeff, *advector, *advector_orig) flux_first_pass(vectmp_a, advector, advectee, advectee_bc, vec_bc) flux = vectmp_a else: if it == 1: advector_oscilatory = advector advector_nonoscilatory = vectmp_a flux = vectmp_b elif it % 2 == 0: advector_oscilatory = vectmp_a advector_nonoscilatory = vectmp_b flux = vectmp_a else: advector_oscilatory = vectmp_b advector_nonoscilatory = vectmp_a flux = vectmp_b antidiff(advector_nonoscilatory, advectee, advectee_bc, advector_oscilatory, vec_bc, g_factor, g_factor_bc) flux_subsequent(flux, advectee, advectee_bc, advector_nonoscilatory, vec_bc) if flux_corrected_transport: fct_beta_down(beta_down, flux, vec_bc, advectee, advectee_bc, psi_min, psi_min_bc, g_factor, g_factor_bc) fct_beta_up(beta_up, flux, vec_bc, advectee, advectee_bc, psi_max, psi_max_bc, g_factor, g_factor_bc) fct_correction(advector_nonoscilatory, vec_bc, beta_down, beta_down_bc, beta_up, beta_up_bc) flux_subsequent(flux, advectee, advectee_bc, advector_nonoscilatory, vec_bc) upwind(advectee, flux, vec_bc, g_factor, g_factor_bc) if non_zero_mu_coeff: advector = advector_orig return (clock() - time) / nt return step
def make_step_impl(options, non_unit_g_factor, grid): n_iters = options.n_iters n_dims = len(grid) halo = options.n_halo non_zero_mu_coeff = options.non_zero_mu_coeff flux_corrected_transport = options.flux_corrected_transport traversals = Traversals(grid, halo, options.jit_flags) upwind = make_upwind(options, non_unit_g_factor, traversals) flux_first_pass = make_flux_first_pass(options, traversals) flux_subsequent = make_flux_subsequent(options, traversals) antidiff = make_antidiff(non_unit_g_factor, options, traversals) laplacian = make_laplacian(non_unit_g_factor, options, traversals) fct_psi_min = make_psi_extremum(min, options, traversals) fct_psi_max = make_psi_extremum(max, options, traversals) fct_beta_down = make_beta(min, non_unit_g_factor, options, traversals) fct_beta_up = make_beta(max, non_unit_g_factor, options, traversals) fct_correction = make_correction(options, traversals) @numba.njit(**options.jit_flags) def axpy(out_meta, out0, out1, a, x_meta, x0, x1, y_meta, y0, y1): out0[:] = a * x0[:] + y0[:] if n_dims > 1: out1[:] = a * x1[:] + y1[:] out_meta[meta_halo_valid] = False @numba.njit(**options.jit_flags) def step(nt, mu_coeff, psi, psi_bc, GC_phys, GC_phys_bc, g_factor, g_factor_bc, vectmp_a, vectmp_a_bc, vectmp_b, vectmp_b_bc, vectmp_c, vectmp_c_bc, psi_min, psi_min_bc, psi_max, psi_max_bc, beta_up, beta_up_bc, beta_down, beta_down_bc ): vec_bc = GC_phys_bc null_vecfield = GC_phys null_vecfield_bc = vec_bc for _ in range(nt): if non_zero_mu_coeff: GC_orig = GC_phys GC_phys = vectmp_c for it in range(n_iters): if it == 0: if flux_corrected_transport: fct_psi_min(psi_min, psi, psi_bc, null_vecfield, null_vecfield_bc) fct_psi_max(psi_max, psi, psi_bc, null_vecfield, null_vecfield_bc) if non_zero_mu_coeff: laplacian(GC_phys, psi, psi_bc, null_vecfield_bc) axpy(*GC_phys, mu_coeff, *GC_phys, *GC_orig) flux_first_pass(vectmp_a, GC_phys, psi, psi_bc, vec_bc) flux = vectmp_a else: if it == 1: GC_unco = GC_phys GC_corr = vectmp_a flux = vectmp_b elif it % 2 == 0: GC_unco = vectmp_a GC_corr = vectmp_b flux = vectmp_a else: GC_unco = vectmp_b GC_corr = vectmp_a flux = vectmp_b antidiff(GC_corr, psi, psi_bc, GC_unco, vec_bc, g_factor, g_factor_bc) flux_subsequent(flux, psi, psi_bc, GC_corr, vec_bc) if flux_corrected_transport: fct_beta_down(beta_down, flux, vec_bc, psi, psi_bc, psi_min, psi_min_bc, g_factor, g_factor_bc) fct_beta_up(beta_up, flux, vec_bc, psi, psi_bc, psi_max, psi_max_bc, g_factor, g_factor_bc) fct_correction(GC_corr, vec_bc, beta_down, beta_down_bc, beta_up, beta_up_bc) flux_subsequent(flux, psi, psi_bc, GC_corr, vec_bc) upwind(psi, flux, vec_bc, g_factor, g_factor_bc) if non_zero_mu_coeff: GC_phys = GC_orig return step