ffc team mailing list archive

[Johan Hake <hake@xxxxxxxxx>]

Hello!

The handed patch allows ffc.jit to compile ufl forms.

Johan

# HG changeset patch
# User "Johan Hake <hake@xxxxxxxxx>"
# Date 1237896663 -3600
# Node ID 5cfa855491ca6ce43b1846eb9a6da861b8505600
# Parent  bf4a10355f828531172d186e4eb8df84e884c1c0
Update to jit to be able to jit a ufl form/element

diff -r bf4a10355f82 -r 5cfa855491ca ffc/jit/jit.py
--- a/ffc/jit/jit.py	Tue Mar 24 00:10:29 2009 +0100
+++ b/ffc/jit/jit.py	Tue Mar 24 13:11:03 2009 +0100
@@ -29,7 +29,10 @@
     from ffc.compiler.uflcompiler import compile as uflcompile
     from ufl.classes import Form as UFLForm
     from ufl.classes import FiniteElementBase
+    from ufl.classes import TestFunction as UFLTestFunction
+    from ufl.objects import dx as UFLdx
 except:
+    print "Error while importing UFL"
     pass
 
 from ffc.compiler.language.algebra import Form, TestFunction
@@ -54,17 +57,17 @@
       options : An option dictionary
     """
 
+    # FIXME: Remove this test later
+    use_ufl = not options is None and "compiler" in options and options["compiler"] == "ufl"
     # Check if we get an element or a form
-    if isinstance(object, FiniteElement) or isinstance(object, MixedElement):
-        return jit_element(object, options)
+    if isinstance(object, FiniteElement) or isinstance(object, MixedElement) or (use_ufl and isinstance(object, FiniteElementBase)):
+        return jit_element(object, options, use_ufl)
     else:
-        return jit_form(object, options)
+        return jit_form(object, options, use_ufl)
 
-def jit_form(form, options=None):
+def jit_form(form, options=None, use_ufl=False):
     "Just-in-time compile the given form"
 
-    # FIXME: Remove this test later
-    use_ufl = not options is None and "compiler" in options and options["compiler"] == "ufl"
     if use_ufl:
         if not isinstance(form, UFLForm):
             form = UFLForm(form)
@@ -126,21 +129,28 @@
     else:
         return compiled_form, module, form.form_data
 
-def jit_element(element, options=None):
+def jit_element(element, options=None, use_ufl=False):
     "Just-in-time compile the given element"
     
     # Check that we get an element
-    if not (isinstance(element, FiniteElement) or isinstance(element, MixedElement)):
+    if not (isinstance(element, FiniteElement) or isinstance(element, MixedElement) or (use_ufl and isinstance(element, FiniteElementBase))):
         raise RuntimeError, "Expecting a finite element."
 
     # Create simplest possible dummy form
-    if element.value_dimension(0) > 1:
-        form = TestFunction(element)[0]*dx
+    if use_ufl:
+        # FIXME: Check that his is correct...
+        if len(element.value_shape()) > 0:
+            form = UFLTestFunction(element)[0]*UFLdx 
+        else:
+            form = UFLTestFunction(element)*UFLdx
     else:
-        form = TestFunction(element)*dx
+        if element.value_dimension(0) > 1:
+            form = TestFunction(element)[0]*dx
+        else:
+            form = TestFunction(element)*dx
 
     # Compile form
-    (compiled_form, module, form_data) = jit_form(form, options)
+    (compiled_form, module, form_data) = jit_form(form, options, use_ufl)
 
     return extract_element_and_dofmap(module)