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 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."""
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 create_function_marker(PHI, W, xlimits, ylimits):
x, y, z = fd.SpatialCoordinate(PHI.ufl_domain())
x_func, y_func, z_func = (
fd.Function(PHI),
fd.Function(PHI),
fd.Function(PHI),
)
with fda.stop_annotating():
x_func.interpolate(x)
y_func.interpolate(y)
z_func.interpolate(z)
domain = "{[i, j]: 0 <= i < f.dofs and 0<= j <= 3}"
instruction = f"""
f[i, j] = 1.0 if (x[i, 0] < {xlimits[1]} and x[i, 0] > {xlimits[0]}) and (y[i, 0] < {ylimits[0]} or y[i, 0] > {ylimits[1]}) and z[i, 0] < 1e-7 else 0.0
"""
I_BC = fd.Function(W)
fd.par_loop(
(domain, instruction),
dx,
{
"f": (I_BC, fd.RW),
"x": (x_func, fd.READ),
"y": (y_func, fd.READ),
"z": (z_func, fd.READ),
},
is_loopy_kernel=True,
)
return I_BC
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(
"""for (int i = 0; i < A.dofs; i++)
for (int j = 0; j < {}; j++)
A[i][j] += 1;
""".format(V.value_size),
firedrake.dx,
{"A": (f, firedrake.INC)})
with f.dat.vec_ro as fv:
return self._V_dof_weights.setdefault(key, fv.copy())
def PeriodicIntervalMesh(ncells, length, comm=COMM_WORLD):
"""Generate a periodic mesh of an interval.
:arg ncells: The number of cells over the interval.
:arg length: The length the interval.
:kwarg comm: Optional communicator to build the mesh on (defaults to
COMM_WORLD).
"""
if ncells < 3:
raise ValueError("1D periodic meshes with fewer than 3 \
cells are not currently supported")
m = CircleManifoldMesh(ncells, comm=comm)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=1)
old_coordinates = m.coordinates
new_coordinates = Function(coord_fs)
periodic_kernel = """
const double pi = 3.141592653589793;
const double eps = 1e-12;
double a = atan2(old_coords[0][1], old_coords[0][0]) / (2*pi);
double b = atan2(old_coords[1][1], old_coords[1][0]) / (2*pi);
int swap = 0;
if ( a >= b ) {
const double tmp = b;
b = a;
a = tmp;
swap = 1;
}
if ( fabs(b) < eps && a < -eps ) {
b = 1.0;
}
if ( a < -eps ) {
a += 1;
}
if ( b < -eps ) {
b += 1;
}
if ( swap ) {
const double tmp = b;
b = a;
a = tmp;
}
new_coords[0][0] = a * L[0];
new_coords[1][0] = b * 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):
"""Build 2D projected mesh of spherical mesh"""
crds_orig = mesh.coordinates
mesh_dg_fs = VectorFunctionSpace(mesh, "DG", 1)
crds_dg = Function(mesh_dg_fs)
crds_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': (crds_dg, WRITE),
'cg': (crds_orig, READ)
})
# lat-lon 'x' = atan2(y, x)
crds_latlon.dat.data[:, 0] = arctan2(crds_dg.dat.data[:, 1],
crds_dg.dat.data[:, 0])
# lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2))
crds_latlon.dat.data[:, 1] = (arcsin(
crds_dg.dat.data[:, 2] /
np_sqrt(crds_dg.dat.data[:, 0]**2 + crds_dg.dat.data[:, 1]**2 +
crds_dg.dat.data[:, 2]**2)))
crds_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, crds_latlon.cell_set,
crds_latlon.dat(op2.RW, crds_latlon.cell_node_map()))
return Mesh(crds_latlon)
def multiplicity(V):
# Lawrence's magic code for calculating dof multiplicities
shapes = (V.finat_element.space_dimension(), np.prod(V.shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
weight = firedrake.Function(V)
firedrake.par_loop((domain, instructions),
firedrake.dx, {"w": (weight, op2.INC)},
is_loopy_kernel=True)
return weight
def extract_comp(mode, comp, result):
instruction = f"""
component[0] = abs(u[i, {comp}])
"""
fd.par_loop(
(domain, instruction),
fd.dx,
{
"u": (coords, fd.READ),
"component": (result, mode)
},
is_loopy_kernel=True,
)
return result
def __init__(self, Vf, Vc, Vf_bcs, Vc_bcs):
self.Vf = Vf
self.Vc = Vc
self.Vf_bcs = Vf_bcs
self.Vc_bcs = Vc_bcs
self.uc = firedrake.Function(Vc)
self.uf = firedrake.Function(Vf)
self.prolong_kernel = self.prolongation_transfer_kernel_action(
Vf, self.uc)
matrix_kernel = self.prolongation_transfer_kernel_action(
Vf, firedrake.TestFunction(Vc))
element_kernel = loopy.generate_code_v2(
matrix_kernel.code).device_code()
element_kernel = element_kernel.replace(
"void expression_kernel", "static void expression_kernel")
dimc = Vc.finat_element.space_dimension() * Vc.value_size
dimf = Vf.finat_element.space_dimension() * Vf.value_size
self.restrict_code = f"""
{element_kernel}
void restriction(double *restrict Rc, const double *restrict Rf, const double *restrict w)
{{
double Afc[{dimf}*{dimc}] = {{0}};
expression_kernel(Afc);
for (int32_t i = 0; i < {dimf}; i++)
for (int32_t j = 0; j < {dimc}; j++)
Rc[j] += Afc[i*{dimc} + j] * Rf[i] / w[i];
}}
"""
self.restrict_kernel = op2.Kernel(self.restrict_code, "restriction")
self.mesh = Vf.mesh()
# Lawrence's magic code for calculating dof multiplicities
shapes = (Vf.finat_element.space_dimension(), np.prod(Vf.shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
self.weight = firedrake.Function(Vf)
firedrake.par_loop((domain, instructions),
firedrake.dx, {"w": (self.weight, op2.INC)},
is_loopy_kernel=True)
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 calculate_max_vel(velocity: fd.Function):
mesh = velocity.ufl_domain()
MAXSP = fd.FunctionSpace(mesh, "R", 0)
maxv = fd.Function(MAXSP)
domain = "{[i, j] : 0 <= i < u.dofs}"
instruction = """
maxv[0] = abs(u[i, 0]) + abs(u[i, 1])
"""
fd.par_loop(
(domain, instruction),
fd.dx,
{
"u": (velocity, fd.READ),
"maxv": (maxv, fd.MAX)
},
is_loopy_kernel=True,
)
maxval = maxv.dat.data[0]
return maxval
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 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())
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)
def clement_interpolant(source, target_space=None, boundary_tag=None):
r"""
Compute the Clement interpolant of a :math:`\mathbb P0`
source field, i.e. take the volume average over
neighbouring cells at each vertex.
:arg source: the :math:`\mathbb P0` source field
:kwarg target_space: the :math:`\mathbb P1` space to
interpolate into
:boundary_tag: optional boundary tag to compute the
Clement interpolant over.
"""
V = source.function_space()
assert V.ufl_element().family() == "Discontinuous Lagrange"
assert V.ufl_element().degree() == 0
rank = len(V.ufl_element().value_shape())
mesh = V.mesh()
dim = mesh.topological_dimension()
P1 = firedrake.FunctionSpace(mesh, "CG", 1)
dX = ufl.dx if boundary_tag is None else ufl.ds(boundary_tag)
if target_space is None:
if rank == 0:
target_space = P1
elif rank == 1:
target_space = firedrake.VectorFunctionSpace(mesh, "CG", 1)
elif rank == 2:
target_space = firedrake.TensorFunctionSpace(mesh, "CG", 1)
else:
raise ValueError(f"Rank-{rank} tensors are not supported.")
else:
assert target_space.ufl_element().family() == "Lagrange"
assert target_space.ufl_element().degree() == 1
target = firedrake.Function(target_space)
# Compute the patch volume at each vertex
if boundary_tag is None:
P0 = firedrake.FunctionSpace(mesh, "DG", 0)
dx = ufl.dx(domain=mesh)
volume = firedrake.assemble(firedrake.TestFunction(P0) * dx)
else:
volume = get_facet_areas(mesh)
patch_volume = firedrake.Function(P1)
kernel = "for (int i=0; i < p.dofs; i++) p[i] += v[0];"
keys = {
"v": (volume, op2.READ),
"p": (patch_volume, op2.INC),
}
firedrake.par_loop(kernel, dX, keys)
# Volume average
keys = {
"s": (source, op2.READ),
"v": (volume, op2.READ),
"t": (target, op2.INC),
}
if rank == 0:
firedrake.par_loop(
"""
for (int i=0; i < t.dofs; i++) {
t[i] += s[0]*v[0];
}
""",
dX,
keys,
)
elif rank == 1:
firedrake.par_loop(
"""
int d = %d;
for (int i=0; i < t.dofs; i++) {
for (int j=0; j < d; j++) {
t[i*d + j] += s[j]*v[0];
}
}
"""
% dim,
dX,
keys,
)
elif rank == 2:
firedrake.par_loop(
"""
int d = %d;
int Nd = d*d;
for (int i=0; i < t.dofs; i++) {
for (int j=0; j < d; j++) {
for (int k=0; k < d; k++) {
t[i*Nd + j*d + k] += s[j*d + k]*v[0];
}
}
}
"""
% dim,
dX,
keys,
)
else:
raise ValueError(f"Rank-{rank} tensors are not supported.")
target.interpolate(target / patch_volume)
if boundary_tag is not None:
target.dat.data_with_halos[:] = np.nan_to_num(target.dat.data_with_halos)
return target
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, 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)
def __init__(
self,
S,
mesh,
beta=1,
gamma=1.0e4,
bcs=None,
dx=dx,
design_domain=None,
solver_parameters=direct_parameters,
output_dir="./",
):
"""
Solver class to regularize the shape derivatives as explained in
Frédéric de Gournay
Velocity Extension for the Level-set Method and Multiple Eigenvalues in Shape Optimization
SIAM J. Control Optim., 45(1), 343–367. (25 pages)
Args:
S ([type]): Function space of the mesh coordinates
mesh ([type]): Mesh
beta ([type], optional): Regularization parameter.
It should be finite multiple of the mesh size.
Defaults to 1.
gamma ([type], optional): Penalty parameter for the penalization of the normal components
of the regularized shape derivatives on the boundary. Defaults to 1.0e4.
bcs ([type], optional): Dirichlet Boundary conditions.
They should be setting the regularized shape derivatives to zero
wherever there are boundary conditions on the original PDE.
Defaults to None.
dx ([type], optional): [description]. Defaults to dx.
design_domain ([type], optional): If we're interested in setting the shape derivatives to
zero outside of the design_domain,
we pass design_domain marker.
This is convenient when we want to fix certain regions in the domain.
Defaults to None.
solver_parameters ([type], optional): Solver options. Defaults to direct_parameters.
output_dir (str, optional): Plot the output somewhere. Defaults to "./".
"""
n = FacetNormal(mesh)
theta, xi = [TrialFunction(S), TestFunction(S)]
self.xi = xi
hmin = min_mesh_size(mesh)
if beta > 20.0 * hmin:
warning(
f"Length scale parameter beta ({beta}) is much larger than the mesh size {hmin}"
)
self.beta_param = Constant(beta)
self.a = (self.beta_param * inner(grad(theta), grad(xi)) +
inner(theta, xi)) * (dx)
if isinstance(mesh.topology, ExtrudedMeshTopology):
ds_reg = ds_b + ds_v + ds_tb + ds_t
else:
ds_reg = ds
self.a += Constant(gamma) * (inner(dot(theta, n), dot(xi, n))) * ds_reg
# Dirichlet boundary conditions equal to zero for regions where we want
# the domain to be static, i.e. zero velocities
if bcs is None:
self.bcs = []
else:
self.bcs = Enlist(bcs)
if design_domain is not None:
# Heaviside step function in domain of interest
V_DG0_B = FunctionSpace(mesh, "DG", 0)
I_B = Function(V_DG0_B)
I_B.assign(1.0)
# Set to zero all the cells within sim_domain
par_loop(
("{[i] : 0 <= i < f.dofs}", "f[i, 0] = 0.0"),
dx(design_domain),
{"f": (I_B, WRITE)},
is_loopy_kernel=True,
)
I_cg_B = Function(S)
dim = S.mesh().geometric_dimension()
# Assume that `A` is a :class:`.Function` in CG1 and `B` is a
# `.Function` in DG0. Then the following code sets each DoF in
# `A` to the maximum value that `B` attains in the cells adjacent to
# that DoF::
par_loop(
(
"{{[i, j] : 0 <= i < A.dofs and 0 <= j < {0} }}".format(
dim),
"A[i, j] = fmax(A[i, j], B[0, 0])",
),
dx,
{
"A": (I_cg_B, RW),
"B": (I_B, READ)
},
is_loopy_kernel=True,
)
import numpy as np
class MyBC(DirichletBC):
def __init__(self, V, value, markers):
# Call superclass init
# We provide a dummy subdomain id.
super(MyBC, self).__init__(V, value, 0)
# Override the "nodes" property which says where the boundary
# condition is to be applied.
self.nodes = np.unique(
np.where(markers.dat.data_ro_with_halos > 0)[0])
self.bcs.append(MyBC(S, 0, I_cg_B))
self.Av = assemble(self.a, bcs=self.bcs)
self.solver_parameters = solver_parameters
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)
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 solve(self, phi: fd.Function, iters: int = 5) -> fd.Function:
marking = fd.Function(self.DG0)
marking_bc_nodes = fd.Function(self.V)
# Mark cells cut by phi(x) = 0
domain = "{[i, j]: 0 <= i < b.dofs}"
instructions = """
<float64> min_value = 1e20
<float64> max_value = -1e20
for i
min_value = fmin(min_value, b[i, 0])
max_value = fmax(max_value, b[i, 0])
end
a[0, 0] = 1.0 if (min_value < 0 and max_value > 0) else 0.0
"""
fd.par_loop(
(domain, instructions),
dx,
{
"a": (marking, RW),
"b": (phi, READ)
},
is_loopy_kernel=True,
)
# Mark the nodes in the marked cells
fd.par_loop(
("{[i] : 0 <= i < A.dofs}", "A[i, 0] = fmax(A[i, 0], B[0, 0])"),
dx,
{
"A": (marking_bc_nodes, RW),
"B": (marking, READ)
},
is_loopy_kernel=True,
)
# Mark the nodes in the marked cells
# Project the gradient of phi on the cut cells
self.phi.assign(phi)
V = self.V
rho, sigma = fd.TrialFunction(V), fd.TestFunction(V)
a = rho * sigma * marking * dx
L_proj = (self.phi / sqrt(inner(grad(self.phi), grad(self.phi))) *
marking * sigma * dx)
bc_proj = BCOut(V, fd.Constant(0.0), marking_bc_nodes)
self.A_proj = fd.assemble(a, tensor=self.A_proj, bcs=bc_proj)
b_proj = fd.assemble(L_proj, bcs=bc_proj)
solver_proj = fd.LinearSolver(self.A_proj,
solver_parameters=self.solver_parameters)
solver_proj.solve(self.phi_int, b_proj)
def nabla_phi_bar(phi):
return sqrt(inner(grad(phi), grad(phi)))
def d1(s):
return fd.Constant(1.0) - fd.Constant(1.0) / s
def d2(s):
return conditional(le(s, fd.Constant(1.0)), s - fd.Constant(1.0),
d1(s))
def residual_phi(phi):
return fd.norm(
fd.assemble(
d2(nabla_phi_bar(phi)) * inner(grad(phi), grad(sigma)) *
dx))
a = inner(grad(rho), grad(sigma)) * dx
L = (inner(
(-d2(nabla_phi_bar(self.phi)) + fd.Constant(1.0)) * grad(self.phi),
grad(sigma),
) * dx)
bc = BCInt(V, self.phi_int, marking_bc_nodes)
phi_sol = fd.Function(V)
A = fd.assemble(a, bcs=bc)
b = fd.assemble(L, bcs=bc)
solver = fd.LinearSolver(A, solver_parameters=self.solver_parameters)
# Solve the Signed distance equation with Picard iteration
bc.apply(phi_sol)
init_res = residual_phi(phi_sol)
res = 1e10
it = 0
while res > init_res or it < iters:
solver.solve(phi_sol, b)
self.phi.assign(phi_sol)
b = fd.assemble(L, bcs=bc, tensor=b)
it += 1
res = residual_phi(phi_sol)
if res > init_res:
fd.warning(
f"Residual in signed distance function increased: {res}, before: {init_res}"
)
return self.phi
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 == "both" and ny == 1 and quadrilateral:
return OneElementThickMesh(nx, Lx, Ly)
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)