gistfile1.py

@stencil
def flux_x(phi):
    return flux(phi[I - 1], phi)

@stencil
def flux_y(phi):
    return flux(phi[J - 1], phi)

@stencil
def step(phi, fx, fy, dx, dy, dt):
    return phi + dt * ((fx[I + 1] - fx) / dx + (fy[J + 1] - fy) / dy), fx, fy

@fencil
def f(phi_in, phi_out, fx_out, fy_out, dx, dy, dt, flx_domain, fly_domain, phi_domain):
    lazy_apply_stencil(flux_x, flx_domain, [fx_out], [phi_in])
    lazy_apply_stencil(flux_y, fly_domain, [fy_out], [phi_in])
    apply_stencil(step, phi_domain, [phi_out], [phi_in, fx_out, fy_out])

    fx = flux_x(phi_in)
    apply_expression(fx, flx_domain, [fx_out])
    # ...
    apply_stencil(step, phi_domain, [phi_out], [phi_in, fx, fy])






@stencil
def step(phi, dx, dy, dt):
    fx = flux_x(phi)  # lift() required in original/non-Hannes model
    fy = flux_y(phi)  # lift() required in original/non-Hannes model
    return phi + dt * ((fx[I + 1] - fx) / dx + (fy[J + 1] - fy) / dy), fx, fy

@fencil
def f(phi_in, phi_out, fx_out, fy_out, dx, dy, dt, flx_domain, fly_domain, phi_domain, edge_domain):
    apply_stencil(step, (phi_domain, flx_domain, fly_domain), [phi_out, fx, fy], [phi_in])

###

@stencil
def step(phi, dx, dy, dt):
    fx = flux_x(phi)  # lift() required in original/non-Hannes model
    fy = flux_y(phi)  # lift() required in original/non-Hannes model
    return phi + dt * ((fx[I + 1] - fx) / dx + (fy[J + 1] - fy) / dy)

@fencil
def f(phi_in, phi_out, fx_out, fy_out, dx, dy, dt, flx_domain, fly_domain, phi_domain):
    apply_stencil(flux_x, fx_domain, [fx_out], [phi_in])
    apply_stencil(flux_y, fy_domain, [fy_out], [phi_in])
    apply_stencil(step, phi_domain, [phi_out], [phi_in])