Skip to content

Commit

Permalink
fix data_allocation.py imports
Browse files Browse the repository at this point in the history
  • Loading branch information
@halungge
halungge committed Jan 17, 2025
1 parent 775c419 commit fbf16d1
Show file tree
Hide file tree
Showing 10 changed files with 120 additions and 150 deletions.
56 changes: 22 additions & 34 deletions ...sphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,34 +97,22 @@ def reference(
def input_data(self, grid) -> dict:
p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim)
p_vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim)
cell_idx = grid.connectivities[dims.E2CDim]
cell_idx_new = data_alloc.flatten_first_two_dims(dims.ECDim, field=cell_idx)
cell_blk = data_alloc.constant_field(grid, 1, dims.EdgeDim, dims.E2CDim, dtype=gtx.int32)
cell_blk_new = data_alloc.flatten_first_two_dims(dims.ECDim, field=cell_blk)
pos_on_tplane_e_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim)
pos_on_tplane_e_1_new = data_alloc.flatten_first_two_dims(
dims.ECDim, field=pos_on_tplane_e_1
cell_idx = data_alloc.flatten_first_two_dims(
dims.ECDim, field=(grid.connectivities[dims.E2CDim])
)
pos_on_tplane_e_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim)
pos_on_tplane_e_2_new = data_alloc.flatten_first_two_dims(
dims.ECDim, field=pos_on_tplane_e_2
)
primal_normal_cell_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim)
primal_normal_cell_1_new = data_alloc.flatten_first_two_dims(
dims.ECDim, field=primal_normal_cell_1
)
dual_normal_cell_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim)
dual_normal_cell_1_new = data_alloc.flatten_first_two_dims(
dims.ECDim, field=dual_normal_cell_1
)
primal_normal_cell_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim)
primal_normal_cell_2_new = data_alloc.flatten_first_two_dims(
dims.ECDim, field=primal_normal_cell_2
)
dual_normal_cell_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim)
dual_normal_cell_2_new = data_alloc.flatten_first_two_dims(
dims.ECDim, field=dual_normal_cell_2
cell_blk = data_alloc.flatten_first_two_dims(
dims.ECDim,
field=(data_alloc.constant_field(grid, 1, dims.EdgeDim, dims.E2CDim, dtype=gtx.int32)),
)
pos_on_tplane_e_1 = data_alloc.random_field(grid, dims.ECDim)
pos_on_tplane_e_2 = data_alloc.random_field(grid, dims.ECDim)
primal_normal_cell_1 = data_alloc.random_field(grid, dims.ECDim)

dual_normal_cell_1 = data_alloc.random_field(grid, dims.ECDim)
primal_normal_cell_2 = data_alloc.random_field(grid, dims.ECDim)

dual_normal_cell_2 = data_alloc.random_field(grid, dims.ECDim)

p_cell_idx = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32)
p_cell_rel_idx_dsl = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32)
p_cell_blk = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32)
Expand All @@ -135,14 +123,14 @@ def input_data(self, grid) -> dict:
return dict(
p_vn=p_vn,
p_vt=p_vt,
cell_idx=cell_idx_new,
cell_blk=cell_blk_new,
pos_on_tplane_e_1=pos_on_tplane_e_1_new,
pos_on_tplane_e_2=pos_on_tplane_e_2_new,
primal_normal_cell_1=primal_normal_cell_1_new,
dual_normal_cell_1=dual_normal_cell_1_new,
primal_normal_cell_2=primal_normal_cell_2_new,
dual_normal_cell_2=dual_normal_cell_2_new,
cell_idx=cell_idx,
cell_blk=cell_blk,
pos_on_tplane_e_1=pos_on_tplane_e_1,
pos_on_tplane_e_2=pos_on_tplane_e_2,
primal_normal_cell_1=primal_normal_cell_1,
dual_normal_cell_1=dual_normal_cell_1,
primal_normal_cell_2=primal_normal_cell_2,
dual_normal_cell_2=dual_normal_cell_2,
p_cell_idx=p_cell_idx,
p_cell_rel_idx_dsl=p_cell_rel_idx_dsl,
p_cell_blk=p_cell_blk,
Expand Down
3 changes: 1 addition & 2 deletions model/atmosphere/advection/tests/advection_tests/test_advection.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,9 +121,8 @@ def test_advection_run_single_step(
)

diagnostic_state = construct_diagnostic_init_state(icon_grid, advection_init_savepoint, ntracer)
prep_adv = construct_prep_adv(icon_grid, advection_init_savepoint)
prep_adv = construct_prep_adv(advection_init_savepoint)
p_tracer_now = advection_init_savepoint.tracer(ntracer)

p_tracer_new = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim)
dtime = advection_init_savepoint.get_metadata("dtime").get("dtime")

Expand Down
5 changes: 2 additions & 3 deletions model/atmosphere/advection/tests/advection_tests/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,21 +88,20 @@ def construct_diagnostic_init_state(

def construct_diagnostic_exit_state(
icon_grid,
savepoint: sb.AdvectionInitSavepoint,
savepoint: sb.AdvectionExitSavepoint,
ntracer: int,
) -> advection_states.AdvectionDiagnosticState:
return advection_states.AdvectionDiagnosticState(
airmass_now=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim),
airmass_new=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim),
grf_tend_tracer=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim),
# FIXME (@halungge) check with David - Savepoint and field do not match
hfl_tracer=savepoint.hfl_tracer(ntracer),
vfl_tracer=savepoint.vfl_tracer(ntracer),
)


def construct_prep_adv(
icon_grid, savepoint: sb.AdvectionInitSavepoint
savepoint: sb.AdvectionInitSavepoint,
) -> advection_states.AdvectionPrepAdvState:
return advection_states.AdvectionPrepAdvState(
vn_traj=savepoint.vn_traj(),
Expand Down
24 changes: 10 additions & 14 deletions model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla4.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,9 @@
import numpy as np
import pytest

import icon4py.model.common.utils.data_allocation as data_alloc
from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import calculate_nabla4
from icon4py.model.common import dimension as dims
from icon4py.model.common.type_alias import vpfloat, wpfloat
from icon4py.model.common.utils.data_allocation import (
random_field,
zero_field,
)
from icon4py.model.common import dimension as dims, type_alias as ta
from icon4py.model.testing.helpers import StencilTest


Expand Down Expand Up @@ -95,17 +91,17 @@ def input_data(self, grid):
if np.any(grid.connectivities[dims.E2C2VDim] == -1):
pytest.xfail("Stencil does not support missing neighbors.")

u_vert = random_field(grid, dims.VertexDim, dims.KDim, dtype=vpfloat)
v_vert = random_field(grid, dims.VertexDim, dims.KDim, dtype=vpfloat)
u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat)
v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat)

primal_normal_vert_v1 = random_field(grid, dims.ECVDim, dtype=wpfloat)
primal_normal_vert_v2 = random_field(grid, dims.ECVDim, dtype=wpfloat)
primal_normal_vert_v1 = data_alloc.random_field(grid, dims.ECVDim, dtype=ta.wpfloat)
primal_normal_vert_v2 = data_alloc.random_field(grid, dims.ECVDim, dtype=ta.wpfloat)

z_nabla2_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat)
inv_vert_vert_length = random_field(grid, dims.EdgeDim, dtype=wpfloat)
inv_primal_edge_length = random_field(grid, dims.EdgeDim, dtype=wpfloat)
z_nabla2_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat)
inv_vert_vert_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat)
inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat)

z_nabla4_e2 = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
z_nabla4_e2 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)

return dict(
u_vert=u_vert,
Expand Down
11 changes: 5 additions & 6 deletions ...here/dycore/tests/dycore_stencil_tests/test_add_interpolated_horizontal_advection_of_w.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,11 @@
import numpy as np
import pytest

import icon4py.model.common.utils.data_allocation as data_alloc
from icon4py.model.atmosphere.dycore.stencils.add_interpolated_horizontal_advection_of_w import (
add_interpolated_horizontal_advection_of_w,
)
from icon4py.model.common import dimension as dims
from icon4py.model.common.type_alias import vpfloat, wpfloat
from icon4py.model.common.utils.data_allocation import random_field
from icon4py.model.common import dimension as dims, type_alias as ta
from icon4py.model.testing.helpers import StencilTest


Expand Down Expand Up @@ -46,9 +45,9 @@ def reference(

@pytest.fixture
def input_data(self, grid):
z_v_grad_w = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
e_bln_c_s = random_field(grid, dims.CEDim, dtype=wpfloat)
ddt_w_adv = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
z_v_grad_w = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)
e_bln_c_s = data_alloc.random_field(grid, dims.CEDim, dtype=ta.wpfloat)
ddt_w_adv = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)

return dict(
e_bln_c_s=e_bln_c_s,
Expand Down
25 changes: 12 additions & 13 deletions ...mosphere/dycore/tests/dycore_stencil_tests/test_compute_advective_normal_wind_tendency.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,7 @@
from icon4py.model.atmosphere.dycore.stencils.compute_advective_normal_wind_tendency import (
compute_advective_normal_wind_tendency,
)
from icon4py.model.common import dimension as dims
from icon4py.model.common.type_alias import vpfloat, wpfloat
from icon4py.model.common import dimension as dims, type_alias as ta
from icon4py.model.testing.helpers import StencilTest


Expand Down Expand Up @@ -88,19 +87,19 @@ def input_data(self, grid):
if np.any(grid.connectivities[dims.E2CDim] == -1):
pytest.xfail("Stencil does not support missing neighbors.")

z_kin_hor_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
coeff_gradekin = data_alloc.random_field(grid, dims.ECDim, dtype=vpfloat)
z_ekinh = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
zeta = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=vpfloat)
vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
f_e = data_alloc.random_field(grid, dims.EdgeDim, dtype=wpfloat)
c_lin_e = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim, dtype=wpfloat)
z_w_con_c_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
z_kin_hor_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)
coeff_gradekin = data_alloc.random_field(grid, dims.ECDim, dtype=ta.vpfloat)
z_ekinh = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)
zeta = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat)
vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)
f_e = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat)
c_lin_e = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat)
z_w_con_c_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)
vn_ie = data_alloc.random_field(
grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat
grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.vpfloat
)
ddqz_z_full_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
ddt_vn_apc = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
ddqz_z_full_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)
ddt_vn_apc = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)

return dict(
z_kin_hor_e=z_kin_hor_e,
Expand Down
47 changes: 23 additions & 24 deletions ...e/dycore/tests/dycore_stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,11 @@
import numpy as np
import pytest

import icon4py.model.common.utils.data_allocation as data_alloc
from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_advection_of_rho_and_theta import (
compute_horizontal_advection_of_rho_and_theta,
)
from icon4py.model.common import dimension as dims
from icon4py.model.common.type_alias import vpfloat, wpfloat
from icon4py.model.common.utils.data_allocation import random_field
from icon4py.model.common import dimension as dims, type_alias as ta
from icon4py.model.testing.helpers import StencilTest


Expand All @@ -36,7 +35,7 @@ def compute_btraj_numpy(
p_dthalf: float,
**kwargs,
) -> tuple[np.array, ...]:
lvn_pos = np.where(p_vn > wpfloat("0.0"), True, False)
lvn_pos = np.where(p_vn > ta.wpfloat("0.0"), True, False)
pos_on_tplane_e_1 = np.expand_dims(pos_on_tplane_e_1, axis=-1)
pos_on_tplane_e_2 = np.expand_dims(pos_on_tplane_e_2, axis=-1)
primal_normal_cell_1 = np.expand_dims(primal_normal_cell_1, axis=-1)
Expand Down Expand Up @@ -185,26 +184,26 @@ def input_data(self, grid):
if np.any(grid.connectivities[dims.E2CDim] == -1):
pytest.xfail("Stencil does not support missing neighbors.")

p_vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat)
p_vt = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
pos_on_tplane_e_1 = random_field(grid, dims.ECDim, dtype=wpfloat)
pos_on_tplane_e_2 = random_field(grid, dims.ECDim, dtype=wpfloat)
primal_normal_cell_1 = random_field(grid, dims.ECDim, dtype=wpfloat)
dual_normal_cell_1 = random_field(grid, dims.ECDim, dtype=wpfloat)
primal_normal_cell_2 = random_field(grid, dims.ECDim, dtype=wpfloat)
dual_normal_cell_2 = random_field(grid, dims.ECDim, dtype=wpfloat)
p_dthalf = wpfloat("2.0")

rho_ref_me = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
theta_ref_me = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat)
z_grad_rth_1 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
z_grad_rth_2 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
z_grad_rth_3 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
z_grad_rth_4 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
z_rth_pr_1 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
z_rth_pr_2 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat)
z_rho_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat)
z_theta_v_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat)
p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat)
p_vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)
pos_on_tplane_e_1 = data_alloc.random_field(grid, dims.ECDim, dtype=ta.wpfloat)
pos_on_tplane_e_2 = data_alloc.random_field(grid, dims.ECDim, dtype=ta.wpfloat)
primal_normal_cell_1 = data_alloc.random_field(grid, dims.ECDim, dtype=ta.wpfloat)
dual_normal_cell_1 = data_alloc.random_field(grid, dims.ECDim, dtype=ta.wpfloat)
primal_normal_cell_2 = data_alloc.random_field(grid, dims.ECDim, dtype=ta.wpfloat)
dual_normal_cell_2 = data_alloc.random_field(grid, dims.ECDim, dtype=ta.wpfloat)
p_dthalf = ta.wpfloat("2.0")

rho_ref_me = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)
theta_ref_me = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat)
z_grad_rth_1 = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)
z_grad_rth_2 = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)
z_grad_rth_3 = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)
z_grad_rth_4 = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)
z_rth_pr_1 = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)
z_rth_pr_2 = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat)
z_rho_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat)
z_theta_v_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat)

return dict(
p_vn=p_vn,
Expand Down
38 changes: 17 additions & 21 deletions ...phere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_15_to_18.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,11 @@
import numpy as np
import pytest

import icon4py.model.common.utils.data_allocation as data_alloc
from icon4py.model.atmosphere.dycore.stencils.fused_velocity_advection_stencil_15_to_18 import (
fused_velocity_advection_stencil_15_to_18,
)
from icon4py.model.common import dimension as dims
from icon4py.model.common.utils.data_allocation import (
random_field,
random_mask,
zero_field,
)
from icon4py.model.testing.helpers import StencilTest

from .test_add_extra_diffusion_for_w_con_approaching_cfl import (
Expand Down Expand Up @@ -184,33 +180,33 @@ def reference(

@pytest.fixture
def input_data(self, grid):
z_w_con_c = random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1})
w = random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1})
coeff1_dwdz = random_field(grid, dims.CellDim, dims.KDim)
coeff2_dwdz = random_field(grid, dims.CellDim, dims.KDim)
z_w_con_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1})
w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1})
coeff1_dwdz = data_alloc.random_field(grid, dims.CellDim, dims.KDim)
coeff2_dwdz = data_alloc.random_field(grid, dims.CellDim, dims.KDim)

z_v_grad_w = random_field(grid, dims.EdgeDim, dims.KDim)
e_bln_c_s = random_field(grid, dims.CEDim)
z_v_grad_w = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim)
e_bln_c_s = data_alloc.random_field(grid, dims.CEDim)

levelmask = random_mask(grid, dims.KDim)
cfl_clipping = random_mask(grid, dims.CellDim, dims.KDim)
owner_mask = random_mask(grid, dims.CellDim)
ddqz_z_half = random_field(grid, dims.CellDim, dims.KDim)
area = random_field(grid, dims.CellDim)
geofac_n2s = random_field(grid, dims.CellDim, dims.C2E2CODim)
levelmask = data_alloc.random_mask(grid, dims.KDim)
cfl_clipping = data_alloc.random_mask(grid, dims.CellDim, dims.KDim)
owner_mask = data_alloc.random_mask(grid, dims.CellDim)
ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim)
area = data_alloc.random_field(grid, dims.CellDim)
geofac_n2s = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim)

z_w_con_c_full = zero_field(grid, dims.CellDim, dims.KDim)
ddt_w_adv = zero_field(grid, dims.CellDim, dims.KDim)
z_w_con_c_full = data_alloc.zero_field(grid, dims.CellDim, dims.KDim)
ddt_w_adv = data_alloc.zero_field(grid, dims.CellDim, dims.KDim)

scalfac_exdiff = 10.0
cfl_w_limit = 3.0
dtime = 2.0

k = zero_field(grid, dims.KDim, dtype=gtx.int32)
k = data_alloc.zero_field(grid, dims.KDim, dtype=gtx.int32)
for level in range(grid.num_levels):
k[level] = level

cell = zero_field(grid, dims.CellDim, dtype=gtx.int32)
cell = data_alloc.zero_field(grid, dims.CellDim, dtype=gtx.int32)
for c in range(grid.num_cells):
cell[c] = c

Expand Down
Loading

0 comments on commit fbf16d1

Please sign in to comment.