Merge pull request #12 from pyroteus/joe/ma-fixed-bc

Fixed boundary option for Monge-Ampere

movement/laplacian.py

@@ -1,5 +1,5 @@
 import firedrake
-from firedrake import PETSc
+from firedrake.petsc import PETSc
 import ufl
 import movement.solver_parameters as solver_parameters
 from movement.mover import PrimeMover

movement/monge_ampere.py

@@ -1,5 +1,5 @@
 import firedrake
-from firedrake import PETSc
+from firedrake.petsc import PETSc
 from pyadjoint import no_annotations
 import ufl
 import numpy as np
@@ -61,6 +61,7 @@ def __init__(self, mesh, monitor_function, **kwargs):
         :kwarg maxiter: maximum number of iterations for the relaxation
         :kwarg rtol: relative tolerance for the residual
         :kwarg dtol: divergence tolerance for the residual
+        :kwarg fix_boundary_nodes: should all boundary nodes remain fixed?
         """
         if monitor_function is None:
             raise ValueError("Please supply a monitor function")
@@ -70,6 +71,7 @@ def __init__(self, mesh, monitor_function, **kwargs):
         self.maxiter = kwargs.pop('maxiter', 1000)
         self.rtol = kwargs.pop('rtol', 1.0e-08)
         self.dtol = kwargs.pop('dtol', 2.0)
+        self.fix_boundary_nodes = kwargs.pop('fix_boundary_nodes', False)
         super().__init__(mesh, monitor_function=monitor_function)
 
         # Create function spaces
@@ -168,37 +170,50 @@ def l2_projector(self):
         n = ufl.FacetNormal(self.mesh)
         bcs = []
         for i in self.mesh.exterior_facets.unique_markers:
-            _n = [firedrake.assemble(n[j]*self.ds(i)) for j in range(self.dim)]
+            if self.fix_boundary_nodes:
+                bcs.append(firedrake.DirichletBC(self.P1_vec, 0, i))
+                continue
+
+            # Check for axis-aligned boundaries
+            _n = [firedrake.assemble(abs(n[j])*self.ds(i)) for j in range(self.dim)]
             if np.allclose(_n, 0.0):
                 raise ValueError(f"Invalid normal vector {_n}")
             else:
                 if self.dim != 2:
                     raise NotImplementedError  # TODO
                 if np.isclose(_n[0], 0.0):
                     bcs.append(firedrake.DirichletBC(self.P1_vec.sub(1), 0, i))
+                    continue
                 elif np.isclose(_n[1], 0.0):
                     bcs.append(firedrake.DirichletBC(self.P1_vec.sub(0), 0, i))
-                else:
-                    # Enforce no mesh movement normal to boundaries
-                    a_bc = ufl.dot(v_cts, n)*ufl.dot(u_cts, n)*self.ds
-                    L_bc = ufl.dot(v_cts, n)*firedrake.Constant(0.0)*self.ds
-                    bcs.append(firedrake.EquationBC(a_bc == L_bc, self.grad_phi, 'on_boundary'))
-
-                    # Allow tangential movement, but only up until the end of boundary segments
-                    s = ufl.perp(n)
-                    a_bc = ufl.dot(v_cts, s)*ufl.dot(u_cts, s)*self.ds
-                    L_bc = ufl.dot(v_cts, s)*ufl.dot(ufl.grad(self.phi_old), s)*self.ds
-                    edges = set(self.mesh.exterior_facets.unique_markers)
-                    if len(edges) > 1:  # NOTE: Assumes that all straight line segments are uniquely tagged
-                        corners = [(i, j) for i in edges for j in edges.difference([i])]
-                        bbc = firedrake.DirichletBC(self.P1_vec, 0, corners)
-                    else:
-                        bbc = None
-                    bcs.append(firedrake.EquationBC(a_bc == L_bc, self.grad_phi, 'on_boundary', bcs=bbc))
+                    continue
+
+            # Enforce no mesh movement normal to boundaries
+            a_bc = ufl.dot(v_cts, n)*ufl.dot(u_cts, n)*self.ds
+            L_bc = ufl.dot(v_cts, n)*firedrake.Constant(0.0)*self.ds
+            bcs.append(firedrake.EquationBC(a_bc == L_bc, self.grad_phi, 'on_boundary'))
+
+            # Allow tangential movement, but only up until the end of boundary segments
+            s = ufl.perp(n)
+            a_bc = ufl.dot(v_cts, s)*ufl.dot(u_cts, s)*self.ds
+            L_bc = ufl.dot(v_cts, s)*ufl.dot(ufl.grad(self.phi_old), s)*self.ds
+            edges = set(self.mesh.exterior_facets.unique_markers)
+            if len(edges) == 0:
+                bbc = None  # Periodic case
+            else:
+                from warnings import warn
+                warn('Have you checked that all straight line segments are uniquely tagged?')
+                corners = [(i, j) for i in edges for j in edges.difference([i])]
+                bbc = firedrake.DirichletBC(self.P1_vec, 0, corners)
+            bcs.append(firedrake.EquationBC(a_bc == L_bc, self.grad_phi, 'on_boundary', bcs=bbc))
 
         # Create solver
         problem = firedrake.LinearVariationalProblem(a, L, self._grad_phi, bcs=bcs)
-        sp = {"ksp_type": "cg"}
+        sp = {
+            "ksp_type": "cg",
+            "pc_type": "bjacobi",
+            "sub_pc_type": "ilu",
+        }
         self._l2_projector = firedrake.LinearVariationalSolver(problem, solver_parameters=sp)
         return self._l2_projector
 

movement/spring.py

@@ -1,5 +1,5 @@
 import firedrake
-from firedrake import PETSc
+from firedrake.petsc import PETSc
 import ufl
 import numpy as np
 import movement.solver_parameters as solver_parameters

test/test_monge_ampere.py

@@ -1,13 +1,19 @@
 from movement import *
 from monitors import *
 import pytest
+import numpy as np
 
 
 @pytest.fixture(params=['relaxation', 'quasi_newton'])
 def method(request):
     return request.param
 
 
+@pytest.fixture(params=[True, False])
+def fix_boundary(request):
+    return request.param
+
+
 def test_uniform_monitor(method, exports=False):
     """
     Test that the mesh mover converges in one
@@ -84,3 +90,30 @@ def test_change_monitor(method, exports=False):
     if exports:
         File("outputs/const.pvd").write(mover.phi, mover.sigma)
     assert np.allclose(coords, mesh.coordinates.dat.data, atol=tol)
+
+
+def test_bcs(method, fix_boundary):
+    """
+    Test that domain boundaries are fixed by
+    the Monge-Ampere movers.
+    """
+    n = 20
+    mesh = UnitSquareMesh(n, n)
+    one = Constant(1.0)
+    bnd = assemble(one*ds(domain=mesh))
+    bnodes = DirichletBC(mesh.coordinates.function_space(), 0, 'on_boundary').nodes
+    bnd_coords = mesh.coordinates.dat.data.copy()[bnodes]
+
+    # Adapt to a ring monitor
+    mover = MongeAmpereMover(mesh, ring_monitor, method=method,
+                             fix_boundary_nodes=fix_boundary)
+    mover.move()
+
+    # Check boundary lengths are preserved
+    bnd_new = assemble(one*ds(domain=mesh))
+    assert np.isclose(bnd, bnd_new)
+
+    # Check boundaries are indeed fixed
+    if fix_boundary:
+        bnd_coords_new = mesh.coordinates.dat.data[bnodes]
+        assert np.allclose(bnd_coords, bnd_coords_new)