# Graupel scheme

###### tags: `functional cycle 12`

shaped by:

## Goals

Full implementation of the Graupel scheme with integration into ICON Fortran. Additionally, GT4Py is extend by features. Workarounds should be used until the features become available.

### Not (yet) goals:
- Performance
- Code beauty within the boundaries of GT4Py (and possible improvements realized in this project)

## Appetite

Proposal: full cycle

## Solution

### Port the scheme

Expand https://github.com/leuty/graupel_standalone/tree/graupel_stub to the full scheme, but only the path that is active in Aquaplanet and MCH production (as discussed in the EXCLAIM developer meeting).

Implementation hints: regularly run the gtfn codegenerator on the partially implemented scheme to detect early if transformation passes can't handle certain constructs.

### Field view extensions and bug fixes

In general, there might still be bugs for edge-cases that are triggered by this scheme. Hot fixes should be implemented with high-priority if no work-arounds exist. It might be convenient to keep a special branch with all required features, if they are implemented, but not merged to the main GT4Py branch.

#### For/while loop: Peter (WIP)

Description: Support point-wise/per-grid-point loops.

Workaround: Use manually unrolled loop.

Alternatives are the proposed `scalar_operator` or the `functional for-loop` discussed in the following.

##### Functional while/for loop

Implement the [jax while_loop](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.while_loop.html).

The simplified satad loop would look like

def cond(x: tuple[Field[...], Field[...],...]): # `...` is used as placeholder, not Python elipses
    counter, maxiter, twork, tworkold, tol = x
    return count < maxiter and twork-tworkold>tol
def body(x: ...):
    counter, maxiter, twork, tworkold, tol = x
    return (counter+1, maxiter, compute_twork_next(...), twork, tol)

def newtonian_loop(...):
    return while_loop(cond, body, (0, 10, twork_init, twork_init+10., 1e-3))

*Steps*

- Implement a `while_loop(cond, body, init)` builtin on iterator level, where `cond` and `body` are functions and `init` is a tuple of values (see [jax while_loop](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.while_loop.html) for details on semantics). First in the embedded backend and later also for the C++ backend. The C++ backend tasks can be done independently of the frontend tasks.
- Implement a `while_loop` builtin in the frontend, where `cond` and `body` are two field operators and `init` is again a tuple of values (may be fields).

    @field_operator
    def cond(counter: int, maxiter: int, twork: Field[...], tworkold: Field[...], tol: float) -> Field[..., bool]:
        return count < maxiter and (twork - tworkold) > tol

    @field_operator
    def body(counter: int, maxiter: int, twork: Field[...], tworkold: Field[...], tol: float) -> tuple[int, int, Field[...], Field[...], float]:
        twork_next = ...
        return (counter+1, maxiter, twork_next, twork, tol)

    @field_operator
    def newtonian_loop(twork_init: Field[...]) -> Field[...]:
        return while_loop(cond, body, (0, 10, twork_init, Inf, 1e-3))

Make sure that the example also works for `twork` being a scalar and `newtonian_loop` called from a scan operator (at least for embedded).
- Improvements: Extend the Parser (`func_to_foast.py`) and lowering (`foast_to_itir.py`) to support nested functions inside of field operators. The current infrastracture should support this without major changes. Nested functions are lowered using the existing `visit_FunctionDefinition` and each definition will - similarly to assignments - become a let statement with the `init-form` being a lambda function.

    @field_operator
    def newtonian_loop(twork_init: float, maxiter: int, tol: float) -> float:
        def cond(counter: int, twork: Field[...], tworkold: Field[...]) -> Field[..., bool]:
            return count < maxiter and (twork - tworkold) > tol

        def body(counter: int, twork: Field[...], tworkold: Field[...]):
            twork_next = ...
            return (counter+1, maxiter, twork_next, twork, tol)
    
        return while_loop(cond, body, (0, twork_init, Inf))
    
    newtonian_loop(twork_init, 10, 1e-3)

- Extend the quickstart guide with a meaningful, simple example. For example an implementation of a [Newton's method](https://en.wikipedia.org/wiki/Newton%27s_method) computing the zeros of function which is parameterized by what is `twork` above.

Tests on all levels required, e.g. tests starting directly from ITIR.

#### Cast --> became astype - Nikki (PR merged)

Description: introduce a `cast` builtin for casting between datatypes.

##### Syntax

**field view**: TODO decide with Enrique `field.astype(dtype)` + `dtype(scalar)` or `cast(field, dtype)` (and the other as syntax sugar later)

**iterator IR**: `cast(dtype)(value_or_column)`, where `dtype` is a string

Final outcome:
- Frontend: astype(obj, new_dtype)
- Backend: CastExpr(obj, new_dtype)

##### Steps

- introduce builtin in field view and iterator IR
- introduce builtin in `embedded` and `tracing` of iterator IR
- extend `gtfn` code generator
- needs to be tested at all levels

#### Implement bitwise operators: Nikki (PR merged)

Negation through the use of `tilde` (~)

def bitwise_negation(cond: boolean, inp0: Field, inp1: Field):
    result = inp0 if ~cond else inp1
    return result

Lower to `not`.

Implementationg of XOR with use of `apostrophe` (^)

def bitwise_xor(cond1: boolean, cond2: boolean, inp0: Field, inp1: Field):
    result = inp0 if cond1^cond2 else inp1
    return result

Add xor operation in itir.

#### More math functions - Nikki (PR merged)

Added front-end tests

Implement all functions from Cartesian GT4Py https://github.com/GridTools/gt4py/blob/adec0c3bf36fd2660ee766ec38114733672150ad/src/gt4py/gtscript.py#L39

#### index()-builtin

Description: Add an `index(Dim)` builtin which returns the index of the current iteration point.

##### Syntax proposal

@field_operator
def foo():
    return index(I)

returns a field of `I` indices from `start` to (including) `end-1` if the domain is `I:{start,end}`.

##### ITIR/backend implementation

TODO: see other document, where the index iterator is introduced.

#### Other features currently under development

The following other features are already available as PRs:
- scalar if
- passing scalars to scan
- calling functions from scan

#### GT4Py features not targetted for this project

Reductions in the vertical (aka "optimized returning CellDim field"). Workaround is to propagate the last level back (backward scan) or read the last level of a scan outside of the DSL.

### ICON Fortran integration

Try if icon4pygen integration is feasible for this kind of program, otherwise use custom bindings generated similar to https://github.com/GridTools/gt4py/pull/608 (Hannes will support the implementation).

## Rabbit holes

Potentially all new GT4Py features are blockers, therefore the implementation of the scheme should be pushed forward using available workarounds.