added a loopy kernel and limiter for the mean mixing ratio

gusto/core/kernels.py

@@ -113,6 +113,54 @@ def apply(self, field, field_in):
                  {"field": (field, WRITE),
                   "field_in": (field_in, READ)})
 
+class MeanMixingRatioWeights():
+    """
+    Finds the lambda values for blending a mixing ratio and its
+    mean DG0 field in the MeanMixingRatioLimiter.
+
+    First, each cell is looped over and the minimum value is computed
+    """
+
+    def __init__(self, V):
+        """
+        Args:
+            V (:class:`FunctionSpace`): The space of the field to be clipped.
+        """
+        # Using DG1-equispaced, with 4 DOFs per cell
+        shapes = {'nDOFs': V.finat_element.space_dimension(),
+                  'nDOFs_base': int(V.finat_element.space_dimension() / 4)}
+        domain = "{{[i]: 0 <= i < {nDOFs_base}}}".format(**shapes)
+
+        instrs = ("""
+                  <float64> min_value = 0.0
+                  for i
+                      min_value = fmin(mX_field[i*4], mX_field[i*4+1])
+                      min_value = fmin(min_value, mX_field[i*4+2])
+                      min_value = fmin(min_value, mX_field[i*4+3])
+                      if min_value < 0.0
+                          lamda[i] = -min_value/(mean_field[i] - min_value)
+                      else
+                          lamda[i] = 0.0
+                      end
+                  end
+                  """)
+
+        self._kernel = (domain, instrs)
+
+    def apply(self, lamda, mX_field, mean_field):
+        """
+        Performs the par loop.
+
+        Args:
+            w (:class:`Function`): the field in which to store the weights. This
+                lives in the continuous target space.
+        """
+        par_loop(self._kernel, dx,
+                 {"lamda": (lamda, WRITE),
+                  "mX_field": (mX_field, READ),
+                  "mean_field": (mean_field, READ)})
+
+
 
 class MinKernel():
     """Finds the minimum DoF value of a field."""

gusto/spatial_methods/augmentation.py

@@ -20,9 +20,9 @@
     time_derivative, transport, transporting_velocity, TransportEquationType,
     logger, prognostic, mass_weighted
 )
-from gusto.spatial_methods.transport_methods import DGUpwind
+from gusto.spatial_methods.limiters import MeanLimiter, MixedFSLimiter
 import copy
-
+import numpy as np
 
 class Augmentation(object, metaclass=ABCMeta):
     """
@@ -55,6 +55,13 @@ def update(self, x_in_mixed):
 
         pass
 
+    def limit(self, x_in_mixed):
+        """
+        Apply any special limiting as part of the augmentation
+        """
+
+        pass
+
 
 class VorticityTransport(Augmentation):
     """
@@ -255,52 +262,64 @@ class MeanMixingRatio(Augmentation):
     Args:
         domain (:class:`Domain`): The domain object.
         eqns (:class:`PrognosticEquationSet`): The overarching equation set.
-        mixing_ratio (:class: list): A list of mixing ratios that 
-        are to have augmented mean mixing ratio fields.
-        OR, keep as a single mixing ratio, but define 
-        multiple augmentations?
-        START with just a single.
+        mX_names (:class: list): A list of mixing ratios that 
+        require augmented mean mixing ratio fields.
     """
 
     def __init__(
-            self, domain, eqns, mX_name
+            self, domain, eqns, mX_names
     ):
 
         self.name = 'mean_mixing_ratio'
-
-        self.mX_name = mX_name
-        self.mean_name = 'mean_'+mX_name
+        self.mX_names = mX_names
+        self.mX_num = len(mX_names)
 
+        # Store information about original equation set
         self.eqn_orig = eqns
-        self.domain=domain
+        self.domain = domain
         exist_spaces = eqns.spaces
         self.idx_orig = len(exist_spaces)
 
-        # Change this to adjust for a list
-        mean_idx = self.idx_orig
-        self.mean_idx = mean_idx
-
-        # Extract the mixing ratio in question:
-        mX_idx = eqns.field_names.index(mX_name)
-        self.mX_idx = mX_idx
-
         # Define the mean mixing ratio on the DG0 space
         DG0 = FunctionSpace(domain.mesh, "DG", 0)
-        mX_space = eqns.spaces[mX_idx]
-
-        # But, what if the mixing ratios are in 
-        # different function spaces???
-
-        self.mX_in = Function(mX_space)
-        self.mean = Function(DG0)
-        self.compute_mean = Projector(self.mX_in, self.mean)
-
-         # Set up the scheme for the mean mixing ratio
-
-        mean_mX = Function(DG0, name=self.mean_name)
-        mean_space = DG0
-        exist_spaces.append(mean_space)
 
+        # Set up fields and names for each mixing ratio
+        self.mean_names = []
+        self.mean_idxs = []
+        self.mX_idxs = []
+        mX_spaces = []
+        mean_spaces = []
+        self.limiters = []
+        #self.sublimiters = {}
+
+        for i in range(self.mX_num):
+            mX_name = mX_names[i]
+            print(mX_names)
+            print(mX_name)
+            self.mean_names.append('mean_'+mX_name)
+            mean_spaces.append(DG0)
+            exist_spaces.append(DG0)
+            self.mean_idxs.append(self.idx_orig + i)
+
+            # Extract the mixing ratio in question:
+            mX_idx = eqns.field_names.index(mX_name)
+            mX_spaces.append(eqns.spaces[mX_idx])
+            self.mX_idxs.append(mX_idx)
+
+            # Define a limiter
+            self.limiters.append(MeanLimiter(eqns.spaces[mX_idx]))
+            #self.sublimiters.update({mX_name: MeanLimiter(eqns.spaces[mX_idx])})
+
+        #self.limiter = MixedFSLimiter(self.eqn_orig, self.sublimiters)
+        #self.limiter = MixedFSLimiter(sublimiters)
+
+        # Contruct projectors for computing the mean mixing ratios
+        self.mX_ins = [Function(mX_spaces[i]) for i in range(self.mX_num)]
+        self.mean_outs = [Function(mean_spaces[i]) for i in range(self.mX_num)]
+        self.compute_means = [Projector(self.mX_ins[i], self.mean_outs[i]) \
+                               for i in range(self.mX_num)]
+
+        # Create the new mixed function space:
         self.fs = MixedFunctionSpace(exist_spaces)
         self.X = Function(self.fs)
         self.tests = TestFunctions(self.fs)
@@ -315,14 +334,24 @@ def setup_residual(self, spatial_methods, equation):
         # Copy spatial method for the mixing ratio onto the 
         # mean mixing ratio.
 
-        old_residual = equation.residual
+        orig_residual = equation.residual
 
         # Copy the mean mixing ratio residual terms:
-        mean_residual = old_residual.label_map(
-            lambda t: t.get(prognostic) == self.mX_name,
-            map_if_false=drop
-        )
-        mean_residual = prognostic.update_value(mean_residual, self.mean_name)
+        for i in range(self.mX_num):
+            if i == 0:
+                mean_residual = orig_residual.label_map(
+                    lambda t: t.get(prognostic) == self.mX_names[i],
+                    map_if_false=drop
+                )
+                mean_residual = prognostic.update_value(mean_residual, self.mean_names[i])
+            else:
+                mean_residual_term = orig_residual.label_map(
+                    lambda t: t.get(prognostic) == self.mX_names[i],
+                    map_if_false=drop
+                )
+                mean_residual_term = prognostic.update_value(mean_residual_term,\
+                                                              self.mean_names[i])
+                mean_residual = mean_residual + mean_residual_term
 
         print('\n in setup_transport')
 
@@ -334,19 +363,19 @@ def setup_residual(self, spatial_methods, equation):
             #print('\n', idx)
             #print(field)
 
-            prog = split(self.X)[idx]
+            #prog = split(self.X)[idx]
 
             #print('\n residual term before change')
             #print(old_residual.label_map(
             #    lambda t: t.get(prognostic) == field,
             #    map_if_false=drop
             #).form)
 
-            old_residual = old_residual.label_map(
+            orig_residual = orig_residual.label_map(
                 lambda t: t.get(prognostic) == field,
                 map_if_true=replace_subject(self.X, old_idx=idx, new_idx = idx)
             )
-            old_residual = old_residual.label_map(
+            orig_residual = orig_residual.label_map(
                 lambda t: t.get(prognostic) == field,
                 map_if_true=replace_test_function(self.tests, old_idx=idx, new_idx=idx)
             )
@@ -361,19 +390,25 @@ def setup_residual(self, spatial_methods, equation):
 
         #print('\n mean mX residual after change')
         #print(mean_residual.form)
-        mean_residual = mean_residual.label_map(
-            all_terms,
-            replace_subject(self.X, old_idx=self.mX_idx, new_idx=self.mean_idx)
-        )
-        mean_residual = mean_residual.label_map(
-            all_terms,
-            replace_test_function(self.tests, old_idx=self.mX_idx, new_idx=self.mean_idx)
-        )
+
+        # Update the subject and test functions for the
+        # mean mixing ratios
+        for i in range(self.mX_num):
+            mean_residual = mean_residual.label_map(
+                all_terms,
+                replace_subject(self.X, old_idx=self.mX_idxs[i], new_idx=self.mean_idxs[i])
+            )
+            mean_residual = mean_residual.label_map(
+                all_terms,
+                replace_test_function(self.tests, old_idx=self.mX_idxs[i], new_idx=self.mean_idxs[i])
+            )
+
+
         #print('\n mean mX residual after change')
         #print(mean_residual.form)
 
         # Form the new residual
-        residual = old_residual + mean_residual
+        residual = orig_residual + mean_residual
         self.residual = subject(residual, self.X)
 
         #Check these two forms
@@ -383,6 +418,7 @@ def setup_residual(self, spatial_methods, equation):
 
 
 
+
     def setup_transport_old(self, spatial_methods):
         mX_spatial_method = next(method for method in spatial_methods if method.variable == self.mX_name)
 
@@ -427,6 +463,7 @@ def pre_apply(self, x_in):
 
     def post_apply(self, x_out):
         """
+        Apply the limiters.
         Sets the output fields, i.e. not the mean mixing ratios
 
         Args:
@@ -438,18 +475,28 @@ def post_apply(self, x_out):
 
     def update(self, x_in_mixed):
         """
-        ...
+        Compute the mean mixing ratio field by projecting the mixing 
+        ratio from DG1 into DG0.
+
+        SHouldn't this be a conservative projection??!!
+        This requires density fields...
 
         Args:
             x_in_mixed (:class:`Function`): The mixed function to update.
         """
-
-        # Compute the mean mixing ratio
-        # How do I do this?
-        self.mX_in.assign(x_in_mixed.subfunctions[self.mX_idx])
-
-        # Project this into the lowest order space:
-        self.compute_mean.project()
-
-        self.x_in.subfunctions[self.mean_idx].assign(self.mean)
-        pass
+        print('in update')
+        for i in range(self.mX_num):
+            self.mX_ins[i].assign(x_in_mixed.subfunctions[self.mX_idxs[i]])
+            self.compute_means[i].project()
+            self.x_in.subfunctions[self.mean_idxs[i]].assign(self.mean_outs[i])
+
+    def limit(self, x_in_mixed):
+        # Ensure non-negativity by applying the blended limiter
+        for i in range(self.mX_num):
+            print('limiting within the augmentation')
+            mX_field = x_in_mixed.subfunctions[self.mX_idxs[i]]
+            mean_field = x_in_mixed.subfunctions[self.mean_idxs[i]]
+            print(np.min(x_in_mixed.subfunctions[self.mX_idxs[i]].dat.data))
+            self.limiters[i].apply(mX_field, mean_field)
+            print(np.min(x_in_mixed.subfunctions[self.mX_idxs[i]].dat.data))
+            #x_in_mixed.subfunctions[self.mX_idxs[i]].assign(mX_field)