diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py index 17202a8bc0..69b30128a6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py @@ -132,7 +132,7 @@ from icon4py.model.atmosphere.dycore.update_dynamical_exner_time_increment import ( update_dynamical_exner_time_increment, ) -from icon4py.model.atmosphere.dycore.update_mass_flux import update_mass_flux +from icon4py.model.atmosphere.dycore.update_mass_volume_flux import update_mass_volume_flux from icon4py.model.atmosphere.dycore.update_mass_flux_weighted import update_mass_flux_weighted from icon4py.model.atmosphere.dycore.update_theta_v import update_theta_v from icon4py.model.atmosphere.dycore.velocity.velocity_advection import VelocityAdvection @@ -1927,7 +1927,7 @@ def run_corrector_step( offset_provider={}, ) log.debug(f"corrector start stencil 58") - update_mass_flux.with_backend(backend)( + update_mass_volume_flux.with_backend(backend)( z_contr_w_fl_l=z_fields.z_contr_w_fl_l, rho_ic=diagnostic_state_nh.rho_ic, vwind_impl_wgt=self.metric_state_nonhydro.vwind_impl_wgt, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_volume_flux.py similarity index 77% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_volume_flux.py index 9c22b97e86..4c88ebeb71 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_volume_flux.py @@ -20,40 +20,45 @@ @field_operator -def _update_mass_flux( +def _update_mass_volume_flux( z_contr_w_fl_l: Field[[CellDim, KDim], wpfloat], rho_ic: Field[[CellDim, KDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], w: Field[[CellDim, KDim], wpfloat], mass_flx_ic: Field[[CellDim, KDim], wpfloat], + vol_flx_ic: Field[[CellDim, KDim], wpfloat], r_nsubsteps: wpfloat, -) -> Field[[CellDim, KDim], wpfloat]: +) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: """Formerly known as _mo_solve_nonhydro_stencil_58.""" - mass_flx_ic_wp = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) - return mass_flx_ic_wp + z_a = r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) + mass_flx_ic_wp = mass_flx_ic + z_a + vol_flx_ic_wp = vol_flx_ic + z_a / rho_ic + return mass_flx_ic_wp, vol_flx_ic_wp @program(grid_type=GridType.UNSTRUCTURED) -def update_mass_flux( +def update_mass_volume_flux( z_contr_w_fl_l: Field[[CellDim, KDim], wpfloat], rho_ic: Field[[CellDim, KDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], w: Field[[CellDim, KDim], wpfloat], mass_flx_ic: Field[[CellDim, KDim], wpfloat], + vol_flx_ic: Field[[CellDim, KDim], wpfloat], r_nsubsteps: wpfloat, horizontal_start: int32, horizontal_end: int32, vertical_start: int32, vertical_end: int32, ): - _update_mass_flux( + _update_mass_volume_flux( z_contr_w_fl_l, rho_ic, vwind_impl_wgt, w, mass_flx_ic, + vol_flx_ic, r_nsubsteps, - out=mass_flx_ic, + out=(mass_flx_ic, vol_flx_ic), domain={ CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end), diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_volume_flux.py similarity index 78% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_volume_flux.py index a556776547..f0084ac11d 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_volume_flux.py @@ -15,15 +15,18 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.update_mass_flux import update_mass_flux +from icon4py.model.atmosphere.dycore.update_mass_volume_flux import update_mass_volume_flux from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field from icon4py.model.common.type_alias import wpfloat class TestMoSolveNonhydroStencil58(StencilTest): - PROGRAM = update_mass_flux - OUTPUTS = ("mass_flx_ic",) + PROGRAM = update_mass_volume_flux + OUTPUTS = ( + "mass_flx_ic", + "vol_flx_ic", + ) @staticmethod def reference( @@ -33,12 +36,15 @@ def reference( vwind_impl_wgt: np.array, w: np.array, mass_flx_ic: np.array, + vol_flx_ic: np.array, r_nsubsteps, **kwargs, ) -> dict: vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) - mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) - return dict(mass_flx_ic=mass_flx_ic) + z_a = r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) + mass_flx_ic = mass_flx_ic + z_a + vol_flx_ic = vol_flx_ic + z_a / rho_ic + return dict(mass_flx_ic=mass_flx_ic, vol_flx_ic=vol_flx_ic) @pytest.fixture def input_data(self, grid): @@ -47,6 +53,7 @@ def input_data(self, grid): vwind_impl_wgt = random_field(grid, CellDim, dtype=wpfloat) w = random_field(grid, CellDim, KDim, dtype=wpfloat) mass_flx_ic = random_field(grid, CellDim, KDim, dtype=wpfloat) + vol_flx_ic = random_field(grid, CellDim, KDim, dtype=wpfloat) r_nsubsteps = 7.0 return dict( @@ -55,6 +62,7 @@ def input_data(self, grid): vwind_impl_wgt=vwind_impl_wgt, w=w, mass_flx_ic=mass_flx_ic, + vol_flx_ic=vol_flx_ic, r_nsubsteps=r_nsubsteps, horizontal_start=int32(0), horizontal_end=int32(grid.num_cells),