[next] Bug in embedded implementation of premap

[egparedes]

This issue has been originally discovered by @nfarabullini . I'm only reporting it there since she is on vacation.

In a new Icon4py PR (https://github.com/C2SM/icon4py/pull/472/files) the field-vire embedded implementation of a premap operation crashes for unclear reasons since it was working previously for iterator embedded, therefore this is most likely a bug.

This is the definition of the operator crashing:

C2E2CODim = Dimension("C2E2CO", DimensionKind.LOCAL)
C2E2CO = FieldOffset("C2E2CO", source=CellDim, target=(CellDim, C2E2CODim))

@field_operator
def _compute_weighted_cell_neighbor_sum(
    field: Field[[CellDim, KDim], wpfloat],
    c_bln_avg: Field[[CellDim, C2E2CODim], wpfloat],
) -> Field[[CellDim, KDim], wpfloat]:
    field_avg = neighbor_sum(field(C2E2CO) * c_bln_avg, axis=C2E2CODim)
    return field_avg

[havogt]

Reading again the description of the implementation, it seems we overlooked the case X2X for a Field[Dims[X,Y],...]. In this case the domain of the connectivity is partially (but not fully contained in the domain of the field).

[havogt]

This line

gt4py/src/gt4py/next/embedded/nd_array_field.py

Line 497 in 9d1e4e9

if self.domain.dim_index(self.codomain) is None

is problematic.
It makes the above case a ALTER_STRUCT only. If we would implement the behavior that is described, we should raise an exception here. But that wouldn't solve the issue. Let's re-discuss how the above case fits in the categorization of connectivities.
Side-note: If I force ALTER_STRUCT | ALTER_DIMS on the connectivity the example works.

[jcanton]

@nfarabullini and I think this issue may have appeared again in this PR to port z_ifc and topography smoothing in icon4py (https://github.com/C2SM/icon4py/pull/564/files)

The stencil in question is nabla2_scalar in topography.py with its tests in test_topography.py
The tests with StencilTest and artificial geofac_n2s fields pass, but when trying to use a "real" geofac_n2s (either from the interpolation savepoint or by computing it - see the two tests below) the following error pops up:

>           topo.nabla2_scalar(
                psi_c=psi_c,
                geofac_n2s=geofac_n2s,
                nabla2_psi_c=nabla2_psi_c,
                offset_provider={"C2E2CO":grid.get_offset_provider("C2E2CO"),}
            )

model/common/tests/grid_tests/test_topography.py:43: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
_external_src/gt4py/src/gt4py/next/ffront/decorator.py:241: in __call__
    ctx.run(self.definition_stage.definition, *rewritten_args, **kwargs)
model/common/src/icon4py/model/common/grid/topography.py:36: in nabla2_scalar
    _nabla2_scalar(
_external_src/gt4py/src/gt4py/next/ffront/decorator.py:629: in __call__
    return embedded_operators.field_operator_call(op, args, kwargs)
_external_src/gt4py/src/gt4py/next/embedded/operators.py:116: in field_operator_call
    res = ctx.run(op, *args, **kwargs)
_external_src/gt4py/src/gt4py/next/embedded/operators.py:29: in __call__
    return self.fun(*args, **kwargs)
model/common/src/icon4py/model/common/grid/topography.py:26: in _nabla2_scalar
    nabla2_psi_c = neighbor_sum(psi_c(C2E2CO) * geofac_n2s, axis=C2E2CODim)
_external_src/gt4py/src/gt4py/next/embedded/nd_array_field.py:323: in __call__
    return functools.reduce(
_external_src/gt4py/src/gt4py/next/embedded/nd_array_field.py:324: in <lambda>
    lambda field, current_index_field: field.premap(current_index_field),
_external_src/gt4py/src/gt4py/next/embedded/nd_array_field.py:313: in premap
    return _reshuffling_premap(self, *cast(list[NdArrayConnectivityField], conn_fields))
_external_src/gt4py/src/gt4py/next/embedded/nd_array_field.py:632: in _reshuffling_premap
    conn_ndarray = xp.transpose(conn_ndarray, transposed_axes)
.venv/lib/python3.10/site-packages/numpy/_core/fromnumeric.py:669: in transpose
    return _wrapfunc(a, 'transpose', axes)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

obj = array([[[15],
        [ 4],
        [ 3],
        [ 0]],

       [[16],
        [ 5],
        [ 4],
        [ 1]],

  ...13],
        [ 1],
        [12],
        [16]],

       [[14],
        [ 2],
        [13],
        [17]]], dtype=int32)
method = 'transpose', args = ([0, 2],), kwds = {}, bound = <built-in method transpose of numpy.ndarray object at 0x1287baaf0>

    def _wrapfunc(obj, method, *args, **kwds):
        bound = getattr(obj, method, None)
        if bound is None:
            return _wrapit(obj, method, *args, **kwds)
    
        try:
>           return bound(*args, **kwds)
E           ValueError: axes don't match array

.venv/lib/python3.10/site-packages/numpy/_core/fromnumeric.py:57: ValueError

[nfarabullini]

the same error occurs in this PR: C2SM/icon4py#567

>       solve_nonhydro.time_step(
            diagnostic_state_nh=diagnostic_state_nh,
            prognostic_state_ls=prognostic_state_ls,
            prep_adv=prep_adv,
            divdamp_fac_o2=initial_divdamp_fac,
            dtime=dtime,
            l_recompute=recompute,
            l_init=linit,
            nnew=nnew,
            nnow=nnow,
            lclean_mflx=clean_mflx,
            lprep_adv=lprep_adv,
            at_first_substep=jstep_init == 0,
            at_last_substep=jstep_init == (ndyn_substeps - 1),
        )
atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py:800: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py:864: in time_step
    self.run_predictor_step(
atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py:956: in run_predictor_step
    self.velocity_advection.run_predictor_step(
atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py:228: in run_predictor_step
    self.compute_tangential_wind(
.tox/run_model_tests/src/gt4py/src/gt4py/next/ffront/decorator.py:245: in __call__
    ctx.run(self.definition_stage.definition, *args, **kwargs)
atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_tangential_wind.py:38: in compute_tangential_wind
    _compute_tangential_wind(
.tox/run_model_tests/src/gt4py/src/gt4py/next/ffront/decorator.py:624: in __call__
    return embedded_operators.field_operator_call(op, args, kwargs)
.tox/run_model_tests/src/gt4py/src/gt4py/next/embedded/operators.py:116: in field_operator_call
    res = ctx.run(op, *args, **kwargs)
.tox/run_model_tests/src/gt4py/src/gt4py/next/embedded/operators.py:29: in __call__
    return self.fun(*args, **kwargs)
atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_tangential_wind.py:24: in _compute_tangential_wind
    vt_wp = neighbor_sum(rbf_vec_coeff_e * vn(E2C2E), axis=E2C2EDim)
.tox/run_model_tests/src/gt4py/src/gt4py/next/embedded/nd_array_field.py:323: in __call__
    return functools.reduce(
.tox/run_model_tests/src/gt4py/src/gt4py/next/embedded/nd_array_field.py:324: in <lambda>
    lambda field, current_index_field: field.premap(current_index_field),
.tox/run_model_tests/src/gt4py/src/gt4py/next/embedded/nd_array_field.py:313: in premap
    return _reshuffling_premap(self, *cast(list[NdArrayConnectivityField], conn_fields))
.tox/run_model_tests/src/gt4py/src/gt4py/next/embedded/nd_array_field.py:632: in _reshuffling_premap
    conn_ndarray = xp.transpose(conn_ndarray, transposed_axes)
.tox/run_model_tests/lib/python3.10/site-packages/numpy/_core/fromnumeric.py:669: in transpose
    return _wrapfunc(a, 'transpose', axes)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
obj = array([[[   -1],
        [   -1],
        [   -1],
        [   -1]],
       [[   -1],
        [   -1],
        [   -1...
        [31557],
        [31555]],
       [[ 8727],
        [18767],
        [31555],
        [31556]]], dtype=int32)
method = 'transpose', args = ([0, 2],), kwds = {}
bound = <built-in method transpose of numpy.ndarray object at 0x7fff4c537570>
    def _wrapfunc(obj, method, *args, **kwds):
        bound = getattr(obj, method, None)
        if bound is None:
            return _wrapit(obj, method, *args, **kwds)
    
        try:
>           return bound(*args, **kwds)
E           ValueError: axes don't match array
.tox/run_model_tests/lib/python3.10/site-packages/numpy/_core/fromnumeric.py:57: ValueError
------------------------------ Captured log setup ------------------------------