Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Debug dycore #611

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

Debug dycore #611

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e78e3f8
Uncommented dycore corrector
muellch Oct 30, 2024
6940830
Uncombine some stencils in velocity advection
muellch Oct 30, 2024
2947703
Uncombined some more vel adv stencils
muellch Oct 30, 2024
5ba684f
Make changes as harmless as possible
muellch Nov 28, 2024
ab41fa4
uncombine some more
muellch Dec 1, 2024
240f985
Uncombine everything needed to compute vn
muellch Dec 4, 2024
8a323d4
cast some lower bounds to to int32
muellch Dec 18, 2024
50755fc
Fix horizontal extend of mass_flx_ic stencils in corrector
muellch Jan 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion base-requirements-dev.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# VCS
-e git+https://github.com/GridTools/gt4py.git@icon4py_20241113#egg=gt4py # use tagged release until #596 & gt4py#1738 is merged
git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python
git+https://github.com/GridTools/serialbox@a1b9334f2330f2deb432e44a5fd7907db08b0076#egg=serialbox&subdirectory=src/serialbox-python

# PyPI
bump2version>=1.0.1
Expand Down
128 changes: 113 additions & 15 deletions model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,16 +19,34 @@
from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_wp import (
init_cell_kdim_field_with_zero_wp,
)
from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_vp import (
init_cell_kdim_field_with_zero_vp,
)

from icon4py.model.atmosphere.dycore.stencils.accumulate_prep_adv_fields import (
accumulate_prep_adv_fields,
)
from icon4py.model.atmosphere.dycore.stencils.add_analysis_increments_from_data_assimilation import (
add_analysis_increments_from_data_assimilation,
)
from icon4py.model.atmosphere.dycore.stencils.set_theta_v_prime_ic_at_lower_boundary import (
set_theta_v_prime_ic_at_lower_boundary,
)
from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import (
extrapolate_temporally_exner_pressure,
)
from icon4py.model.atmosphere.dycore.stencils.add_analysis_increments_to_vn import (
add_analysis_increments_to_vn,
)
from icon4py.model.atmosphere.dycore.stencils.compute_first_vertical_derivative import (
compute_first_vertical_derivative,
)
from icon4py.model.atmosphere.dycore.stencils.interpolate_to_half_levels_vp import (
interpolate_to_half_levels_vp,
)
from icon4py.model.atmosphere.dycore.stencils.interpolate_to_surface import (
interpolate_to_surface,
)
from icon4py.model.atmosphere.dycore.stencils.add_temporal_tendencies_to_vn import (
add_temporal_tendencies_to_vn,
)
Expand Down Expand Up @@ -471,6 +489,15 @@ def __init__(
self._compute_theta_and_exner = compute_theta_and_exner.with_backend(self._backend)
self._compute_exner_from_rhotheta = compute_exner_from_rhotheta.with_backend(self._backend)
self._update_theta_v = update_theta_v.with_backend(self._backend)
self._compute_first_vertical_derivative = (
compute_first_vertical_derivative.with_backend(self._backend)
)
self._interpolate_to_half_levels_vp = (
interpolate_to_half_levels_vp.with_backend(self._backend)
)
self._interpolate_to_surface = (
interpolate_to_surface.with_backend(self._backend)
)
self._init_two_cell_kdim_fields_with_zero_vp = (
init_two_cell_kdim_fields_with_zero_vp.with_backend(self._backend)
)
Expand Down Expand Up @@ -504,6 +531,12 @@ def __init__(
self._backend
)
)
self._extrapolate_temporally_exner_pressure = (
extrapolate_temporally_exner_pressure.with_backend(self._backend)
)
self._set_theta_v_prime_ic_at_lower_boundary = (
set_theta_v_prime_ic_at_lower_boundary.with_backend(self._backend)
)
self._compute_hydrostatic_correction_term = (
compute_hydrostatic_correction_term.with_backend(self._backend)
)
Expand Down Expand Up @@ -582,6 +615,9 @@ def __init__(
self._init_cell_kdim_field_with_zero_wp = init_cell_kdim_field_with_zero_wp.with_backend(
self._backend
)
self._init_cell_kdim_field_with_zero_vp = init_cell_kdim_field_with_zero_vp.with_backend(
self._backend
)
self._update_mass_flux_weighted = update_mass_flux_weighted.with_backend(self._backend)
self._compute_z_raylfac = dycore_utils.compute_z_raylfac.with_backend(self._backend)
self._predictor_stencils_2_3 = nhsolve_stencils.predictor_stencils_2_3.with_backend(
Expand Down Expand Up @@ -994,7 +1030,20 @@ def run_predictor_step(
# - $\exnerprime{\n}{\c}{\k}$ : exner - exner_ref_mc
# - $\exnerprime{\n-1}{\c}{\k}$ : exner_pr
#
self._predictor_stencils_2_3(
#self._predictor_stencils_2_3(
# exner_exfac=self._metric_state_nonhydro.exner_exfac,
# exner=prognostic_state[nnow].exner,
# exner_ref_mc=self._metric_state_nonhydro.exner_ref_mc,
# exner_pr=diagnostic_state_nh.exner_pr,
# z_exner_ex_pr=self.z_exner_ex_pr,
# horizontal_start=self._start_cell_lateral_boundary_level_3,
# horizontal_end=self._end_cell_halo,
# vertical_start=0,
# vertical_end=self._grid.num_levels + 1,
# offset_provider={},
#)

self._extrapolate_temporally_exner_pressure(
exner_exfac=self._metric_state_nonhydro.exner_exfac,
exner=prognostic_state[nnow].exner,
exner_ref_mc=self._metric_state_nonhydro.exner_ref_mc,
Expand All @@ -1003,6 +1052,15 @@ def run_predictor_step(
horizontal_start=self._start_cell_lateral_boundary_level_3,
horizontal_end=self._end_cell_halo,
vertical_start=0,
vertical_end=self._grid.num_levels,
offset_provider={},
)

self._init_cell_kdim_field_with_zero_wp(
field_with_zero_wp=self.z_exner_ex_pr,
horizontal_start=self._start_cell_lateral_boundary_level_3,
horizontal_end=self._end_cell_halo,
vertical_start=self._grid.num_levels,
vertical_end=self._grid.num_levels + 1,
offset_provider={},
)
Expand Down Expand Up @@ -1033,17 +1091,50 @@ def run_predictor_step(
# - $\exnerprime{\ntilde}{\c}{\k}$ : z_exner_ex_pr
# - $1 / \Dz{\k}$ : inv_ddqz_z_full
#
self._predictor_stencils_4_5_6(
wgtfacq_c_dsl=self._metric_state_nonhydro.wgtfacq_c,
z_exner_ex_pr=self.z_exner_ex_pr,
z_exner_ic=self.z_exner_ic,
#self._predictor_stencils_4_5_6(
# wgtfacq_c_dsl=self._metric_state_nonhydro.wgtfacq_c,
# z_exner_ex_pr=self.z_exner_ex_pr,
# z_exner_ic=self.z_exner_ic,
# wgtfac_c=self._metric_state_nonhydro.wgtfac_c,
# inv_ddqz_z_full=self._metric_state_nonhydro.inv_ddqz_z_full,
# z_dexner_dz_c_1=self.z_dexner_dz_c_1,
# horizontal_start=self._start_cell_lateral_boundary_level_3,
# horizontal_end=self._end_cell_halo,
# vertical_start=max(1, self._vertical_params.nflatlev),
# vertical_end=self._grid.num_levels + 1,
# offset_provider=self._grid.offset_providers,
#)

self._interpolate_to_surface(
wgtfacq_c=self._metric_state_nonhydro.wgtfacq_c,
interpolant=self.z_exner_ex_pr,
interpolation_to_surface=self.z_exner_ic,
horizontal_start=self._start_cell_lateral_boundary_level_3,
horizontal_end=self._end_cell_halo,
vertical_start=self._grid.num_levels,
vertical_end=self._grid.num_levels + 1,
offset_provider=self._grid.offset_providers,
)

self._interpolate_to_half_levels_vp(
wgtfac_c=self._metric_state_nonhydro.wgtfac_c,
interpolant=self.z_exner_ex_pr,
interpolation_to_half_levels_vp=self.z_exner_ic,
horizontal_start=self._start_cell_lateral_boundary_level_3,
horizontal_end=self._end_cell_halo,
vertical_start=gtx.int32(max(1, self._vertical_params.nflatlev)),
vertical_end=self._grid.num_levels,
offset_provider=self._grid.offset_providers,
)

self._compute_first_vertical_derivative(
z_exner_ic=self.z_exner_ic,
inv_ddqz_z_full=self._metric_state_nonhydro.inv_ddqz_z_full,
z_dexner_dz_c_1=self.z_dexner_dz_c_1,
horizontal_start=self._start_cell_lateral_boundary_level_3,
horizontal_end=self._end_cell_halo,
vertical_start=max(1, self._vertical_params.nflatlev),
vertical_end=self._grid.num_levels + 1,
vertical_start=gtx.int32(max(1, self._vertical_params.nflatlev)),
vertical_end=self._grid.num_levels,
offset_provider=self._grid.offset_providers,
)

Expand Down Expand Up @@ -1076,17 +1167,24 @@ def run_predictor_step(
)

# Perturbation theta at top and surface levels
self._predictor_stencils_11_lower_upper(
wgtfacq_c_dsl=self._metric_state_nonhydro.wgtfacq_c,
self._init_cell_kdim_field_with_zero_vp(
field_with_zero_vp=self.z_theta_v_pr_ic,
horizontal_start=self._start_cell_lateral_boundary_level_3,
horizontal_end=self._end_cell_halo,
vertical_start=0,
vertical_end=1,
offset_provider=self._grid.offset_providers,
)

self._set_theta_v_prime_ic_at_lower_boundary(
wgtfacq_c=self._metric_state_nonhydro.wgtfacq_c,
z_rth_pr=self.z_rth_pr_2,
theta_ref_ic=self._metric_state_nonhydro.theta_ref_ic,
z_theta_v_pr_ic=self.z_theta_v_pr_ic,
theta_v_ic=diagnostic_state_nh.theta_v_ic,
k_field=self.k_field,
nlev=self._grid.num_levels,
horizontal_start=self._start_cell_lateral_boundary_level_3,
horizontal_end=self._end_cell_halo,
vertical_start=0,
vertical_start=self._grid.num_levels,
vertical_end=self._grid.num_levels + 1,
offset_provider=self._grid.offset_providers,
)
Expand Down Expand Up @@ -2147,7 +2245,7 @@ def run_corrector_step(
self._init_cell_kdim_field_with_zero_wp(
field_with_zero_wp=prep_adv.mass_flx_ic,
horizontal_start=self._start_cell_lateral_boundary,
horizontal_end=self._end_cell_nudging,
horizontal_end=self._end_cell_lateral_boundary_level_4,
vertical_start=0,
vertical_end=self._grid.num_levels + 1,
offset_provider={},
Expand All @@ -2163,8 +2261,8 @@ def run_corrector_step(
mass_flx_ic=prep_adv.mass_flx_ic,
r_nsubsteps=r_nsubsteps,
horizontal_start=self._start_cell_lateral_boundary,
horizontal_end=self._end_cell_nudging,
vertical_start=0,
horizontal_end=self._end_cell_lateral_boundary_level_4,
vertical_start=1,
vertical_end=self._grid.num_levels,
offset_provider={},
)
Expand Down
Loading