Defining a fem.Function using a conditional ufl statement

Hi Dokken,

Thank you for the quick reply!

I still get an error when using fem.Expression, although it is a different one.

Are using Bessel functions allowed for Expressions?

Many thanks,
Katie

import numpy as np
from dolfinx import mesh, fem
import ufl
from mpi4py import MPI


domain = mesh.create_unit_interval(MPI.COMM_WORLD, 100)
V = fem.FunctionSpace(domain, ("Lagrange", 1))

def analytic_rz(x):   

    # Simplified parameters for the problem
    
    r1 = 0.1
    r2 = 0.2
    r3 = 0.3

    mu0 = 4 * np.pi * 1e-7

    C = 2e8
    I = 100

    gamma = np.sqrt(1j * C)

    alpha2 = I # constants that are already known
    beta1 = 0
    alpha_ext = 0

    alpha1 = (
        mu0
        * r2
        * alpha2
        / (ufl.bessel_I(1, r1) + r1 * gamma * 0.5 * ( ufl.bessel_I(0, r1) + ufl.bessel_I(1, r1)))
    )
    beta2 = r1 * (alpha1 * ufl.bessel_I(1, r1) - 0.5 * mu0 * r1 * alpha2)
    beta_ext = r2 * (beta2 / r2 - 0.5 * mu0 * r2 * alpha2)

    A = ufl.conditional(ufl.And(0 <= x[0], x[0] <= r1), alpha1 * ufl.bessel_I(1, x[0] / gamma ), \
                        ufl.conditional(ufl.And(r1 <= x[0], x[0] <= r2), 0.5 * mu0 * x[0] * alpha2 + beta2 / x[0], beta_ext / x[0]))
    return A


x = ufl.SpatialCoordinate(domain)
analytic_A = fem.Function(V)
analytic_expr = fem.Expression(analytic_rz(x), V.element.interpolation_points())
analytic_A.interpolate(analytic_expr)
RuntimeError                              Traceback (most recent call last)
Cell In[2], line 45
     43 x = ufl.SpatialCoordinate(domain)
     44 analytic_A = fem.Function(V)
---> 45 analytic_expr = fem.Expression(analytic_rz(x), V.element.interpolation_points())
     46 analytic_A.interpolate(analytic_expr)

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/dolfinx/fem/function.py:149, in Expression.__init__(self, e, X, comm, form_compiler_options, jit_options, dtype)
    146 else:
    147     raise RuntimeError(f"Unsupported scalar type {dtype} for Expression.")
--> 149 self._ufcx_expression, module, self._code = jit.ffcx_jit(comm, (e, _X),
    150                                                          form_compiler_options=form_compiler_options,
    151                                                          jit_options=jit_options)
    152 self._ufl_expression = e
    154 # Prepare coefficients data. For every coefficient in expression
    155 # take its C++ object.

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/dolfinx/jit.py:56, in mpi_jit_decorator.<locals>.mpi_jit(comm, *args, **kwargs)
     51 @functools.wraps(local_jit)
     52 def mpi_jit(comm, *args, **kwargs):
     53 
     54     # Just call JIT compiler when running in serial
     55     if comm.size == 1:
---> 56         return local_jit(*args, **kwargs)
     58     # Default status (0 == ok, 1 == fail)
     59     status = 0

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/dolfinx/jit.py:210, in ffcx_jit(ufl_object, form_compiler_options, jit_options)
    208     r = ffcx.codegeneration.jit.compile_coordinate_maps([ufl_object], options=p_ffcx, **p_jit)
    209 elif isinstance(ufl_object, tuple) and isinstance(ufl_object[0], ufl.core.expr.Expr):
--> 210     r = ffcx.codegeneration.jit.compile_expressions([ufl_object], options=p_ffcx, **p_jit)
    211 else:
    212     raise TypeError(type(ufl_object))

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/jit.py:247, in compile_expressions(expressions, options, cache_dir, timeout, cffi_extra_compile_args, cffi_verbose, cffi_debug, cffi_libraries)
    245     except Exception:
    246         pass
--> 247     raise e
    249 obj, module = _load_objects(cache_dir, module_name, expr_names)
    250 return obj, module, (decl, impl)

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/jit.py:238, in compile_expressions(expressions, options, cache_dir, timeout, cffi_extra_compile_args, cffi_verbose, cffi_debug, cffi_libraries)
    235     for name in expr_names:
    236         decl += expression_template.format(name=name)
--> 238     impl = _compile_objects(decl, expressions, expr_names, module_name, p, cache_dir,
    239                             cffi_extra_compile_args, cffi_verbose, cffi_debug, cffi_libraries)
    240 except Exception as e:
    241     try:
    242         # remove c file so that it will not timeout next time

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/jit.py:260, in _compile_objects(decl, ufl_objects, object_names, module_name, options, cache_dir, cffi_extra_compile_args, cffi_verbose, cffi_debug, cffi_libraries)
    256 import ffcx.compiler
    258 # JIT uses module_name as prefix, which is needed to make names of all struct/function
    259 # unique across modules
--> 260 _, code_body = ffcx.compiler.compile_ufl_objects(ufl_objects, prefix=module_name, options=options)
    262 ffibuilder = cffi.FFI()
    263 ffibuilder.set_source(module_name, code_body, include_dirs=[ffcx.codegeneration.get_include_path()],
    264                       extra_compile_args=cffi_extra_compile_args, libraries=cffi_libraries)

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/compiler.py:107, in compile_ufl_objects(ufl_objects, object_names, prefix, options, visualise)
    105 # Stage 3: code generation
    106 cpu_time = time()
--> 107 code = generate_code(ir, options)
    108 _print_timing(3, time() - cpu_time)
    110 # Stage 4: format code

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/codegeneration.py:56, in generate_code(ir, options)
     54 code_integrals = [integral_generator(integral_ir, options) for integral_ir in ir.integrals]
     55 code_forms = [form_generator(form_ir, options) for form_ir in ir.forms]
---> 56 code_expressions = [expression_generator(expression_ir, options) for expression_ir in ir.expressions]
     57 code_file_pre, code_file_post = file_generator(options)
     58 return CodeBlocks(file_pre=[code_file_pre], elements=code_finite_elements, dofmaps=code_dofmaps,
     59                   integrals=code_integrals, forms=code_forms, expressions=code_expressions,
     60                   file_post=[code_file_post])

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/codegeneration.py:56, in <listcomp>(.0)
     54 code_integrals = [integral_generator(integral_ir, options) for integral_ir in ir.integrals]
     55 code_forms = [form_generator(form_ir, options) for form_ir in ir.forms]
---> 56 code_expressions = [expression_generator(expression_ir, options) for expression_ir in ir.expressions]
     57 code_file_pre, code_file_post = file_generator(options)
     58 return CodeBlocks(file_pre=[code_file_pre], elements=code_finite_elements, dofmaps=code_dofmaps,
     59                   integrals=code_integrals, forms=code_forms, expressions=code_expressions,
     60                   file_post=[code_file_post])

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/C/expressions.py:37, in generator(ir, options)
     34 d["name_from_uflfile"] = ir.name_from_uflfile
     35 d["factory_name"] = ir.name
---> 37 parts = eg.generate()
     39 CF = CFormatter(options["scalar_type"])
     40 d["tabulate_expression"] = CF.c_format(parts)

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/expression_generator.py:41, in ExpressionGenerator.generate(self)
     39 # Generate the tables of geometry data that are needed
     40 parts += self.generate_geometry_tables()
---> 41 parts += self.generate_piecewise_partition()
     43 all_preparts = []
     44 all_quadparts = []

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/expression_generator.py:145, in ExpressionGenerator.generate_piecewise_partition(self)
    142 F = self.ir.integrand[self.quadrature_rule]["factorization"]
    144 arraysymbol = L.Symbol("sp", dtype=L.DataType.SCALAR)
--> 145 parts = self.generate_partition(arraysymbol, F, "piecewise")
    146 parts = L.commented_code_list(parts, "Unstructured piecewise computations")
    147 return parts

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/expression_generator.py:348, in ExpressionGenerator.generate_partition(self, symbol, F, mode)
    346 # Mapping UFL operator to target language
    347 self._ufl_names.add(v._ufl_handler_name_)
--> 348 vexpr = L.ufl_to_lnodes(v, *vops)
    350 # Create a new intermediate for each subexpression
    351 # except boolean conditions and its childs
    352 if isinstance(parent_exp, ufl.classes.Condition):
    353     # Skip intermediates for 'x' and 'y' in x<y
    354     # Avoid the creation of complex valued intermediates

File ~/anaconda3/envs/env3-10-complex/lib/python3.10/site-packages/ffcx/codegeneration/lnodes.py:920, in ufl_to_lnodes(operator, *args)
    918     return _ufl_call_lookup[optype](operator, *args)
    919 else:
--> 920     raise RuntimeError(f"Missing lookup for expr type {optype}.")

RuntimeError: Missing lookup for expr type <class 'ufl.mathfunctions.BesselI'>.