Split up kernel generation

python/dune/perftool/pdelab/localoperator.py

@@ -3,6 +3,8 @@ from os.path import splitext
 
 import logging
 
+import numpy as np
+
 from dune.perftool.options import (get_form_option,
                                    get_option,
                                    option_switch)
@@ -707,81 +709,7 @@ def cgen_class_from_cache(tag, members=[]):
     return Class(basename, base_classes=base_classes, members=[constructor] + members + pm + decls, tparam_decls=tparams)
 
 
-def generate_localoperator_kernels(operator):
-    logger = logging.getLogger(__name__)
-
-    data = get_global_context_value("data")
-    original_form = data.object_by_name[get_form_option("form")]
-    from dune.perftool.ufl.preprocess import preprocess_form
-
-    if get_form_option("adjoint"):
-        # Generate adjoint operator
-        #
-        # The jacobian of the adjoint form is just the jacobian of the
-        # original form with test and ansazt function swapped. A a
-        # linear form you have to subtract the derivative of the
-        # objective function w.r.t the ansatz function to get the
-        # final residual formulation of the adjoint.
-        #
-        # TODO: This might only be true for linear problems. Adjust
-        # documentation as knowledge improves ;)
-        assert get_form_option("objective_function") is not None
-        assert get_form_option("control") is False
-
-        from ufl import derivative, adjoint, action, replace
-        from ufl.classes import Coefficient
-
-        # Jacobian of the adjoint form
-        jacform = derivative(original_form, original_form.coefficients()[0])
-        adjoint_jacform = adjoint(jacform)
-
-        # Derivative of objective function w.r.t. state
-        objective = data.object_by_name[get_form_option("objective_function")]
-        objective_jacobian = derivative(objective, objective.coefficients()[0])
-
-        # Replace coefficient belonging to ansatz function with new coefficient
-        element = objective.coefficients()[0].ufl_element()
-        coeff = Coefficient(element, count=3)
-        objective_jacobian = replace(objective_jacobian, {objective.coefficients()[0]: coeff})
-        if len(adjoint_jacform.coefficients()) > 0:
-            adjoint_jacform = replace(adjoint_jacform, {adjoint_jacform.coefficients()[0]: coeff})
-
-        # Residual of the adjoint form
-        adjoint_form = action(adjoint_jacform, original_form.coefficients()[0])
-        adjoint_form = adjoint_form + objective_jacobian
-
-        # Update form and original_form
-        original_form = adjoint_form
-        form = preprocess_form(adjoint_form).preprocessed_form
-
-    elif get_form_option("control"):
-        # Generate control operator
-        #
-        # This is the normal form derived w.r.t. the control
-        # variable. Right now this is olny implemented for scalar
-        # controls.
-        assert get_form_option("control_variable") is not None
-
-        control = data.object_by_name[get_form_option("control_variable")]
-        assert control.ufl_shape is ()
-
-        from ufl import diff, replace
-        from ufl.classes import Coefficient
-        control_form = diff(original_form, control)
-        # element = control_form.coefficients()[0].ufl_element()
-        # coeff = Coefficient(element, count=3)
-        # control_form = replace(control_form, {control_form.coefficients()[0]: coeff})
-        original_form = control_form
-        form = preprocess_form(control_form).preprocessed_form
-
-    else:
-        form = preprocess_form(original_form).preprocessed_form
-
-
-    # Reset the generation cache
-    from dune.perftool.generation import delete_cache_items
-    delete_cache_items()
-
+def local_operator_default_settings(operator, form):
     # Manage includes and base classes that we always need
     include_file('dune/pdelab/gridfunctionspace/gridfunctionspace.hh', filetag="operatorfile")
     include_file('dune/pdelab/localoperator/idefault.hh', filetag="operatorfile")
@@ -825,10 +753,9 @@ def generate_localoperator_kernels(operator):
         base_class('Dune::PDELab::InstationaryLocalOperatorDefaultMethods<{}>'
                    .format(rf), classtag="operator")
 
-    # Have a data structure collect the generated kernels
-    operator_kernels = {}
 
-    logger.info("generate_localoperator_kernels: create residual methods")
+def generate_residual_kernels(form):
+    logger = logging.getLogger(__name__)
     with global_context(form_type='residual'):
         # Generate the necessary residual methods
         for measure in set(i.integral_type() for i in form.integrals()):
@@ -883,63 +810,181 @@ def generate_localoperator_kernels(operator):
                                              classtag="operator",
                                              )
 
-                operator_kernels[(measure, 'residual')] = kernel
+                return {(measure, 'residual'): kernel}
 
-    # Generate the necessary jacobian methods
-    if not get_form_option("numerical_jacobian"):
-        logger.info("generate_localoperator_kernels: create jacobian methods")
-        from ufl import derivative
-        jacform = derivative(original_form, original_form.coefficients()[0])
 
-        from dune.perftool.ufl.preprocess import preprocess_form
-        jacform = preprocess_form(jacform).preprocessed_form
+def generate_jacobian_kernels(form, original_form):
+    logger = logging.getLogger(__name__)
 
-        with global_context(form_type="jacobian"):
-            if get_form_option("generate_jacobians"):
-                for measure in set(i.integral_type() for i in jacform.integrals()):
-                    logger.info("generate_localoperator_kernels: measure {}".format(measure))
-                    with global_context(integral_type=measure):
-                        with global_context(kernel=assembler_routine_name()):
-                            kernel = [k for k in get_backend(interface="generate_kernels_per_integral")(jacform.integrals_by_type(measure))]
-                    operator_kernels[(measure, 'jacobian')] = kernel
+    from ufl import derivative
+    jacform = derivative(original_form, original_form.coefficients()[0])
+
+    from dune.perftool.ufl.preprocess import preprocess_form
+    jacform = preprocess_form(jacform).preprocessed_form
+
+    operator_kernels = {}
+    with global_context(form_type="jacobian"):
+        if get_form_option("generate_jacobians"):
+            for measure in set(i.integral_type() for i in jacform.integrals()):
+                logger.info("generate_localoperator_kernels: measure {}".format(measure))
+                with global_context(integral_type=measure):
+                    from dune.perftool.pdelab.signatures import assembler_routine_name
+                    with global_context(kernel=assembler_routine_name()):
+                        kernel = [k for k in get_backend(interface="generate_kernels_per_integral")(jacform.integrals_by_type(measure))]
+                operator_kernels[(measure, 'jacobian')] = kernel
+
+            # Generate dummy functions for those kernels, that vanished in the differentiation process
+            # We *could* solve this problem by using lambda_* terms but we do not really want that, so
+            # we use empty jacobian assembly methods instead
+            alpha_measures = set(i.integral_type() for i in form.integrals())
+            jacobian_measures = set(i.integral_type() for i in jacform.integrals())
+            for it in alpha_measures - jacobian_measures:
+                with global_context(integral_type=it):
+                    from dune.perftool.pdelab.signatures import assembly_routine_signature
+                    operator_kernels[(it, 'jacobian')] = [LoopyKernelMethod(assembly_routine_signature(), kernel=None)]
+
+    # Jacobian apply methods for matrix-free computations
+    if get_form_option("matrix_free"):
+        # The apply vector has reserved index 1 so we directly use Coefficient class from ufl
+        from ufl import Coefficient
+        apply_coefficient = Coefficient(form.coefficients()[0].ufl_element(), 1)
+
+        # Create application of jacobian on vector
+        from ufl import action
+        jac_apply_form = action(jacform, apply_coefficient)
+
+        # Create kernel for jacobian application
+        with global_context(form_type="jacobian_apply"):
+            for measure in set(i.integral_type() for i in jac_apply_form.integrals()):
+                with global_context(integral_type=measure):
+                    with global_context(kernel=assembler_routine_name()):
+                        kernel = [k for k in get_backend(interface="generate_kernels_per_integral")(jac_apply_form.integrals_by_type(measure))]
+                operator_kernels[(measure, 'jacobian_apply')] = kernel
 
                 # Generate dummy functions for those kernels, that vanished in the differentiation process
                 # We *could* solve this problem by using lambda_* terms but we do not really want that, so
                 # we use empty jacobian assembly methods instead
                 alpha_measures = set(i.integral_type() for i in form.integrals())
-                jacobian_measures = set(i.integral_type() for i in jacform.integrals())
-                for it in alpha_measures - jacobian_measures:
+                jacobian_apply_measures = set(i.integral_type() for i in jac_apply_form.integrals())
+                for it in alpha_measures - jacobian_apply_measures:
                     with global_context(integral_type=it):
                         from dune.perftool.pdelab.signatures import assembly_routine_signature
-                        operator_kernels[(it, 'jacobian')] = [LoopyKernelMethod(assembly_routine_signature(), kernel=None)]
-
-        # Jacobian apply methods for matrix-free computations
-        if get_form_option("matrix_free"):
-            # The apply vector has reserved index 1 so we directly use Coefficient class from ufl
-            from ufl import Coefficient
-            apply_coefficient = Coefficient(form.coefficients()[0].ufl_element(), 1)
-
-            # Create application of jacobian on vector
-            from ufl import action
-            jac_apply_form = action(jacform, apply_coefficient)
-
-            # Create kernel for jacobian application
-            with global_context(form_type="jacobian_apply"):
-                for measure in set(i.integral_type() for i in jac_apply_form.integrals()):
-                    with global_context(integral_type=measure):
-                        with global_context(kernel=assembler_routine_name()):
-                            kernel = [k for k in get_backend(interface="generate_kernels_per_integral")(jac_apply_form.integrals_by_type(measure))]
-                    operator_kernels[(measure, 'jacobian_apply')] = kernel
-
-                    # Generate dummy functions for those kernels, that vanished in the differentiation process
-                    # We *could* solve this problem by using lambda_* terms but we do not really want that, so
-                    # we use empty jacobian assembly methods instead
-                    alpha_measures = set(i.integral_type() for i in form.integrals())
-                    jacobian_apply_measures = set(i.integral_type() for i in jac_apply_form.integrals())
-                    for it in alpha_measures - jacobian_apply_measures:
-                        with global_context(integral_type=it):
-                            from dune.perftool.pdelab.signatures import assembly_routine_signature
-                            operator_kernels[(it, 'jacobian_apply')] = [LoopyKernelMethod(assembly_routine_signature(), kernel=None)]
+                        operator_kernels[(it, 'jacobian_apply')] = [LoopyKernelMethod(assembly_routine_signature(), kernel=None)]
+
+    return operator_kernels
+
+
+def generate_localoperator_kernels(operator):
+    logger = logging.getLogger(__name__)
+
+    data = get_global_context_value("data")
+    original_form = data.object_by_name[get_form_option("form")]
+    from dune.perftool.ufl.preprocess import preprocess_form
+
+    if get_form_option("adjoint"):
+        # Generate adjoint operator
+        #
+        # The jacobian of the adjoint form is just the jacobian of the
+        # original form with test and ansazt function swapped. A a
+        # linear form you have to subtract the derivative of the
+        # objective function w.r.t the ansatz function to get the
+        # final residual formulation of the adjoint.
+        #
+        # TODO: This might only be true for linear problems. Adjust
+        # documentation as knowledge improves ;)
+        assert get_form_option("objective_function") is not None
+        assert get_form_option("control") is False
+
+        from ufl import derivative, adjoint, action, replace
+        from ufl.classes import Coefficient
+
+        # Jacobian of the adjoint form
+        jacform = derivative(original_form, original_form.coefficients()[0])
+        adjoint_jacform = adjoint(jacform)
+
+        # Derivative of objective function w.r.t. state
+        objective = data.object_by_name[get_form_option("objective_function")]
+        objective_jacobian = derivative(objective, objective.coefficients()[0])
+
+        # Replace coefficient belonging to ansatz function with new coefficient
+        element = objective.coefficients()[0].ufl_element()
+        coeff = Coefficient(element, count=3)
+        objective_jacobian = replace(objective_jacobian, {objective.coefficients()[0]: coeff})
+        if len(adjoint_jacform.coefficients()) > 0:
+            adjoint_jacform = replace(adjoint_jacform, {adjoint_jacform.coefficients()[0]: coeff})
+
+        # Residual of the adjoint form
+        adjoint_form = action(adjoint_jacform, original_form.coefficients()[0])
+        adjoint_form = adjoint_form + objective_jacobian
+
+        # Update form and original_form
+        original_form = adjoint_form
+        form = preprocess_form(adjoint_form).preprocessed_form
+
+    elif get_form_option("control"):
+        # Generate control operator
+        #
+        # This is the normal form derived w.r.t. the control
+        # variable.
+        assert get_form_option("control_variable") is not None
+        controls = [data.object_by_name[ctrl.strip()] for ctrl in get_form_option("control_variable").split(",")]
+        assert len(controls) == 1
+
+        # We need to transform numpy ints to python native ints
+        def _unravel(flat_index, shape):
+            multi_index = np.unravel_index(flat_index, shape)
+            multi_index = tuple(int(i) for i in multi_index)
+            return multi_index
+
+        forms = []
+        for control in controls:
+            shape = control.ufl_shape
+            flat_length = np.prod(shape)
+            for i in range(flat_length):
+                c = control[_unravel(i, shape)]
+                control_form = diff(original_form, control)
+                forms.append(preprocess_form(control_form).preprocessed_form)
+
+        # Used to create local operator default settings
+        form = preprocess_form(original_form).preprocessed_form
+
+        # control = data.object_by_name[get_form_option("control_variable")]
+        # assert control.ufl_shape is ()
+
+        # from ufl import diff, replace
+        # from ufl.classes import Coefficient
+        # control_form = diff(original_form, control)
+        # # element = control_form.coefficients()[0].ufl_element()
+        # # coeff = Coefficient(element, count=3)
+        # # control_form = replace(control_form, {control_form.coefficients()[0]: coeff})
+        # original_form = control_form
+        # form = preprocess_form(control_form).preprocessed_form
+
+    else:
+        form = preprocess_form(original_form).preprocessed_form
+
+
+    # Reset the generation cache
+    from dune.perftool.generation import delete_cache_items
+    delete_cache_items()
+
+    # Have a data structure collect the generated kernels
+    operator_kernels = {}
+
+    # Generate things needed for all local operator files
+    local_operator_default_settings(operator, form)
+
+    if get_form_option("control"):
+        logger.info("generate_localoperator_kernels: create methods for control operator")
+        operator_kernels.update(generate_control_kernels(forms))
+    else:
+        logger.info("generate_localoperator_kernels: create residual methods")
+        operator_kernels.update(generate_residual_kernels(form))
+
+        # Generate the necessary jacobian methods
+        if not get_form_option("numerical_jacobian"):
+            logger.info("generate_localoperator_kernels: create jacobian methods")
+            operator_kernels.update(generate_jacobian_kernels(form, original_form))
 
     # Return the set of generated kernels
     return operator_kernels