# GT4Py Vertical Computations

###### tags: `Felix’s remains` `IR`

## Idea: Stencils Operate on Whole Columns → Vertical Solvers are Stencils

# Some stencil functions working on a whole column

# Question: would it better to use another decorator for vertical stencils? Something like 'solver' or 'iterated_stencil' ?

@stencil(order=FORWARD)  
def forward(a, b, c, d):
    with interval(0, 1):
        cp = c / b
        dp = d / b
    with interval(1, None):
        cp = c / (b - a * cp[K - 1])
        dp = (d - a * dp[K - 1]) / (b - a * cp[K - 1])
    return cp, dp
        
@stencil(order=BACKWARD)
def backward(cp, dp):
    with interval(-1, None):
        x = dp
    with interval(0, -1):
        x = dp - cp * x[K + 1]
    return x
    

@stencil
def lap(inp):
    with interval(...):
        out = -4 * inp + inp[I - 1] + inp[I + 1] + inp[J - 1] + inp[J + 1]
    return out
    
    
# composition of stencil functions
   
@stencil
def solve_tridiag(a, b, c, d):
    cp, dp = forward(a, b, c, d)
    x = backward(cp, dp)
    return x
    
   
@stencil
def lap_of_solution(a, b, c, d):
    x = solve_tridiag(a, b, c, d)
    l = lap(x)
    return l
    
    
# splitting of vertical interval. Does this make sense?
@stencil
def partial_lap(inp):
    with interval(None, 3):
        out = lap(inp)
    with interval(3, None):
        out = inp
    return out
    
    
# simple fencil functions for testing components
    
@fencil
def everywhere_tridiag(a: InField, b: InField, c: InField, d: InField, x: OutField):
    with horizontal_region[:, :]:
        x = solve_tridiag(a, b, c)
        
@fencil
def everywhere_lap(inp: InField, out: OutField):
    with horizontal_region[:, :]:
        out = lap(inp)
        

@fencil
def everywhere_lap_of_solution(a: InField, b: InField, c: InField, d: InField, out: OutField):
    with horizontal_region[:, :]:
        out = lap_of_solution(a, b, c, d)
        

# Another possible way to test stencils
cp, dp = forward[100,100,10](a, b, c, d) # Allocate new output fields
foo[100,100,10](a, b, out=(cp, dp))      # Passing pre-allocatd outputs ala Numpy


# fancy fencil function including boundary conditions etc.
        
@fencil
def fancy_thing(a: InField, b: InField, c: InField, d: InField, out: OutField):
    with horizontal_region[1:-1, 1:-1]:
        out = lap_of_solution(a, b, c, d)
    with horizontal_region[:1, 1:-1]:
        x = solve_tridiag(a, b, c, d)
        # laplacian of solution with zero-gradient BC
        out = -4 * x + 2 * x[I + 1] + x[J + 1] + x[J - 1]
    with (horizontal_region[-1:, 1:-1] |
            horizontal_region[:, :1] | 
            horizontal_region[:, -1]):
        out = 0
    

## Example Fencil Fusion

import numpy as np

def lift(stencil):
    def lifted(*accessors):
        def accessor(**outer_offsets):
            def wrap(acc):
                def a(**offsets):
                    dims = set(outer_offsets) | set(offsets)
                    return acc(
                        **{
                            dim: outer_offsets.get(dim, 0) +
                            offsets.get(dim, 0)
                            for dim in dims
                        })
                return a
            return stencil(*(wrap(accessor) for accessor in accessors))
        return accessor
    return lifted

def apply_fencil(fencil, shape, *args):
    out = np.empty(shape)
    def create_accessor(array, index):
        def accessor(**offsets):
            if array.ndim == 1:
                return array[index + offsets.get("i", 0)]
            return array[tuple(i + offsets.get(a, 0)
                               for i, a in zip(index, "ijk"))]
        return accessor
    for stencil, domain in fencil:
        for index in domain:
            accessors = (create_accessor(arg, index) for arg in args)
            out[index] = stencil(*accessors)
    return out

def stencil(x):
    return x(i=-1) + x(i=+1)

def bc(x):
    return 0

fencil = [(bc, range(0, 1)), (stencil, range(1, 9)), (bc, range(9, 10))]
x = np.random.uniform(size=10)
y = apply_fencil(fencil, x.shape, x)
z = apply_fencil(fencil, y.shape, y)
print(z)

# Incorrect!
# composed_fencil = [(lambda x: bc(lift(bc)(x)), range(0, 1)),
#                    (lambda x: stencil(lift(stencil)(x)), range(1, 9)),
#                    (lambda x: bc(lift(bc)(x)), range(9, 10))]
# z2 = apply_fencil(composed_fencil, x.shape, x)
# print(z2)

# Correct
composed_fencil = [(lambda x: bc(lift(bc(x))), range(0, 1)),
                   # new stencil which is not a simple composition!
                   (lambda x: lift(bc)(x)(i=-1) + lift(stencil)(x)(i=+1), range(1, 2)),
                   (lambda x: stencil(lift(stencil)(x)), range(2, 8)),
                   # new stencil which is not a simple composition!
                   (lambda x: lift(stencil)(x)(i=-1) + lift(bc)(x)(i=+1), range(8, 9)),
                   (lambda x: bc(lift(bc)(x)), range(9, 10))]
z3 = apply_fencil(composed_fencil, x.shape, x)
print(z3)

Let us express bc composition in the previous example via stencil composition:

# let as add extra input field -- is_border : bool

def stencil(x):
    return x(i=-1) + x(i=+1)

def bc(x):
    return 0

def bc_stencil(x, is_border):
    return bc(x) if is_border else stencil(x)

def bc_stencil2(interval_id):
   def res(x) :
       return bc(x) if is_border(interval_id) else stencil(x)
   return res


# now lift(bc_stencil) works with expected results 

composed_fencil = [(
    lambda x, b : bc_stencil(lift(bc_stencil)(x, b), b),
    range(0, 10)
)]

### Outcome

Do not try to do composition of the fencil that reperesents `bc` directly. Instead we need to be able automatically generate the field with the computation area ID and a single stencil with the accessor to that filed as an additional parameter.