# ICON stencils in field view examples

## Diffusion stencil 5

### Dusk version

from dusk.script import *

lb, nudging, interior, halo, end = HorizontalDomains(0, 1000, 2000, 3000, 4000)

@stencil
def mo_nh_diffusion_stencil_05(
    u_vert: Field[Vertex, K],
    v_vert: Field[Vertex, K],
    primal_normal_vert_v1: Field[Edge > Cell > Vertex],
    primal_normal_vert_v2: Field[Edge > Cell > Vertex],
    z_nabla2_e: Field[Edge, K],
    inv_vert_vert_length: Field[Edge],
    inv_primal_edge_length: Field[Edge],
    z_nabla4_e2: Field[Edge, K]
):
    nabv_tang: Field[Edge, K]
    nabv_norm: Field[Edge, K]
    with domain.upward.across[nudging:halo]:
        nabv_tang = sum_over(
            Edge > Cell > Vertex,
            u_vert * primal_normal_vert_v1 + v_vert * primal_normal_vert_v2,
            weights=[1.0, 1.0, 0.0, 0.0],
        )
        nabv_norm = sum_over(
            Edge > Cell > Vertex,
            u_vert * primal_normal_vert_v1 + v_vert * primal_normal_vert_v2,
            weights=[0.0, 0.0, 1.0, 1.0],
        )
        z_nabla4_e2 = 4.0 * (
            (nabv_norm - 2.0 * z_nabla2_e) * inv_vert_vert_length ** 2
            + (nabv_tang - 2.0 * z_nabla2_e) * inv_primal_edge_length ** 2
        )

### Field view with functions

import gt4py as gt

domain = {K: range(0,65), Edge: range(nudging, halo)}

@stencil
def vn_vert(
    weights: Field[Diamond],
    u_vert: Field[Vertex, K],
    v_vert: Field[Vertex, K],
    primal_normal_vert_v1: Field[Edge, Diamond],
    primal_normal_vert_v2: Field[Edge, Diamond]
) -> Field[Edge, Diamond]:

    return gt.sum(
            (
                u_vert(Diamond) * primal_normal_vert_v1
                + v_vert(Diamond) * primal_normal_vert_v2
            )
            * weights, # TODO need to discuss if we implicitly broadcast here or not
            axis=Diamond,
        )

@stencil
def nabv_tang(
    u_vert: Field[Vertex, K],
    v_vert: Field[Vertex, K],
    primal_normal_vert_v1: Field[Edge, Diamond],
    primal_normal_vert_v2: Field[Edge, Diamond]
) -> Field[Edge, K]:

    weights: Field[Diamond] = gt.weights([[1.0, 1.0, 0.0, 0.0]])
    return vn_vert(weights, u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2)

@stencil
def nabv_norm(
    u_vert: Field[Vertex, K],
    v_vert: Field[Vertex, K],
    primal_normal_vert_v1: Field[Edge, Diamond],
    primal_normal_vert_v2: Field[Edge, Diamond]
) -> Field[Edge, K]:

    weights: Field[Diamond] = gt.weights([[0.0, 0.0, 1.0, 1.0]])
    return vn_vert(weights, u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2)

@stencil
def z_nabla4_e2(
    nabv_norm: Field[Edge, K],
    nabv_tang: Field[Edge, K],
    z_nabla2_e: Field[Edge, K],
    inv_vert_vert_length: Field[Edge],
    inv_primal_edge_length: Field[Edge]
) -> Field[Edge, K]:

    return 4.0 * (
        (nabv_norm - 2.0 * z_nabla2_e) * inv_vert_vert_length ** 2
        + (nabv_tang - 2.0 * z_nabla2_e) * inv_primal_edge_length ** 2
    )

@fencil
def mo_nh_diffusion_stencil_05(
    u_vert: Field[Vertex, K],
    v_vert: Field[Vertex, K],
    primal_normal_vert_v1: Field[Edge, Diamond],
    primal_normal_vert_v2: Field[Edge, Diamond],
    z_nabla2_e: Field[Edge, K],
    inv_vert_vert_length: Field[Edge],
    inv_primal_edge_length: Field[Edge],
    z_nabla4_e2: Field[Edge, K]
) -> Field[Edge, K]:

    nabv_tang = nabv_tang(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2)
    nabv_norm = nabv_norm(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2)
    z_nabla4_e2 = z_nabla4_e2(
        nabv_norm, nabv_tang, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length
    )

    return z_nabla4_e2[domain]

### Field view closest to dusk

Diamond = gt.ICON.nbhdChain(Edge, Cell, Vertex)
# Hannes: currently
Diamond = offset("Diamond")

@stencil
def mo_nh_diffusion_stencil_05(
    u_vert: Field[Vertex, K],
    v_vert: Field[Vertex, K],
    primal_normal_vert_v1: Field[Diamond],
    primal_normal_vert_v2: Field[Diamond],
    z_nabla2_e: Field[Edge, K],
    inv_vert_vert_length: Field[Edge],
    inv_primal_edge_length: Field[Edge],
):
    nabv_tang = sum(
        (
            u_vert(Diamond) * primal_normal_vert_v1
            + v_vert(Diamond) * primal_normal_vert_v2
        )
        * Field(Diamond)([1.0, 1.0, 0.0, 0.0]),
        axis=Diamond,
    )
    nabv_norm = sum(
        (
            u_vert(Diamond) * primal_normal_vert_v1
            + v_vert(Diamond) * primal_normal_vert_v2
        )
        * (Field(Diamond)([0.0, 0.0, 1.0, 1.0])),
        axis=Diamond,
    )
    z_nabla4_e2 = 4.0 * (
        (nabv_norm - 2.0 * z_nabla2_e) * inv_vert_vert_length ** 2
        + (nabv_tang - 2.0 * z_nabla2_e) * inv_primal_edge_length ** 2
    )
    return z_nabla4_e2


def main(
    k_size,
    nudging,
    halo,
    u_vert,
    v_vert,
    primal_normal_vert_v1,
    primal_normal_vert_v2,
    z_nabla2_e,
    inv_vert_vert_length,
    inv_primal_edge_length,
    z_nabla4_e2,
):
    domain = {K: range(0, k_size), Edge: range(nudging, halo)}
    mo_nh_diffusion_stencil_05[domain](
        u_vert,
        v_vert,
        primal_normal_vert_v1,
        primal_normal_vert_v2,
        z_nabla2_e,
        inv_vert_vert_length,
        inv_primal_edge_length,
        out=z_nabla4_e2,
    )

#### Alternative reduction syntax
This is a sketch and we can discuss if we can improve.

    nabv_tang = sum(
        (
            u_vert(v) * primal_normal_vert_v1[i]
            + v_vert(v) * primal_normal_vert_v2[i]
        )
        * [1.0, 1.0, 0.0, 0.0][i] for i, v in enumerate(EdgeCellVertex)
    )

## Diffusion stencil 14
### Dusk version

from dusk.script import *


lb, nudging, interior, halo, end = HorizontalDomains(0, 1000, 2000, 3000, 4000)
thresh_tdiff = Global("thresh_tdiff")


@stencil
def mo_nh_diffusion_stencil_14(
    theta_v: Field[Cell, K],
    theta_ref_mc: Field[Cell, K],
    enh_diffu_3d: Field[Cell, K],
):
    tdiff: Field[Cell, K]
    trefdiff: Field[Cell, K]

    with domain.upward[-2:].across[nudging-1:halo-1]:

        tdiff = theta_v - sum_over(Cell > Edge > Cell, theta_v[Cell > Edge > Cell]) / 3

        trefdiff = (
            theta_ref_mc
            - sum_over(Cell > Edge > Cell, theta_ref_mc[Cell > Edge > Cell]) / 3
        )

        if (
            tdiff - trefdiff < thresh_tdiff
            and trefdiff < 0
            # or tdiff - trefdiff < 1.5 * thresh_tdiff MR: can't find this conditional in the ICON code!?
        ):
            enh_diffu_3d = (thresh_tdiff - tdiff + trefdiff) * 5e-4
        else:
            enh_diffu_3d = -1.7976931348623157e+308

### Field view closest to dusk
As a test use full type notation here, including `float`.

Triforce = gt.ICON.nbhdChain(Cell, Edge, Cell)
# Hannes: currently
Triforce = offset("Triforce")

@stencil
def mo_nh_diffusion_stencil_14(
    theta_v: Field[[Cell, K], float],
    theta_ref_mc: Field[[Cell, K], float],
    thresh_tdiff: float,
    enh_diffu_3d: Field[[Cell, K], float],
):
    tdiff = theta_v - sum(theta_v[Triforce], axis=Triforce) / 3
    trefdiff = theta_ref_mc - sum(theta_ref_mc[Triforce], axis=Triforce) / 3

    enh_diff_3d = gt.ternary(
        tdiff - trefdiff < thresh_tdiff & trefdiff < 0,
        (thresh_tdiff - tdiff + trefdiff) * 5e-4,
        -1.7976931348623157e+308,
    )
    
    return enh_diffu_3d

def main(
    k_size,
    nudging,
    halo,
    theta_v,
    theta_ref_mc,
    thresh_tdiff,
    enh_diffu_3d,
):
    domain = {K: range(k_size - 2, k_size), Cell: range(nudging - 1, halo - 1)}
    mo_nh_diffusion_stencil_14[domain](
        theta_v,
        theta_ref_mc,
        thresh_tdiff,
        out=enh_diffu_3d,
    )

## Diffusion stencil 18
### Dusk version

from dusk.script import *
from dusk.test import stencil_test

lb, nudging, interior, halo, end = HorizontalDomains(0, 1000, 2000, 3000, 4000)


@stencil
def mo_nh_diffusion_stencil_18(
    # This mask field is not in the original Fortran, since it builds a generic index list that contains
    # the result of masking the entire domain
    mask: Field[Cell, K],
    zd_vertidx: IndexField[Origin+Cell > Edge > Cell, K],
    zd_vertidx_p1: IndexField[Origin+Cell > Edge > Cell, K],
    zd_diffcoef: Field[Cell, K],
    geofac_n2s: Field[Origin + Cell > Edge > Cell],
    # In fortran vcoef does not contain center. It will require copying it into a new sparse field
    # that contains center with value 1, before calling codegen.
    vcoef: Field[Origin+Cell > Edge > Cell, K],
    theta_v: Field[Cell, K],
    z_temp: Field[Cell, K]):

    with domain.upward.across[nudging:end]:
        if(mask):
            z_temp += zd_diffcoef * \
                sum_over(Origin+Cell > Edge > Cell, geofac_n2s *
                        (vcoef*theta_v[Origin+Cell > Edge > Cell, zd_vertidx] +
                        (1-vcoef)*theta_v[Origin+Cell > Edge > Cell, zd_vertidx_p1 + 1]))

### Field view closest to dusk
As a test use full type notation here, including `float`.

Triforce = gt.ICON.nbhdChain(Cell, Edge, Cell)
# Hannes: currently
Triforce = offset("Triforce")

@stencil
def mo_nh_diffusion_stencil_18(
    mask: Field[[Cell, K], bool],
    zd_vertidx: Field[[Triforce, K], int],
    zd_vertidx_p1: Field[[Triforce, K], int],
    zd_diffcoef: Field[[Cell, K], float],
    geofac_n2s: Field[[Triforce], float],
    vcoef: Field[[Triforce, K], float],
    theta_v: Field[[Cell, K], float],
    z_temp: Field[[Cell, K], float],
):
    z_temp = gt.where(
        mask,
        z_temp +
        zd_diffcoef *
        sum(
            geofac_n2s * (vcoef * theta_v[Triforce, zd_vertidx] + (1 - vcoef) * theta_v[Triforce, zd_vertidx_p1 + 1]),
            axis=Triforce
        ),
        z_temp
    )
    
    return z_temp

def main(
    k_size,
    nudging,
    end,
    mask,
    zd_vertidx,
    zd_vertidx_p1,
    zd_diffcoef,
    geofac_n2s,
    vcoef,
    theta_v,
    z_temp,
):
    domain = {K: range(0, k_size), Cell: range(nudging, end)}
    mo_nh_diffusion_stencil_18[domain](
        mask,
        zd_vertidx,
        zd_vertidx_p1,
        zd_diffcoef,
        geofac_n2s,
        vcoef,
        theta_v,
        out=z_temp,
    )