Iterator Example

iterator_example.py

from __future__ import annotations
import numpy as np
from functional.iterator.embedded import MDIterator, np_as_located_field, lift, shift, deref, NeighborTableOffsetProvider
from fvm_advection_sphere.common import *

# v-e-v-e-v
#   c   c
num_vertices = 10
num_edges = num_vertices-1
num_cells = num_vertices-1

e2v_conn = NeighborTableOffsetProvider(np.asarray([[i, i+1] for i in range(num_edges)]), Edge, Vertex, 2)
c2e_conn = NeighborTableOffsetProvider(np.asarray([[i] for i in range(num_cells)]), Cell, Edge, 1)
offset_provider = {
    "E2V": e2v_conn,
    "C2E": c2e_conn
}

vertex_field = np_as_located_field(Vertex)(np.arange(0., 10., 1.))

cell_it: Iterator[[Cell], [Vertex]] = MDIterator(vertex_field, {"Cell": 1}, offset_provider=offset_provider)
vertex_it: Iterator[[Vertex], [Vertex]] = shift("E2V", 1)(shift("C2E", 0)(cell_it))
val = deref(vertex_it)  # == 2.0

def edge_average(edge_it: Iterator[[Edge], [Vertex]]) -> float:
    """Stencil mapping from iterators to values"""
    return (deref(shift("E2V", 0)(edge_it))+deref(shift("E2V", 1)(edge_it)))/2.

# lift rules:
#  shift(...)(lift(stencil)(it1, it2)) == lift(stencil)(shift(...)(it1), shift(...)(it2))
#  deref(lift(stencil)(it1, it2)) == stencil(it1, it2)
edge_average_it: Iterator[[Cell], [Edge]] = lift(edge_average)(cell_it)

edge_average = deref(shift("C2E", 0)(edge_average_it))  # == 1.5

breakpoint()