Move around parts of FML

gusto/__init__.py

@@ -1,3 +1,18 @@
+# This monkey patch is required so that Gusto can redefine the UFL perp.
+# It should be removed if the UFL perp gets its own operator
+# The main thing to be aware of is if we could now do (and should avoid doing)
+# from gusto import _monkey_patch_ufl()
+def _monkey_patch_ufl():
+    from ufl.algorithms.apply_algebra_lowering import LowerCompoundAlgebra
+
+    def perp(self, o, a):
+        from firedrake import as_vector
+        return as_vector([-a[1], a[0]])
+    LowerCompoundAlgebra.perp = perp
+
+
+_monkey_patch_ufl()
+
 from gusto.active_tracers import *                   # noqa
 from gusto.common_forms import *                     # noqa
 from gusto.configuration import *                    # noqa

gusto/domain.py

@@ -6,9 +6,9 @@
 
 from gusto.coordinates import Coordinates
 from gusto.function_spaces import Spaces, check_degree_args
+from gusto.perp import perp
 from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross,
-                       as_vector, inner, interpolate, VectorFunctionSpace,
-                       Function)
+                       inner, interpolate, VectorFunctionSpace, Function)
 import numpy as np
 
 
@@ -108,7 +108,7 @@ def __init__(self, mesh, dt, family, degree=None,
             kvec[dim-1] = 1.0
             self.k = Constant(kvec)
             if dim == 2:
-                self.perp = lambda u: as_vector([-u[1], u[0]])
+                self.perp = perp
 
         # -------------------------------------------------------------------- #
         # Set up coordinates

gusto/equations.py

@@ -7,11 +7,10 @@
                        DirichletBC, conditional, SpatialCoordinate,
                        split, Constant, action)
 from gusto.fields import PrescribedFields
-from gusto.fml.form_manipulation_labelling import Term, all_terms, keep, drop, Label
-from gusto.labels import (subject, time_derivative, transport, prognostic,
-                          replace_subject, linearisation,
-                          name, pressure_gradient, coriolis, perp,
-                          replace_trial_function, hydrostatic)
+from gusto.fml import (Term, all_terms, keep, drop, Label, subject, name,
+                       replace_subject, replace_trial_function)
+from gusto.labels import (time_derivative, transport, prognostic, hydrostatic,
+                          linearisation, pressure_gradient, coriolis)
 from gusto.thermodynamics import exner_pressure
 from gusto.common_forms import (advection_form, continuity_form,
                                 vector_invariant_form, kinetic_energy_form,
@@ -700,17 +699,11 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None,
             V = FunctionSpace(domain.mesh, 'CG', 1)
             f = self.prescribed_fields('coriolis', V).interpolate(fexpr)
             coriolis_form = coriolis(subject(
-                prognostic(f*inner(domain.perp(u), w)*dx, 'u'), self.X))
-            if not domain.on_sphere:
-                coriolis_form = perp(coriolis_form, domain.perp)
+                prognostic(f*inner(domain.perp(u), w)*dx, "u"), self.X))
             # Add linearisation
             if self.linearisation_map(coriolis_form.terms[0]):
-                linear_coriolis = perp(
-                    coriolis(
-                        subject(prognostic(f*inner(domain.perp(u_trial), w)*dx, 'u'), self.X)
-                    ), domain.perp)
-                if not domain.on_sphere:
-                    linear_coriolis = perp(linear_coriolis, domain.perp)
+                linear_coriolis = coriolis(
+                    subject(prognostic(f*inner(domain.perp(u_trial), w)*dx, 'u'), self.X))
                 coriolis_form = linearisation(coriolis_form, linear_coriolis)
             residual += coriolis_form
 

gusto/fml/__init__.py

@@ -1 +1,2 @@
-from gusto.fml.form_manipulation_labelling import *        # noqa
+from gusto.fml.form_manipulation_language import *        # noqa
+from gusto.fml.replacement import *                       # noqa

gusto/fml/form_manipulation_labelling.py → gusto/fml/form_manipulation_language.py

@@ -3,18 +3,24 @@
 import ufl
 import functools
 import operator
-from firedrake import Constant
+from firedrake import Constant, Function
+
+
+__all__ = ["Label", "Term", "LabelledForm", "identity", "drop", "all_terms",
+           "keep", "subject", "name"]
 
 # ---------------------------------------------------------------------------- #
 # Core routines for filtering terms
 # ---------------------------------------------------------------------------- #
-
 identity = lambda t: t
 drop = lambda t: None
 all_terms = lambda t: True
 keep = identity
 
 
+# ---------------------------------------------------------------------------- #
+# Term class
+# ---------------------------------------------------------------------------- #
 class Term(object):
     """A Term object contains a form and its labels."""
 
@@ -74,7 +80,9 @@ def __add__(self, other):
         Returns:
             :class:`LabelledForm`: a labelled form containing the terms.
         """
-        if other is None:
+        if self is NullTerm:
+            return other
+        if other is None or other is NullTerm:
             return self
         elif isinstance(other, Term):
             return LabelledForm(self, other)
@@ -112,10 +120,6 @@ def __mul__(self, other):
         Returns:
             :class:`Term`: the product of the term with the quantity.
         """
-        if type(other) in (float, int):
-            other = Constant(other)
-        elif type(other) not in [Constant, ufl.algebra.Product]:
-            return NotImplemented
         return Term(other*self.form, self.labels)
 
     __rmul__ = __mul__
@@ -133,16 +137,16 @@ def __truediv__(self, other):
         Returns:
             :class:`Term`: the quotient of the term divided by the quantity.
         """
-        if type(other) in (float, int, Constant, ufl.algebra.Product):
-            other = Constant(1.0 / other)
-            return self * other
-        else:
-            return NotImplemented
+        return self * (Constant(1.0) / other)
 
 
+# This is necessary to be the initialiser for functools.reduce
 NullTerm = Term(None)
 
 
+# ---------------------------------------------------------------------------- #
+# Labelled form class
+# ---------------------------------------------------------------------------- #
 class LabelledForm(object):
     """
     A form, broken down into terms that pair individual forms with labels.
@@ -208,33 +212,24 @@ def __sub__(self, other):
             return LabelledForm(*self, Constant(-1.)*other)
         elif type(other) is LabelledForm:
             return LabelledForm(*self, *[Constant(-1.)*t for t in other])
-        elif type(other) is ufl.algebra.Product:
-            return LabelledForm(*self, Term(Constant(-1.)*other))
         elif other is None:
             return self
         else:
-            return NotImplemented
+            # Make new Term for other and subtract it
+            return LabelledForm(*self, Term(Constant(-1.)*other))
 
     def __mul__(self, other):
         """
         Multiplies this labelled form by another quantity.
 
         Args:
             other (float, :class:`Constant` or :class:`ufl.algebra.Product`):
-                the quantity to multiply this labelled form by. If it is a float
-                or int then it is converted to a :class:`Constant` before the
-                multiplication. All terms in the form are multiplied.
+                the quantity to multiply this labelled form by. All terms in
+                the form are multiplied.
 
         Returns:
             :class:`LabelledForm`: the product of all terms with the quantity.
         """
-        if type(other) in (float, int):
-            other = Constant(other)
-        # UFL can cancel constants to a Zero type which needs treating separately
-        elif type(other) is ufl.constantvalue.Zero:
-            other = Constant(0.0)
-        elif type(other) not in [Constant, ufl.algebra.Product]:
-            return NotImplemented
         return self.label_map(all_terms, lambda t: Term(other*t.form, t.labels))
 
     def __truediv__(self, other):
@@ -243,18 +238,13 @@ def __truediv__(self, other):
 
         Args:
             other (float, :class:`Constant` or :class:`ufl.algebra.Product`):
-                the quantity to divide this labelled form by. If it is a float
-                or int then it is converted to a :class:`Constant` before the
-                division. All terms in the form are divided.
+                the quantity to divide this labelled form by. All terms in the
+                form are divided.
 
         Returns:
             :class:`LabelledForm`: the quotient of all terms with the quantity.
         """
-        if type(other) in (float, int, Constant, ufl.algebra.Product):
-            other = Constant(1.0 / other)
-            return self * other
-        else:
-            return NotImplemented
+        return self * (Constant(1.0) / other)
 
     __rmul__ = __mul__
 
@@ -287,10 +277,11 @@ def label_map(self, term_filter, map_if_true=identity,
                              filter(lambda t: t is not None,
                                     (map_if_true(t) if term_filter(t) else
                                      map_if_false(t) for t in self.terms)),
-                             # TODO: Not clear what the initialiser should be!
-                             # No initialiser means label_map can't work if everything is false
-                             # None is a problem as cannot be added to Term
-                             # NullTerm works but will need dropping ...
+                             # Need to set an initialiser, otherwise the label_map
+                             # won't work if the term_filter is False for everything
+                             # None does not work, as then we add Terms to None
+                             # and the addition operation is defined from None
+                             # rather than the Term. NullTerm solves this.
                              NullTerm))
 
         # Drop the NullTerm
@@ -429,3 +420,11 @@ def update_value(self, target, new):
                 return target
         else:
             raise ValueError("Unable to relabel %s" % target)
+
+
+# ---------------------------------------------------------------------------- #
+# Some common labels
+# ---------------------------------------------------------------------------- #
+
+subject = Label("subject", validator=lambda value: type(value) == Function)
+name = Label("name", validator=lambda value: type(value) == str)