feat[cartesian]: numpy backend support for K iterator access

[FlorianDeconinck]

Description

Implement the K iterator access in stencil code for the numpy backend by using a pre-defined k_mask as a broadcastable reference to K.

The code

        with computation(PARALLEL), interval(...):
            if K == 2:
                field_A = 20.20
            field_B = float(K + offsets)

generates

def run(*, field_A, field_B, offsets, _domain_, _origin_):

    # --- begin domain boundary shortcuts ---
    _di_, _dj_, _dk_ = 0, 0, 0
    _dI_, _dJ_, _dK_ = _domain_
    k_mask = np.empty(_domain_, dtype=int)
    k_mask[:, :] = np.arange(_dK_, dtype=int)
    # --- end domain padding ---

    field_A = Field(field_A, _origin_["field_A"], (True, True, True))
    field_B = Field(field_B, _origin_["field_B"], (True, True, True))
    offsets = Field(offsets, _origin_["offsets"], (False, False, True))

    with np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore"):

        # --- begin vertical block ---
        k, K = _dk_, _dK_

        # --- begin horizontal block --
        i, I = _di_ - 0, _dI_ + 0
        j, J = _dj_ - 0, _dJ_ + 0

        field_A[i:I, j:J, k:K] = np.where(
            (k_mask[i:I, j:J, k:K] == np.int64(2)), np.float64(20.2), field_A[i:I, j:J, k:K]
        )
        field_B[i:I, j:J, k:K] = ufuncs.float64(
            (k_mask[i:I, j:J, k:K] + offsets[k:K].astype(np.int64))
        )
        # --- end horizontal block --

        # --- end vertical block ---

Requirements

• All fixes and/or new features come with corresponding tests.

[twicki]

This is very lean, thanks a lot!