Increment form for implicit RK added and tested

gusto/time_discretisation/explicit_runge_kutta.py

@@ -123,13 +123,9 @@ def __init__(self, domain, butcher_matrix, field_name=None,
                          limiter=limiter, options=options,
                          augmentation=augmentation)
         self.butcher_matrix = butcher_matrix
-        self.nbutcher = int(np.shape(self.butcher_matrix)[0])
+        self.nStages = int(np.shape(self.butcher_matrix)[0])
         self.rk_formulation = rk_formulation
 
-    @property
-    def nStages(self):
-        return self.nbutcher
-
     def setup(self, equation, apply_bcs=True, *active_labels):
         """
         Set up the time discretisation based on the equation.
@@ -163,7 +159,7 @@ def solver(self):
             for stage in range(self.nStages):
                 # setup linear solver using lhs and rhs defined in derived class
                 problem = NonlinearVariationalProblem(
-                    self.lhs[stage].form - self.rhs[stage].form,
+                    self.res[stage].form,
                     self.field_i[stage+1], bcs=self.bcs
                 )
                 solver_name = self.field_name+self.__class__.__name__+str(stage)
@@ -176,7 +172,7 @@ def solver(self):
 
         elif self.rk_formulation == RungeKuttaFormulation.linear:
             problem = NonlinearVariationalProblem(
-                self.lhs - self.rhs[0], self.x1, bcs=self.bcs
+                self.res[0], self.x1, bcs=self.bcs
             )
             solver_name = self.field_name+self.__class__.__name__
             solver = NonlinearVariationalSolver(
@@ -186,7 +182,7 @@ def solver(self):
 
             # Set up problem for final step
             problem_last = NonlinearVariationalProblem(
-                self.lhs - self.rhs[1], self.x1, bcs=self.bcs
+                self.res[1], self.x1, bcs=self.bcs
             )
             solver_name = self.field_name+self.__class__.__name__+'_last'
             solver_last = NonlinearVariationalSolver(
@@ -202,54 +198,21 @@ def solver(self):
             )
 
     @cached_property
-    def lhs(self):
-        """Set up the discretisation's left hand side (the time derivative)."""
+    def res(self):
+        """Set up the discretisation's residual."""
 
         if self.rk_formulation == RungeKuttaFormulation.increment:
-            l = self.residual.label_map(
+            residual = self.residual.label_map(
                 lambda t: t.has_label(time_derivative),
                 map_if_true=replace_subject(self.x_out, old_idx=self.idx),
                 map_if_false=drop)
-
-            return l.form
-
-        elif self.rk_formulation == RungeKuttaFormulation.predictor:
-            lhs_list = []
-            for stage in range(self.nStages):
-                l = self.residual.label_map(
-                    lambda t: t.has_label(time_derivative),
-                    map_if_true=replace_subject(self.field_i[stage+1], old_idx=self.idx),
-                    map_if_false=drop)
-                lhs_list.append(l)
-
-            return lhs_list
-
-        if self.rk_formulation == RungeKuttaFormulation.linear:
-            l = self.residual.label_map(
-                lambda t: t.has_label(time_derivative),
-                map_if_true=replace_subject(self.x1, old_idx=self.idx),
-                map_if_false=drop)
-
-            return l.form
-
-        else:
-            raise NotImplementedError(
-                'Runge-Kutta formulation is not implemented'
-            )
-
-    @cached_property
-    def rhs(self):
-        """Set up the time discretisation's right hand side."""
-
-        if self.rk_formulation == RungeKuttaFormulation.increment:
             r = self.residual.label_map(
                 all_terms,
                 map_if_true=replace_subject(self.x1, old_idx=self.idx))
 
-            r = r.label_map(
+            residual += r.label_map(
                 lambda t: t.has_label(time_derivative),
-                map_if_true=drop,
-                map_if_false=lambda t: -1*t)
+                map_if_true=drop)
 
             # If there are no active labels, we may have no terms at this point
             # So that we can still do xnp1 = xn, put in a zero term here
@@ -261,19 +224,22 @@ def rhs(self):
                     # Drop label from this
                     map_if_true=lambda t: time_derivative.remove(t),
                     map_if_false=drop)
-                r += null_term
+                residual += null_term
 
-            return r.form
+            return residual.form
 
         elif self.rk_formulation == RungeKuttaFormulation.predictor:
-            rhs_list = []
-
+            residual_list = []
             for stage in range(self.nStages):
+                residual = self.residual.label_map(
+                    lambda t: t.has_label(time_derivative),
+                    map_if_true=replace_subject(self.field_i[stage+1], self.idx),
+                    map_if_false=drop)
                 r = self.residual.label_map(
                     all_terms,
                     map_if_true=replace_subject(self.field_i[0], old_idx=self.idx))
 
-                r = r.label_map(
+                residual -= r.label_map(
                     lambda t: t.has_label(time_derivative),
                     map_if_true=keep,
                     map_if_false=lambda t: -self.butcher_matrix[stage, 0]*self.dt*t)
@@ -285,14 +251,16 @@ def rhs(self):
                         map_if_false=replace_subject(self.field_i[i], old_idx=self.idx)
                     )
 
-                    r -= self.butcher_matrix[stage, i]*self.dt*r_i
-
-                rhs_list.append(r)
+                    residual += self.butcher_matrix[stage, i]*self.dt*r_i
+                residual_list.append(residual)
 
-            return rhs_list
-
-        elif self.rk_formulation == RungeKuttaFormulation.linear:
+            return residual_list
 
+        if self.rk_formulation == RungeKuttaFormulation.linear:
+            time_term = self.residual.label_map(
+                lambda t: t.has_label(time_derivative),
+                map_if_true=replace_subject(self.x1, self.idx),
+                map_if_false=drop)
             r = self.residual.label_map(
                 lambda t: t.has_label(time_derivative),
                 map_if_true=replace_subject(self.x0, old_idx=self.idx),
@@ -331,8 +299,9 @@ def rhs(self):
                 map_if_true=keep,
                 map_if_false=lambda t: -self.dt*t
             )
-
-            return r_all_but_last.form, r.form
+            res = time_term - r
+            res_all_but_last = time_term - r_all_but_last
+            return res_all_but_last.form, res.form
 
         else:
             raise NotImplementedError(

gusto/time_discretisation/imex_runge_kutta.py

@@ -61,7 +61,7 @@ class IMEXRungeKutta(TimeDiscretisation):
 
     def __init__(self, domain, butcher_imp, butcher_exp, field_name=None,
                  linear_solver_parameters=None, nonlinear_solver_parameters=None,
-                 limiter=None, options=None):
+                 limiter=None, options=None, augmentation=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -82,10 +82,13 @@ def __init__(self, domain, butcher_imp, butcher_exp, field_name=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
+            augmentation (:class:`Augmentation`): allows the equation solved in
+                this time discretisation to be augmented, for instances with
+                extra terms of another auxiliary variable. Defaults to None.
         """
         super().__init__(domain, field_name=field_name,
                          solver_parameters=nonlinear_solver_parameters,
-                         options=options)
+                         options=options, augmentation=augmentation)
         self.butcher_imp = butcher_imp
         self.butcher_exp = butcher_exp
         self.nStages = int(np.shape(self.butcher_imp)[1])
@@ -127,16 +130,6 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
         self.xs = [Function(self.fs) for i in range(self.nStages)]
 
-    @cached_property
-    def lhs(self):
-        """Set up the discretisation's left hand side (the time derivative)."""
-        return super(IMEXRungeKutta, self).lhs
-
-    @cached_property
-    def rhs(self):
-        """Set up the discretisation's right hand side (the time derivative)."""
-        return super(IMEXRungeKutta, self).rhs
-
     def res(self, stage):
         """Set up the discretisation's residual for a given stage."""
         # Add time derivative terms  y_s - y^n for stage s
@@ -226,7 +219,7 @@ def solvers(self):
     @cached_property
     def final_solver(self):
         """Set up a solver for the final solve to evaluate time level n+1."""
-        # setup solver using lhs and rhs defined in derived class
+        # setup solver using residual (res) defined in derived class
         problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs)
         solver_name = self.field_name+self.__class__.__name__
         return NonlinearVariationalSolver(problem, solver_parameters=self.linear_solver_parameters, options_prefix=solver_name)
@@ -269,7 +262,7 @@ class IMEX_Euler(IMEXRungeKutta):
     """
     def __init__(self, domain, field_name=None,
                  linear_solver_parameters=None, nonlinear_solver_parameters=None,
-                 limiter=None, options=None):
+                 limiter=None, options=None, augmentation=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -286,13 +279,16 @@ def __init__(self, domain, field_name=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
+            augmentation (:class:`Augmentation`): allows the equation solved in
+                this time discretisation to be augmented, for instances with
+                extra terms of another auxiliary variable. Defaults to None.
         """
         butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]])
         butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]])
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
                          linear_solver_parameters=linear_solver_parameters,
                          nonlinear_solver_parameters=nonlinear_solver_parameters,
-                         limiter=limiter, options=options)
+                         limiter=limiter, options=options, augmentation=augmentation)
 
 
 class IMEX_ARS3(IMEXRungeKutta):
@@ -313,7 +309,7 @@ class IMEX_ARS3(IMEXRungeKutta):
     """
     def __init__(self, domain, field_name=None,
                  linear_solver_parameters=None, nonlinear_solver_parameters=None,
-                 limiter=None, options=None):
+                 limiter=None, options=None, augmentation=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -330,6 +326,9 @@ def __init__(self, domain, field_name=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
+            augmentation (:class:`Augmentation`): allows the equation solved in
+                this time discretisation to be augmented, for instances with
+                extra terms of another auxiliary variable. Defaults to None.
         """
         g = (3. + np.sqrt(3.))/6.
         butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]])
@@ -338,7 +337,7 @@ def __init__(self, domain, field_name=None,
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
                          linear_solver_parameters=linear_solver_parameters,
                          nonlinear_solver_parameters=nonlinear_solver_parameters,
-                         limiter=limiter, options=options)
+                         limiter=limiter, options=options, augmentation=augmentation)
 
 
 class IMEX_ARK2(IMEXRungeKutta):
@@ -359,7 +358,7 @@ class IMEX_ARK2(IMEXRungeKutta):
     """
     def __init__(self, domain, field_name=None,
                  linear_solver_parameters=None, nonlinear_solver_parameters=None,
-                 limiter=None, options=None):
+                 limiter=None, options=None, augmentation=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -376,6 +375,9 @@ def __init__(self, domain, field_name=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
+            augmentation (:class:`Augmentation`): allows the equation solved in
+                this time discretisation to be augmented, for instances with
+                extra terms of another auxiliary variable. Defaults to None.
         """
         g = 1. - 1./np.sqrt(2.)
         d = 1./(2.*np.sqrt(2.))
@@ -385,7 +387,7 @@ def __init__(self, domain, field_name=None,
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
                          linear_solver_parameters=linear_solver_parameters,
                          nonlinear_solver_parameters=nonlinear_solver_parameters,
-                         limiter=limiter, options=options)
+                         limiter=limiter, options=options, augmentation=augmentation)
 
 
 class IMEX_SSP3(IMEXRungeKutta):
@@ -404,7 +406,7 @@ class IMEX_SSP3(IMEXRungeKutta):
     """
     def __init__(self, domain, field_name=None,
                  linear_solver_parameters=None, nonlinear_solver_parameters=None,
-                 limiter=None, options=None):
+                 limiter=None, options=None, augmentation=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -421,14 +423,17 @@ def __init__(self, domain, field_name=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
+            augmentation (:class:`Augmentation`): allows the equation solved in
+                this time discretisation to be augmented, for instances with
+                extra terms of another auxiliary variable. Defaults to None.
         """
         g = 1. - (1./np.sqrt(2.))
         butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.), (1./6.), (2./3.)]])
         butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.), (1./6.), (2./3.)]])
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
                          linear_solver_parameters=linear_solver_parameters,
                          nonlinear_solver_parameters=nonlinear_solver_parameters,
-                         limiter=limiter, options=options)
+                         limiter=limiter, options=options, augmentation=augmentation)
 
 
 class IMEX_Trap2(IMEXRungeKutta):
@@ -447,7 +452,7 @@ class IMEX_Trap2(IMEXRungeKutta):
     """
     def __init__(self, domain, field_name=None,
                  linear_solver_parameters=None, nonlinear_solver_parameters=None,
-                 limiter=None, options=None):
+                 limiter=None, options=None, augmentation=None):
         """
         Args:
             domain (:class:`Domain`): the model's domain object, containing the
@@ -464,11 +469,14 @@ def __init__(self, domain, field_name=None,
                 options to either be passed to the spatial discretisation, or
                 to control the "wrapper" methods, such as Embedded DG or a
                 recovery method. Defaults to None.
+            augmentation (:class:`Augmentation`): allows the equation solved in
+                this time discretisation to be augmented, for instances with
+                extra terms of another auxiliary variable. Defaults to None.
         """
         e = 0.
         butcher_imp = np.array([[0., 0., 0., 0.], [e, 0., 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]])
         butcher_exp = np.array([[0., 0., 0., 0.], [1., 0., 0., 0.], [0.5, 0.5, 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0.5, 0.]])
         super().__init__(domain, butcher_imp, butcher_exp, field_name,
                          linear_solver_parameters=linear_solver_parameters,
                          nonlinear_solver_parameters=nonlinear_solver_parameters,
-                         limiter=limiter, options=options)
+                         limiter=limiter, options=options, augmentation=augmentation)