Normal (#768)

Generalize function.normal so that it can be used to evaluate the
boundary normal of manifolds in higher dimensional space.

nutils/evaluable.py

@@ -1094,28 +1094,28 @@ def evalf(evalargs):
         return weights
 
 
-class Normal(Array):
-    'normal'
+class Orthonormal(Array):
+    'make a vector orthonormal to a subspace'
 
-    __slots__ = 'lgrad',
+    __slots__ = '_basis', '_vector'
 
-    def __init__(self, lgrad: Array):
-        assert isinstance(lgrad, Array) and lgrad.ndim >= 2 and _equals_simplified(lgrad.shape[-2], lgrad.shape[-1]) and lgrad.dtype != complex, f'lgrad={lgrad!r}'
-        self.lgrad = lgrad
-        super().__init__(args=(lgrad,), shape=lgrad.shape[:-1], dtype=float)
+    def __init__(self, basis: Array, vector: Array):
+        assert isinstance(basis, Array) and basis.ndim >= 2 and basis.dtype != complex, f'basis={basis!r}'
+        assert isinstance(vector, Array) and vector.ndim >= 1 and vector.dtype != complex, f'vector={vector!r}'
+        assert equalshape(basis.shape[:-1], vector.shape)
+        self._basis = basis
+        self._vector = vector
+        super().__init__(args=(basis, vector), shape=vector.shape, dtype=float)
 
     def _simplified(self):
         if _equals_scalar_constant(self.shape[-1], 1):
-            return Sign(Take(self.lgrad, constant(0)))
-        unaligned, where = unalign(self.lgrad, naxes=self.ndim - 1)
+            return Sign(self._vector)
+        basis, vector, where = unalign(self._basis, self._vector, naxes=self.ndim - 1)
         if len(where) < self.ndim - 1:
-            return align(Normal(unaligned), (*where, self.ndim - 1), self.shape)
+            return align(Orthonormal(basis, vector), (*where, self.ndim - 1), self.shape)
 
     @staticmethod
-    def evalf(lgrad):
-        n = lgrad[..., -1]
-        # orthonormalize n to G
-        G = lgrad[..., :-1]
+    def evalf(G, n):
         GG = numpy.einsum('...ki,...kj->...ij', G, G)
         v1 = numpy.einsum('...ij,...i->...j', G, n)
         v2 = numpy.linalg.solve(GG, v1)
@@ -1125,9 +1125,47 @@ def evalf(lgrad):
     def _derivative(self, var, seen):
         if _equals_scalar_constant(self.shape[-1], 1):
             return zeros(self.shape + var.shape)
-        G = self.lgrad[..., :-1]
+
+        # definitions:
+        #
+        # P := I - G (G^T G)^-1 G^T (orthogonal projector)
+        # n := P v (orthogonal projection of v)
+        # N := n / |n| (self: orthonormal projection of v)
+        #
+        # identities:
+        #
+        #   P^T = P          N^T N = 1
+        #   P P = P          P N = N
+        #   P G = P Q = 0    G^T N = Q^T N = 0
+        #
+        # derivatives:
+        #
+        # P' = Q P + P Q^T where Q := -G (G^T G)^-1 G'^T
+        # n' = P' v + P v'
+        #    = Q n + P (Q^T v + v')
+        # N' = (I - N N^T) n' / |n|
+        #    = (I - N N^T) (Q n / |n| + P (Q^T v + v') / |n|)
+        #    = Q N + (P - N N^T) (Q^T v + v') / |n|
+
+        G = self._basis
         invGG = inverse(einsum('Aki,Akj->Aij', G, G))
-        return -einsum('Ail,Alj,Ak,AkjB->AiB', G, invGG, self, derivative(G, var, seen))
+
+        Q = -einsum('Aim,Amn,AjnB->AijB', G, invGG, derivative(G, var, seen))
+        QN = einsum('Ai,AjiB->AjB', self, Q)
+
+        if _equals_simplified(G.shape[-1], G.shape[-2] - 1): # dim(kern(G)) = 1
+            # In this situation, since N is a basis for the kernel of G, we
+            # have the identity P == N N^T which cancels the entire second term
+            # of N' along with any reference to v', reducing it to N' = Q N.
+            return QN
+
+        v = self._vector
+        P = Diagonalize(ones(self.shape)) - einsum('Aim,Amn,Ajn->Aij', G, invGG, G)
+        Z = P - einsum('Ai,Aj->Aij', self, self) # P - N N^T
+
+        return QN + einsum('A,AiB->AiB',
+            power(einsum('Ai,Aij,Aj->A', v, P, v), -.5),
+            einsum('Aij,AjB->AiB', Z, einsum('Ai,AijB->AjB', v, Q) + derivative(v, var, seen)))
 
 
 class Constant(Array):

nutils/function.py

@@ -1139,8 +1139,8 @@ def lower(self, args: LowerArgs) -> evaluable.Array:
         geom = self._geom.lower(args)
         spaces_dim = builtins.sum(args.transform_chains[space][0][0].todims for space in self._geom.spaces)
         normal_dim = spaces_dim - builtins.sum(args.transform_chains[space][0][0].fromdims for space in self._geom.spaces)
-        if self._geom.shape[-1] != spaces_dim:
-            raise ValueError('The dimension of geometry must equal the sum of the dimensions of the given spaces.')
+        if self._geom.shape[-1] < spaces_dim:
+            raise ValueError('The dimension of geometry must equal or larger than the sum of the dimensions of the given spaces.')
         if normal_dim == 0:
             raise ValueError('Cannot compute the normal because the dimension of the normal space is zero.')
         elif normal_dim > 1:
@@ -1158,9 +1158,9 @@ def lower(self, args: LowerArgs) -> evaluable.Array:
                 assert normal is None and chain.todims == chain.fromdims + 1
                 basis = evaluable.einsum('Aij,jk->Aik', rgrad, evaluable.TransformBasis(chain, index))
                 tangents.append(basis[..., :chain.fromdims])
-                normal = basis[..., chain.fromdims:]
+                normal = basis[..., chain.fromdims]
         assert normal is not None
-        return evaluable.Normal(evaluable.concatenate((*tangents, normal), axis=-1))
+        return evaluable.Orthonormal(evaluable.concatenate(tangents, axis=-1), normal)
 
 
 class _ExteriorNormal(Array):

tests/test_evaluable.py

@@ -554,9 +554,10 @@ def _check(name, op, n_op, *arg_values, hasgrad=True, zerograd=False, ndim=2):
 _check('take-duplicate', lambda f: evaluable.Take(f, evaluable.constant([0, 3, 0])), lambda a: a[:, [0, 3, 0]], ANY(2, 4))
 _check('choose', lambda a, b, c: evaluable.Choose(a % 2, (b, c)), lambda a, b, c: numpy.choose(a % 2, [b, c]), INT(3, 3), ANY(3, 3), ANY(3, 3))
 _check('slice', lambda a: evaluable.asarray(a)[::2], lambda a: a[::2], ANY(5, 3))
-_check('normal1d', lambda a: evaluable.Normal(a), lambda a: numpy.sign(a[..., 0]), NZ(3, 1, 1))
-_check('normal2d', lambda a: evaluable.Normal(a), lambda a: numpy.stack([Q[:, -1]*numpy.sign(R[-1, -1]) for ai in a for Q, R in [numpy.linalg.qr(ai, mode='complete')]], axis=0), POS(1, 2, 2)+numpy.eye(2))
-_check('normal3d', lambda a: evaluable.Normal(a), lambda a: numpy.stack([Q[:, -1]*numpy.sign(R[-1, -1]) for ai in a for Q, R in [numpy.linalg.qr(ai, mode='complete')]], axis=0), POS(2, 3, 3)+numpy.eye(3))
+_check('normal1d', evaluable.Orthonormal, lambda G, a: numpy.sign(a), numpy.zeros([3,1,0]), NZ(3, 1))
+_check('normal2d', evaluable.Orthonormal, lambda G, a: numeric.normalize(a - numeric.normalize(G[...,0]) * ((a * numeric.normalize(G[...,0])).sum(-1))[...,numpy.newaxis]), POS(3, 2, 1), ANY(3, 2))
+_check('normal3d', evaluable.Orthonormal, lambda G, a: numeric.normalize(a - numpy.einsum('pij,pj->pi', G, numpy.linalg.solve(numpy.einsum('pki,pkj->pij', G, G), numpy.einsum('pij,pi->pj', G, a)))), POS(2, 3, 2) + numpy.eye(3)[:,:2], ANY(2, 3))
+_check('normalmanifold', evaluable.Orthonormal, lambda G, a: numeric.normalize(a - numpy.einsum('pij,pj->pi', G, numpy.linalg.solve(numpy.einsum('pki,pkj->pij', G, G), numpy.einsum('pij,pi->pj', G, a)))), POS(2, 3, 1), ANY(2, 3))
 _check('loopsum1', lambda: evaluable.loop_sum(evaluable.loop_index('index', 3), evaluable.loop_index('index', 3)), lambda: numpy.array(3))
 _check('loopsum2', lambda a: evaluable.loop_sum(a, evaluable.loop_index('index', 2)), lambda a: 2*a, ANY(3, 4, 2, 4))
 _check('loopsum3', lambda a: evaluable.loop_sum(evaluable.get(a, 0, evaluable.loop_index('index', 3)), evaluable.loop_index('index', 3)), lambda a: numpy.sum(a, 0), ANY(3, 4, 2, 4))

tests/test_function.py

@@ -945,6 +945,18 @@ def test_normal_3d(self):
         for bnd, n in ('right', [1, 0]), ('left', [-1, 0]), ('top', [0, 1]), ('bottom', [0, -1]):
             self.assertEvalAlmostEqual(domain.boundary[bnd], self.normal(x), numpy.array(n)[numpy.newaxis, :, numpy.newaxis])
 
+    def test_normal_manifold(self):
+        domain, geom = mesh.rectilinear([1]*2)
+        x = numpy.stack([geom[0], geom[1], geom[0]**2 - geom[1]**2])
+        n = self.normal(x) # boundary normal
+        N = self.normal(x, geom) # exterior normal
+        k = -.5 * self.div(N, x, -1) # mean curvature
+        dA = function.jacobian(x, 2)
+        dL = function.jacobian(x, 1)
+        v1 = domain.integrate(2 * k * N * dA, degree=16)
+        v2 = domain.boundary.integrate(n * dL, degree=16)
+        self.assertAllAlmostEqual(v1, v2) # divergence theorem in curved space
+
     def test_dotnorm(self):
         domain, x = mesh.rectilinear([1]*2)
         x = 2*x-0.5
@@ -994,11 +1006,11 @@ def test_different_argument_dtype(self):
            ngrad=function.ngrad,
            nsymgrad=function.nsymgrad)
 derivative('method',
-           normal=lambda geom: function.Array.cast(geom).normal(),
+           normal=lambda geom, refgeom=None: function.Array.cast(geom).normal(refgeom),
            tangent=lambda geom, vec: function.Array.cast(geom).tangent(vec),
            dotnorm=lambda vec, geom: function.Array.cast(vec).dotnorm(geom),
            grad=lambda arg, geom: function.Array.cast(arg).grad(geom),
-           div=lambda arg, geom: function.Array.cast(arg).div(geom),
+           div=lambda arg, geom, ndims=0: function.Array.cast(arg).div(geom, ndims),
            curl=lambda arg, geom: function.Array.cast(arg).curl(geom),
            laplace=lambda arg, geom: function.Array.cast(arg).laplace(geom),
            symgrad=lambda arg, geom: function.Array.cast(arg).symgrad(geom),