Test strong vertical integral with pyop2

test_vert_int.py

from thetis import *

"""
Computes a "strong" vertical integral of a field with pyop2
by integrating along the vertical 1D edges of the prism.
"""


def strong_integral_2d(source, output_2d, z_coords):
    """Computes strong integral of a P1DGxP1DG field onto P1DG 2D field"""
    fs_3d = source.function_space()
    fs_2d = output_2d.function_space()
    # number of nodes in vertical direction
    n_vert_nodes = len(fs_3d.fiat_element.B.entity_closure_dofs()[1][0])

    nodes = fs_3d.bt_masks['geometric'][0]
    idx = op2.Global(len(nodes), nodes, dtype=np.int32, name='node_idx')
    kernel = op2.Kernel("""
        void my_kernel(double **func, double **func2d, double **z, int *idx) {
            for ( int d = 0; d < %(nodes)d; d++ ) {
                double h = z[idx[d]+1][0] - z[idx[d]+0][0];
                for ( int c = 0; c < %(func_dim)d; c++ ) {
                    for ( int e = 0; e < %(v_nodes)d; e++ ) {
                        func2d[d][c] += 0.5*h*func[idx[d]+e][c];
                    }
                }
            }
        }""" % {'nodes': output_2d.cell_node_map().arity,
                'func_dim': output_2d.function_space().dim,
                'v_nodes': n_vert_nodes},
        'my_kernel')

    op2.par_loop(
        kernel, fs_3d.mesh().cell_set,
        source.dat(op2.READ, fs_3d.cell_node_map()),
        output_2d.dat(op2.INC, fs_2d.cell_node_map()),
        z_coords.dat(op2.READ, z_coords.function_space().cell_node_map()),
        idx(op2.READ),
        iterate=op2.ALL)


def strong_integral_3d_p1(source, output_3d, z_coords):
    """Computes strong integral of a P1DGxP1DG onto P1DGxP1 field"""
    fs_3d = source.function_space()
    fs_out = output_3d.function_space()
    # number of nodes in vertical direction
    n_vert_nodes = len(fs_3d.fiat_element.B.entity_closure_dofs()[1][0])

    nodes = fs_3d.bt_masks['geometric'][0]
    idx = op2.Global(len(nodes), nodes, dtype=np.int32, name='node_idx')
    nodes_out = fs_out.bt_masks['geometric'][0]
    idx_out = op2.Global(len(nodes_out), nodes_out, dtype=np.int32, name='node_idx')
    kernel = op2.Kernel("""
        void my_kernel(double **func, double **out3d, double **z, int *idx, int *idx_out) {
            for ( int d = 0; d < %(nodes)d; d++ ) {
                double h = z[idx[d]+1][0] - z[idx[d]+0][0];
                for ( int c = 0; c < %(func_dim)d; c++ ) {
                    // we integrate f(x) = (b-a)*x/h + a
                    // from x=0 to x=h
                    double a = func[idx[d]+0][c];
                    double b = func[idx[d]+1][c];
                    // bottom of elem x=0
                    out3d[idx_out[d]+0][c] += 0.0;
                    // top of elem x=h
                    out3d[idx_out[d]+1][c] += 0.5*h*(b-a) + h*a + out3d[idx_out[d]+0][c];
                }
            }
        }""" % {'nodes': len(nodes),
                'func_dim': output_3d.function_space().dim,
                'v_nodes': n_vert_nodes},
        'my_kernel')

    op2.par_loop(
        kernel, fs_3d.mesh().cell_set,
        source.dat(op2.READ, fs_3d.cell_node_map()),
        output_3d.dat(op2.INC, fs_out.cell_node_map()),
        z_coords.dat(op2.READ, z_coords.function_space().cell_node_map()),
        idx(op2.READ),
        idx_out(op2.READ),
        iterate=op2.ALL)


def strong_integral_3d_p2(source, output_3d, z_coords):
    """Computes strong integral of a P1DGxP1DG onto P1DGxP2 field"""
    fs_3d = source.function_space()
    fs_out = output_3d.function_space()
    # number of nodes in vertical direction
    n_vert_nodes = len(fs_out.fiat_element.B.entity_closure_dofs()[1][0])

    nodes = fs_3d.bt_masks['geometric'][0]
    idx = op2.Global(len(nodes), nodes, dtype=np.int32, name='node_idx')
    nodes_out = fs_out.bt_masks['geometric'][0]
    idx_out = op2.Global(len(nodes_out), nodes_out, dtype=np.int32, name='node_idx')
    kernel = op2.Kernel("""
        void my_kernel(double **func, double **out3d, double **z, int *idx, int *idx_out) {
            for ( int d = 0; d < %(nodes)d; d++ ) {
                double h = z[idx[d]+1][0] - z[idx[d]+0][0];
                for ( int c = 0; c < %(func_dim)d; c++ ) {
                    // we integrate f(x) = (b-a)*x/h + a
                    // from x=0 to x=h
                    double a = func[idx[d]+0][c];
                    double b = func[idx[d]+1][c];
                    // bottom of elem x=0
                    out3d[idx_out[d]+0][c] += 0.0;
                    // top of elem x=h
                    out3d[idx_out[d]+1][c] += 0.5*h*(b-a) + h*a + out3d[idx_out[d]+0][c];
                    // middle of elem x= 0.5h
                    out3d[idx_out[d]+2][c] += 0.125*h*(b-a) + 0.5*h*a + out3d[idx_out[d]+0][c];
                }
            }
        }""" % {'nodes': len(nodes),
                'func_dim': output_3d.function_space().dim,
                'v_nodes': n_vert_nodes},
        'my_kernel')

    op2.par_loop(
        kernel, fs_3d.mesh().cell_set,
        source.dat(op2.READ, fs_3d.cell_node_map()),
        output_3d.dat(op2.INC, fs_out.cell_node_map()),
        z_coords.dat(op2.READ, z_coords.function_space().cell_node_map()),
        idx(op2.READ),
        idx_out(op2.READ),
        iterate=op2.ALL)

# ---------------
# test
# ---------------

mesh2d = UnitSquareMesh(2, 2)
nlayers = 8
mesh = ExtrudedMesh(mesh2d, nlayers, 1.0/nlayers)
# deform mesh
mesh.coordinates.dat.data[:, 2] *= 0.1 + 0.9*mesh.coordinates.dat.data[:, 1]

P1 = FunctionSpace(mesh, 'CG', 1, vfamily='CG', vdegree=1)
P1xP2 = FunctionSpace(mesh, 'CG', 1, vfamily='CG', vdegree=2)
V = FunctionSpace(mesh, 'DG', 1, vfamily='DG', vdegree=1)
Vv = VectorFunctionSpace(mesh, 'DG', 1, vfamily='DG', vdegree=1)
Vv_int = VectorFunctionSpace(mesh, 'DG', 1, vfamily='CG', vdegree=1)
Vv_int_p2 = VectorFunctionSpace(mesh, 'DG', 1, vfamily='CG', vdegree=2)
V_2d = FunctionSpace(mesh2d, 'DG', 1)
Vv_2d = VectorFunctionSpace(mesh2d, 'DG', 1)

f = Function(Vv, name='input')
f_int = Function(Vv_int, name='integral 3d')
f_int_p2 = Function(Vv_int_p2, name='integral 3d')
f_int_2d = Function(Vv_2d, name='integral')
f_int_w = Function(Vv_int_p2, name='weak integral 3d')
f_int_w_2d = Function(Vv_2d, name='weak integral')

xyz = SpatialCoordinate(mesh)
f.interpolate(as_vector((xyz[0], xyz[2], 1)))

# z coords
z_func = Function(P1, name='z coords').interpolate(xyz[2])

strong_integral_2d(f, f_int_2d, z_func)
strong_integral_3d_p1(f, f_int, z_func)
strong_integral_3d_p2(f, f_int_p2, z_func)

print 'function', f.dat.data.min(), f.dat.data.max()
print 'integral', f_int_2d.dat.data.min(), f_int_2d.dat.data.max()
print 'int 3d  ', f_int.dat.data.min(), f_int.dat.data.max()

out_f = File('function.pvd')
out_f.write(f)
out_f.write(f)  # output twice - easier to refresh on paraview
out_f_int_2d = File('integral_2d.pvd')
out_f_int_2d.write(f_int_2d)
out_f_int_2d.write(f_int_2d)
out_f_int = File('integral_3d.pvd')
out_f_int.write(f_int)
out_f_int.write(f_int_p2)  # second one is the p2 solution

# -----------------------------------------------------
# compute the same in weak sense with Thetis
# -----------------------------------------------------

integrator = VerticalIntegrator(f, f_int_w,
                                bottom_to_top=True,
                                bnd_value=Constant((0.0, 0.0, 0.0)),
                                average=False)
extract_surf = SubFunctionExtractor(f_int_w, f_int_w_2d, boundary='top', elem_facet='top')
integrator.solve()
extract_surf.solve()

print 'function', f.dat.data.min(), f.dat.data.max()
print 'integral', f_int_w_2d.dat.data.min(), f_int_w_2d.dat.data.max()  # should be "0.0 2.0"

out_f_int_w_2d = File('weak_integral_2d.pvd')
out_f_int_w_2d.write(f_int_w_2d)
out_f_int_w_2d.write(f_int_w_2d)
out_f_int_w = File('weak_integral_3d.pvd')
out_f_int_w.write(f_int_w)
out_f_int_w.write(f_int_w)