Allow mixing of logical and physical coordinates in expressions - V2

sympde/topology/mapping.py

@@ -858,6 +858,157 @@ def eval(cls, F, v):
         v = M*v
         return Tuple(*v)
 
+#==============================================================================
+def get_kind_type(kind):
+
+    from .datatype import SpaceType, dtype_space_registry
+
+    if kind is None:
+        kind = 'undefined'
+
+    if isinstance(kind, str):
+        kind_str = kind.lower()
+        assert kind_str in ['h1', 'hcurl', 'hdiv', 'l2', 'undefined']
+        kind_type = dtype_space_registry[kind_str]
+
+    elif isinstance(kind, SpaceType):
+        kind_type = kind
+
+    else:
+        raise TypeError('Expecting kind to be string or SpaceType')
+
+    return kind_type
+
+#==============================================================================
+class PullBack(Expr):
+    is_commutative = False
+
+    def __new__(cls, expr, *, mapping=None, domain=None, kind='h1', evaluate=False):
+
+        kind = get_kind_type(kind)
+
+        if evaluate:
+            return cls.eval(expr, mapping=mapping, domain=domain, kind=kind)
+
+        obj = Expr.__new__(cls, expr, mapping, kind)
+        return obj
+
+    @property
+    def expr(self):
+        return self.args[0]
+
+    @property
+    def mapping(self):
+        return self.args[1]
+
+    @property
+    def kind(self):
+        return self.args[2]
+
+    @classmethod
+    def eval(cls, expr, domain, mapping, kind):
+
+        kind = get_kind_type(kind)
+
+        if isinstance(expr, (ScalarFunction, VectorFunction)):
+            u       = expr
+            u_hat   = get_logical_test_function(u)
+            space   = u.space
+            mapping = space.domain.mapping
+
+            # By definition u is the PushForward of u_hat. (This is implicit.)
+            # Hence, if the PullBack and the PushForward are of the same kind,
+            # they cancel each other.
+            if kind == u.space.kind:
+                return u_hat
+
+            # If the kinds do not match, continue processing...
+
+        if isinstance(expr, PushForward):
+            # If PullBack and PushForward are of the same kind, they simplify.
+            if kind == expr.kind:
+                return expr.arg
+
+            # Otherwise, evaluate PushForward explicitly and redefine expression
+            else:
+                expr = PushForward(
+                    expr.arg,
+                    mapping = expr.mapping,
+                    kind = expr.kind,
+                    evaluate = True
+                )
+
+        # In all other cases use LogicalExpr.eval() and apply inverse formulas
+        # TODO: get domain
+        log_expr = LogicalExpr.eval(expr, domain)
+
+        # TODO: verify what this means
+        if space.is_broken:
+            assert mapping is not None
+
+        if kind in (H1SpaceType(), UndefinedSpaceType()):
+            return log_expr
+
+        J = mapping.jacobian
+
+        if kind is HcurlSpaceType():
+            return J.T * log_expr
+
+        if kind is HdivSpaceType():
+            return J.det() * J.inv() * log_expr
+
+        if kind is L2SpaceType():
+            return J.det() * log_expr
+
+        raise ValueError(f'Cannot compute Pull-Back of kind = {kind}')
+
+#==============================================================================
+class PushForward(Expr):
+    is_commutative = False
+
+    def __new__(cls, expr, *, mapping=None, kind='h1', evaluate=False):
+
+        kind = get_kind_type(kind)
+
+        if evaluate:
+            return cls.eval(expr, mapping=mapping, kind=kind)
+
+        obj = Expr.__new__(cls, expr, mapping, kind)
+        return obj
+
+    @property
+    def arg(self):
+        return self.args[0]
+
+    @property
+    def mapping(self):
+        return self.args[1]
+
+    @property
+    def kind(self):
+        return self.args[2]
+
+    @classmethod
+    def eval(cls, expr, mapping, kind):
+
+        kind = get_kind_type(kind)
+
+        if kind in (H1SpaceType(), UndefinedSpaceType()):
+            return expr
+
+        J = mapping.jacobian
+
+        if kind is HcurlSpaceType():
+            return J.inv().T * expr
+
+        if kind is HdivSpaceType():
+            return (J / J.det()) * expr
+
+        if kind is L2SpaceType():
+            return expr / J.det()
+
+        raise ValueError(f'Cannot compute Push-Forward of kind = {kind}')
+
 #==============================================================================
 class LogicalExpr(CalculusFunction):
 
@@ -898,7 +1049,9 @@ def eval(cls, expr, domain, **options):
         from sympde.expr.expr import BilinearForm, LinearForm, BasicForm, Norm
         from sympde.expr.expr import Integral
 
-        types = (ScalarFunction, VectorFunction, DifferentialOperator, Trace, Integral)
+        # Expression types for which special rules are applied
+        types = (ScalarFunction, VectorFunction, DifferentialOperator,
+                Trace, Integral, PushForward)
 
         mapping   = domain.mapping
         dim       = domain.dim
@@ -910,6 +1063,7 @@ def eval(cls, expr, domain, **options):
 
         if not has(expr, types):
             if has(expr, DiffOperator):
+                # TODO: do not rewrite
                 return cls( expr, domain, evaluate=False)
             else:
                 syms = symbols(ph_coords[:dim])
@@ -970,94 +1124,127 @@ def eval(cls, expr, domain, **options):
             newexpr = newexpr.expr.subs(test, PlusInterfaceOperator(test))
             return newexpr
 
+#        elif isinstance(expr, (VectorFunction, ScalarFunction)):
+#            return PullBack(expr, mapping).expr
+
         elif isinstance(expr, (VectorFunction, ScalarFunction)):
-            return PullBack(expr, mapping).expr
+            u       = expr
+            u_hat   = get_logical_test_function(u)
+            space   = u.space
+            kind    = space.kind
+
+            # TODO: verify what this means
+            if space.is_broken:
+                assert mapping is not None
+            else:
+                mapping = space.domain.mapping
+
+            return PushForward(u_hat, mapping=mapping, kind=kind, evaluate=True)
+
+        elif isinstance(expr, PushForward):
+            return PushForward(expr.arg, mapping=expr.mapping, kind=expr.kind,
+                evaluate=True)
 
         elif isinstance(expr, Transpose):
             arg = cls(expr.arg, domain)
             return Transpose(arg)
-
-        elif isinstance(expr, grad):
-            arg = expr.args[0]
-            if isinstance(mapping, InterfaceMapping):
-                if isinstance(arg, MinusInterfaceOperator):
-                    a     = arg.args[0]
-                    mapping = mapping.minus
-                elif isinstance(arg, PlusInterfaceOperator):
-                    a = arg.args[0]
-                    mapping = mapping.plus
-                else:
-                    raise TypeError(arg)
 
-                arg = type(arg)(cls.eval(a, domain))
-            else:
-                arg = cls.eval(arg, domain)
-
-            return mapping.jacobian.inv().T*grad(arg)
+        elif isinstance(expr, grad):
+            arg_log = PullBack(expr.args[0], mapping=mapping, domain=domain, kind='h1', evaluate=True)
+            return PushForward(grad(arg_log), mapping=mapping, kind='hcurl', evaluate=True)
 
         elif isinstance(expr, curl):
-            arg = expr.args[0]
-            if isinstance(mapping, InterfaceMapping):
-                if isinstance(arg, MinusInterfaceOperator):
-                    arg     = arg.args[0]
-                    mapping = mapping.minus
-                elif isinstance(arg, PlusInterfaceOperator):
-                    arg = arg.args[0]
-                    mapping = mapping.plus
-                else:
-                    raise TypeError(arg)
-
-            if isinstance(arg, VectorFunction):
-                arg = PullBack(arg, mapping)
-            else:
-                arg = cls.eval(arg, domain)
-
-            if isinstance(arg, PullBack) and isinstance(arg.kind, HcurlSpaceType):
-                J   = mapping.jacobian
-                arg = arg.test
-                if isinstance(expr.args[0], (MinusInterfaceOperator, PlusInterfaceOperator)):
-                    arg = type(expr.args[0])(arg)
-                if expr.is_scalar:
-                    return (1/J.det())*curl(arg)
-
-                return (J/J.det())*curl(arg)
-            else:
-                raise NotImplementedError('TODO')
+            caller  = lambda e, cls=cls: cls.eval(e, cls.domain)
+            arg_log = PullBack(expr.args[0], mapping=mapping, domain=domain, kind='hcurl', evaluate=True)
+            return PushForward(curl(arg_log), mapping=mapping, kind='hdiv', evaluate=True)
 
         elif isinstance(expr, div):
-            arg = expr.args[0]
-            if isinstance(mapping, InterfaceMapping):
-                if isinstance(arg, MinusInterfaceOperator):
-                    arg     = arg.args[0]
-                    mapping = mapping.minus
-                elif isinstance(arg, PlusInterfaceOperator):
-                    arg = arg.args[0]
-                    mapping = mapping.plus
-                else:
-                    raise TypeError(arg)
+            caller  = lambda e, cls=cls: cls.eval(e, cls.domain)
+            arg_log = PullBack(expr.args[0], mapping=mapping, domain=domain, kind='hdiv', evaluate=True)
+            return PushForward(div(arg_log), mapping=mapping, kind='l2', evaluate=True)
 
-            if isinstance(arg, (ScalarFunction, VectorFunction)):
-                arg = PullBack(arg, mapping)
-            else:
-
-                arg = cls.eval(arg, domain)
-
-            if isinstance(arg, PullBack) and isinstance(arg.kind, HdivSpaceType):
-                J   = mapping.jacobian
-                arg = arg.test
-                if isinstance(expr.args[0], (MinusInterfaceOperator, PlusInterfaceOperator)):
-                    arg = type(expr.args[0])(arg)
-                return (1/J.det())*div(arg)
-            elif isinstance(arg, PullBack):
-                return SymbolicTrace(mapping.jacobian.inv().T*grad(arg.test))
-            else:
-                raise NotImplementedError('TODO')
-
-        elif isinstance(expr, laplace):
-            arg = expr.args[0]
-            v   = cls.eval(grad(arg), domain)
-            v   = mapping.jacobian.inv().T*grad(v)
-            return SymbolicTrace(v)
+#        elif isinstance(expr, grad):
+#            arg = expr.args[0]
+#            print(f'arg = {arg}')
+#            if isinstance(mapping, InterfaceMapping):
+#                if isinstance(arg, MinusInterfaceOperator):
+#                    a     = arg.args[0]
+#                    mapping = mapping.minus
+#                elif isinstance(arg, PlusInterfaceOperator):
+#                    a = arg.args[0]
+#                    mapping = mapping.plus
+#                else:
+#                    raise TypeError(arg)
+#
+#                arg_log = type(arg)(cls.eval(a, domain))
+#            else:
+#                arg_log = cls.eval(arg, domain)
+#
+#            return mapping.jacobian.inv().T * grad(arg_log)
+#
+#        elif isinstance(expr, curl):
+#            arg = expr.args[0]
+#            if isinstance(mapping, InterfaceMapping):
+#                if isinstance(arg, MinusInterfaceOperator):
+#                    arg     = arg.args[0]
+#                    mapping = mapping.minus
+#                elif isinstance(arg, PlusInterfaceOperator):
+#                    arg = arg.args[0]
+#                    mapping = mapping.plus
+#                else:
+#                    raise TypeError(arg)
+#
+#            if isinstance(arg, VectorFunction):
+#                arg = PullBack(arg, mapping)
+#            else:
+#                arg = cls.eval(arg, domain)
+#
+#            if isinstance(arg, PullBack) and isinstance(arg.kind, HcurlSpaceType):
+#                J   = mapping.jacobian
+#                arg = arg.test
+#                if isinstance(expr.args[0], (MinusInterfaceOperator, PlusInterfaceOperator)):
+#                    arg = type(expr.args[0])(arg)
+#                if expr.is_scalar:
+#                    return (1/J.det())*curl(arg)
+#
+#                return (J/J.det())*curl(arg)
+#            else:
+#                raise NotImplementedError('TODO')
+#
+#        elif isinstance(expr, div):
+#            arg = expr.args[0]
+#            if isinstance(mapping, InterfaceMapping):
+#                if isinstance(arg, MinusInterfaceOperator):
+#                    arg     = arg.args[0]
+#                    mapping = mapping.minus
+#                elif isinstance(arg, PlusInterfaceOperator):
+#                    arg = arg.args[0]
+#                    mapping = mapping.plus
+#                else:
+#                    raise TypeError(arg)
+#
+#            if isinstance(arg, (ScalarFunction, VectorFunction)):
+#                arg = PullBack(arg, mapping)
+#            else:
+#
+#                arg = cls.eval(arg, domain)
+#
+#            if isinstance(arg, PullBack) and isinstance(arg.kind, HdivSpaceType):
+#                J   = mapping.jacobian
+#                arg = arg.test
+#                if isinstance(expr.args[0], (MinusInterfaceOperator, PlusInterfaceOperator)):
+#                    arg = type(expr.args[0])(arg)
+#                return (1/J.det())*div(arg)
+#            elif isinstance(arg, PullBack):
+#                return SymbolicTrace(mapping.jacobian.inv().T*grad(arg.test))
+#            else:
+#                raise NotImplementedError('TODO')
+#
+#        elif isinstance(expr, laplace):
+#            arg = expr.args[0]
+#            v   = cls.eval(grad(arg), domain)
+#            v   = mapping.jacobian.inv().T*grad(v)
+#            return SymbolicTrace(v)
 
 #        elif isinstance(expr, hessian):
 #           arg = expr.args[0]