wave.py

if do_plotting:
    ### PLOTTING STUFF ###
    mesh_cg_fs = VectorFunctionSpace(mesh, "CG", 1)
    coords_orig = Function(mesh_cg_fs)
    coords_stretch = Function(mesh_cg_fs)
    coords_latlon = Function(mesh_dg_fs)
    coords_orig.dat.data[:] = mesh.coordinates.dat.data[:]
    coords_stretch.dat.data[:] = mesh_fake.coordinates.dat.data[:]
    # lat-lon 'x' = atan2(y, x)
    coords_latlon.dat.data[:,0] = np.arctan2(coords_dg.dat.data[:,1], coords_dg.dat.data[:,0])
    # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2))
    coords_latlon.dat.data[:,1] = np.arcsin(coords_dg.dat.data[:,2]/np.sqrt(coords_dg.dat.data[:,0]**2 + coords_dg.dat.data[:,1]**2 + coords_dg.dat.data[:,2]**2))
    # lat-lon 'z' = sqrt(x^2 + y^2 + z^2) - a, and scaled by 1/a.
    coords_latlon.dat.data[:,2] = (np.sqrt(coords_dg.dat.data[:,0]**2 + coords_dg.dat.data[:,1]**2 + coords_dg.dat.data[:,2]**2) - a)/a

    kernel = op2.Kernel("""
    #define PI 3.141592653589793
    #define TWO_PI 6.283185307179586
    void splat_coords(double **coords) {
        for ( int i = 0; i < 2; i++ ) {
            double diff0 = (coords[0 + i][0] - coords[2 + i][0]);
            double diff1 = (coords[0 + i][0] - coords[4 + i][0]);
            double diff2 = (coords[2 + i][0] - coords[4 + i][0]);

            if (fabs(diff0) > PI || fabs(diff1) > PI || fabs(diff2) > PI) {
                const int sign0 = coords[0 + i][0] < 0 ? -1 : 1;
                const int sign1 = coords[2 + i][0] < 0 ? -1 : 1;
                const int sign2 = coords[4 + i][0] < 0 ? -1 : 1;
                if (sign0 < 0) {
                    coords[0 + i][0] += TWO_PI;
                }
                if (sign1 < 0) {
                    coords[2 + i][0] += TWO_PI;
                }
                if (sign2 < 0) {
                    coords[4 + i][0] += TWO_PI;
                }
            }
        }
    }""", "splat_coords")

    op2.par_loop(kernel, coords_latlon.cell_set,
                 coords_latlon.dat(op2.RW, coords_latlon.cell_node_map()))
    out_b_s = File("stretch/b.pvd")
    out_p_s = File("stretch/p.pvd")
    out_b_l = File("latlon/b.pvd")
    out_p_l = File("latlon/p.pvd")

    viz_fs = FunctionSpace(mesh, "DG", 1)
    b0_viz = project(b0, viz_fs)
    p0_viz = project(p0, viz_fs)

    mesh.coordinates = coords_stretch
    out_b_s << b0_viz
    out_p_s << p0_viz

    mesh.coordinates = coords_latlon
    out_b_l << b0_viz
    out_p_l << p0_viz

    mesh.coordinates = coords_orig
    ### END PLOTTING STUFF ###