def PeriodicIntervalMesh(ncells, length):
"""Generate a periodic mesh of an interval.
:arg ncells: The number of cells over the interval.
:arg length: The length the interval."""
m = CircleManifoldMesh(ncells)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=1)
old_coordinates = Function(m.coordinates)
new_coordinates = Function(coord_fs)
periodic_kernel = """double Y,pi;
Y = 0.5*(old_coords[0][1]-old_coords[1][1]);
pi=3.141592653589793;
for(int i=0;i<2;i++){
new_coords[i][0] = atan2(old_coords[i][1],old_coords[i][0])/pi/2;
if(new_coords[i][0]<0.) new_coords[i][0] += 1;
if(new_coords[i][0]==0 && Y<0.) new_coords[i][0] = 1.0;
new_coords[i][0] *= L;
}"""
periodic_kernel = periodic_kernel.replace('L', str(length))
par_loop(periodic_kernel, dx,
{"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ)})
m.coordinates = new_coordinates
return m
def PeriodicIntervalMesh(ncells, length):
"""Generate a periodic mesh of an interval.
:arg ncells: The number of cells over the interval.
:arg length: The length the interval."""
if ncells < 3:
raise ValueError("1D periodic meshes with fewer than 3 \
cells are not currently supported")
m = CircleManifoldMesh(ncells)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=1)
old_coordinates = m.coordinates
new_coordinates = Function(coord_fs)
periodic_kernel = """double Y,pi;
Y = 0.5*(old_coords[0][1]-old_coords[1][1]);
pi=3.141592653589793;
for(int i=0;i<2;i++){
new_coords[i][0] = atan2(old_coords[i][1],old_coords[i][0])/pi/2;
if(new_coords[i][0]<0.) new_coords[i][0] += 1;
if(new_coords[i][0]==0 && Y<0.) new_coords[i][0] = 1.0;
new_coords[i][0] *= L[0];
}"""
cL = Constant(length)
par_loop(periodic_kernel, dx,
{"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ),
"L": (cL, READ)})
return mesh.Mesh(new_coordinates)
def get_latlon_mesh(mesh):
coords_orig = mesh.coordinates
mesh_dg_fs = VectorFunctionSpace(mesh, "DG", 1)
coords_dg = Function(mesh_dg_fs)
coords_latlon = Function(mesh_dg_fs)
par_loop("""
for (int i=0; i<3; i++) {
for (int j=0; j<3; j++) {
dg[i][j] = cg[i][j];
}
}
""", dx, {'dg': (coords_dg, WRITE),
'cg': (coords_orig, READ)})
# 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))
coords_latlon.dat.data[:,2] = 0.0
kernel = op2.Kernel("""
#define PI 3.141592653589793
#define TWO_PI 6.283185307179586
void splat_coords(double **coords) {
double diff0 = (coords[0][0] - coords[1][0]);
double diff1 = (coords[0][0] - coords[2][0]);
double diff2 = (coords[1][0] - coords[2][0]);
if (fabs(diff0) > PI || fabs(diff1) > PI || fabs(diff2) > PI) {
const int sign0 = coords[0][0] < 0 ? -1 : 1;
const int sign1 = coords[1][0] < 0 ? -1 : 1;
const int sign2 = coords[2][0] < 0 ? -1 : 1;
if (sign0 < 0) {
coords[0][0] += TWO_PI;
}
if (sign1 < 0) {
coords[1][0] += TWO_PI;
}
if (sign2 < 0) {
coords[2][0] += TWO_PI;
}
}
}""", "splat_coords")
op2.par_loop(kernel, coords_latlon.cell_set,
coords_latlon.dat(op2.RW, coords_latlon.cell_node_map()))
return Mesh(coords_latlon)
def V_dof_weights(self, V):
"""Dof weights for Fortin projection.
:arg V: function space to compute weights for.
:returns: A PETSc Vec.
"""
key = V.dim()
try:
return self._V_dof_weights[key]
except KeyError:
# Compute dof multiplicity for V
# Spin over all (owned) cells incrementing visible dofs by 1.
# After halo exchange, the Vec representation is the
# global Vector counting the number of cells that see each
# dof.
f = firedrake.Function(V)
firedrake.par_loop(("{[i, j]: 0 <= i < A.dofs and 0 <= j < %d}" % V.value_size,
"A[i, j] = A[i, j] + 1"),
firedrake.dx,
{"A": (f, firedrake.INC)},
is_loopy_kernel=True)
with f.dat.vec_ro as fv:
return self._V_dof_weights.setdefault(key, fv.copy())
def PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction="x", quadrilateral=False, reorder=None, comm=COMM_WORLD):
"""Generates RectangleMesh that is periodic in the x or y direction.
:arg nx: The number of cells in the x direction
:arg ny: The number of cells in the y direction
:arg Lx: The extent in the x direction
:arg Ly: The extent in the y direction
:kwarg direction: The direction of the periodicity (default x).
:kwarg quadrilateral: (optional), creates quadrilateral mesh, defaults to False
:kwarg reorder: (optional), should the mesh be reordered
:kwarg comm: Optional communicator to build the mesh on (defaults to
COMM_WORLD).
If direction == "x" the boundary edges in this mesh are numbered as follows:
* 1: plane y == 0
* 2: plane y == Ly
If direction == "y" the boundary edges are:
* 1: plane x == 0
* 2: plane x == Lx
"""
if direction not in ("x", "y"):
raise ValueError("Unsupported periodic direction '%s'" % direction)
# handle x/y directions: na, La are for the periodic axis
na, nb, La, Lb = nx, ny, Lx, Ly
if direction == "y":
na, nb, La, Lb = ny, nx, Ly, Lx
if na < 3:
raise ValueError("2D periodic meshes with fewer than 3 \
cells in each direction are not currently supported")
m = CylinderMesh(na, nb, 1.0, 1.0, longitudinal_direction="z",
quadrilateral=quadrilateral, reorder=reorder, comm=comm)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=2)
old_coordinates = m.coordinates
new_coordinates = Function(coord_fs)
# make x-periodic mesh
# unravel x coordinates like in periodic interval
# set y coordinates to z coordinates
periodic_kernel = """double Y,pi;
Y = 0.0;
for(int i=0; i<old_coords.dofs; i++) {
Y += old_coords[i][1];
}
pi=3.141592653589793;
for(int i=0;i<new_coords.dofs;i++){
new_coords[i][0] = atan2(old_coords[i][1],old_coords[i][0])/pi/2;
if(new_coords[i][0]<0.) new_coords[i][0] += 1;
if(new_coords[i][0]==0 && Y<0.) new_coords[i][0] = 1.0;
new_coords[i][0] *= Lx[0];
new_coords[i][1] = old_coords[i][2]*Ly[0];
}"""
cLx = Constant(La)
cLy = Constant(Lb)
par_loop(periodic_kernel, dx,
{"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ),
"Lx": (cLx, READ),
"Ly": (cLy, READ)})
if direction == "y":
# flip x and y coordinates
operator = np.asarray([[0, 1],
[1, 0]])
new_coordinates.dat.data[:] = np.dot(new_coordinates.dat.data, operator.T)
return mesh.Mesh(new_coordinates)
def PeriodicRectangleMesh(nx, ny, Lx, Ly, direction="both",
quadrilateral=False, reorder=None, comm=COMM_WORLD):
"""Generate a periodic rectangular mesh
:arg nx: The number of cells in the x direction
:arg ny: The number of cells in the y direction
:arg Lx: The extent in the x direction
:arg Ly: The extent in the y direction
:arg direction: The direction of the periodicity, one of
``"both"``, ``"x"`` or ``"y"``.
:kwarg quadrilateral: (optional), creates quadrilateral mesh, defaults to False
:kwarg reorder: (optional), should the mesh be reordered
:kwarg comm: Optional communicator to build the mesh on (defaults to
COMM_WORLD).
If direction == "x" the boundary edges in this mesh are numbered as follows:
* 1: plane y == 0
* 2: plane y == Ly
If direction == "y" the boundary edges are:
* 1: plane x == 0
* 2: plane x == Lx
"""
if direction not in ("both", "x", "y"):
raise ValueError("Cannot have a periodic mesh with periodicity '%s'" % direction)
if direction != "both":
return PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction=direction,
quadrilateral=quadrilateral,
reorder=reorder,
comm=comm)
if nx < 3 or ny < 3:
raise ValueError("2D periodic meshes with fewer than 3 \
cells in each direction are not currently supported")
m = TorusMesh(nx, ny, 1.0, 0.5, quadrilateral=quadrilateral, reorder=reorder,
comm=comm)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=2)
old_coordinates = m.coordinates
new_coordinates = Function(coord_fs)
periodic_kernel = """
double pi = 3.141592653589793;
double eps = 1e-12;
double bigeps = 1e-1;
double phi, theta, Y, Z;
Y = 0.0;
Z = 0.0;
for(int i=0; i<old_coords.dofs; i++) {
Y += old_coords[i][1];
Z += old_coords[i][2];
}
for(int i=0; i<new_coords.dofs; i++) {
phi = atan2(old_coords[i][1], old_coords[i][0]);
if (fabs(sin(phi)) > bigeps)
theta = atan2(old_coords[i][2], old_coords[i][1]/sin(phi) - 1.0);
else
theta = atan2(old_coords[i][2], old_coords[i][0]/cos(phi) - 1.0);
new_coords[i][0] = phi/(2.0*pi);
if(new_coords[i][0] < -eps) {
new_coords[i][0] += 1.0;
}
if(fabs(new_coords[i][0]) < eps && Y < 0.0) {
new_coords[i][0] = 1.0;
}
new_coords[i][1] = theta/(2.0*pi);
if(new_coords[i][1] < -eps) {
new_coords[i][1] += 1.0;
}
if(fabs(new_coords[i][1]) < eps && Z < 0.0) {
new_coords[i][1] = 1.0;
}
new_coords[i][0] *= Lx[0];
new_coords[i][1] *= Ly[0];
}
"""
cLx = Constant(Lx)
cLy = Constant(Ly)
par_loop(periodic_kernel, dx,
{"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ),
"Lx": (cLx, READ),
"Ly": (cLy, READ)})
return mesh.Mesh(new_coordinates)
par_loop("""
double v0[3];
double v1[3];
double n[3];
double com[3];
double dot;
double norm;
norm = 0.0;
dot = 0.0;
// form "x1 - x0" and "x2 - x0" of cell base
for (int i=0; i<3; ++i) {
v0[i] = coords[2][i] - coords[0][i];
v1[i] = coords[4][i] - coords[0][i];
}
for (int i=0; i<3; ++i) {
com[i] = 0.0;
}
// take cross-product to form normal vector
n[0] = v0[1] * v1[2] - v0[2] * v1[1];
n[1] = v0[2] * v1[0] - v0[0] * v1[2];
n[2] = v0[0] * v1[1] - v0[1] * v1[0];
// get (scaled) centre-of-mass of cell
for (int i=0; i<6; ++i) {
com[0] += coords[i][0];
com[1] += coords[i][1];
com[2] += coords[i][2];
}
// is the normal pointing outwards or inwards w.r.t. origin?
for (int i=0; i<3; ++i) {
dot += com[i]*n[i];
}
for (int i=0; i<3; ++i) {
norm += n[i]*n[i];
}
// normalise normal vector and multiply by -1 if dot product was < 0
norm = sqrt(norm);
norm *= (dot < 0.0 ? -1.0 : 1.0);
for (int i=0; i<3; ++i) {
normals[0][i] = n[i] / norm;
}
""", dx,
{'normals': (zhat, WRITE),
'coords': (mesh.coordinates, READ)})
def PeriodicRectangleMesh(nx, ny, Lx, Ly, quadrilateral=False, reorder=None):
"""Generate a periodic rectangular mesh
:arg nx: The number of cells in the x direction
:arg ny: The number of cells in the y direction
:arg Lx: The extent in the x direction
:arg Ly: The extent in the y direction
:kwarg quadrilateral: (optional), creates quadrilateral mesh, defaults to False
:kwarg reorder: (optional), should the mesh be reordered
"""
if nx < 3 or ny < 3:
raise ValueError("2D periodic meshes with fewer than 3 \
cells in each direction are not currently supported")
m = TorusMesh(nx, ny, 1.0, 0.5, quadrilateral=quadrilateral, reorder=reorder)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=2)
old_coordinates = m.coordinates
new_coordinates = Function(coord_fs)
periodic_kernel = """
double pi = 3.141592653589793;
double eps = 1e-12;
double bigeps = 1e-1;
double phi, theta, Y, Z;
Y = 0.0;
Z = 0.0;
for(int i=0; i<old_coords.dofs; i++) {
Y += old_coords[i][1];
Z += old_coords[i][2];
}
for(int i=0; i<new_coords.dofs; i++) {
phi = atan2(old_coords[i][1], old_coords[i][0]);
if (fabs(sin(phi)) > bigeps)
theta = atan2(old_coords[i][2], old_coords[i][1]/sin(phi) - 1.0);
else
theta = atan2(old_coords[i][2], old_coords[i][0]/cos(phi) - 1.0);
new_coords[i][0] = phi/(2.0*pi);
if(new_coords[i][0] < -eps) {
new_coords[i][0] += 1.0;
}
if(fabs(new_coords[i][0]) < eps && Y < 0.0) {
new_coords[i][0] = 1.0;
}
new_coords[i][1] = theta/(2.0*pi);
if(new_coords[i][1] < -eps) {
new_coords[i][1] += 1.0;
}
if(fabs(new_coords[i][1]) < eps && Z < 0.0) {
new_coords[i][1] = 1.0;
}
new_coords[i][0] *= Lx[0];
new_coords[i][1] *= Ly[0];
}
"""
cLx = Constant(Lx)
cLy = Constant(Ly)
par_loop(periodic_kernel, dx,
{"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ),
"Lx": (cLx, READ),
"Ly": (cLy, READ)})
return mesh.Mesh(new_coordinates)
