20030291596

Committed 08 Dec 2025 01:46PM UTC coverage: 84.453% (+1.4%) from 83.044%

Build # 20030291596

Build Type

Pull #801

github

Committed by

schnellerhase

Commit Message

docstring

Pull Request Pull Request #801: Add `numba` backend

Coverage Stats

355 of 386 new or added lines in 22 files covered. (91.97%)

4074 of 4824 relevant lines covered (84.45%)

0.84 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

79.59

/ffcx/codegeneration/numba/implementation.py

# Copyright (C) 2025 Chris Richardson and Paul T. Kühner
#
# This file is part of FFCx. (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Numba implementation for output."""

import numpy as np
from numpy import typing as npt

import ffcx.codegeneration.lnodes as L
from ffcx.codegeneration.utils import dtype_to_scalar_dtype


def build_initializer_lists(values: npt.NDArray) -> str:
    """Build list of values."""
    arr = "["
    if len(values.shape) == 1:
        return "[" + ", ".join(str(v) for v in values) + "]"
    elif len(values.shape) > 1:
        arr += ",\n".join(build_initializer_lists(v) for v in values)
    arr += "]"
    return arr


class Formatter:
    """Implementation for numba output backend."""

    scalar_type: np.dtype
    real_type: np.dtype

    def __init__(self, dtype: npt.DTypeLike) -> None:
        """Initialise."""
        self.scalar_type = np.dtype(dtype)
        self.real_type = dtype_to_scalar_dtype(dtype)

    def _dtype_to_name(self, dtype: L.DataType) -> str:
        """Convert dtype to Python name."""
        if dtype == L.DataType.SCALAR:
            return f"np.{self.scalar_type}"
        if dtype == L.DataType.REAL:
            return f"np.{self.real_type}"
        if dtype == L.DataType.INT:
            return f"np.{np.int32}"
        if dtype == L.DataType.BOOL:
            return f"np.{np.bool}"
        raise ValueError(f"Invalid dtype: {dtype}")

    def format_section(self, section: L.Section) -> str:
        """Format a section."""
        # add new line before section
        comments = self._format_comment_str("------------------------")
        comments += self._format_comment_str(f"Section: {section.name}")
        comments += self._format_comment_str(f"Inputs: {', '.join(w.name for w in section.input)}")
        comments += self._format_comment_str(
            f"Outputs: {', '.join(w.name for w in section.output)}"
        )
        declarations = "".join(self.format(s) for s in section.declarations)

        body = ""
        if len(section.statements) > 0:
            body = "".join(self.format(s) for s in section.statements)

        body += self._format_comment_str("------------------------")
        return comments + declarations + body

    def format_statement_list(self, slist: L.StatementList) -> str:
        """Format a list of statements."""
        output = ""
        for s in slist.statements:
            output += self.format(s)
        return output

    def _format_comment_str(self, comment: str) -> str:
        """Format str to comment string."""
        return f"# {comment} \n"

    def format_comment(self, c: L.Comment) -> str:
        """Format a comment."""
        return self._format_comment_str(c.comment)

    def format_array_decl(self, arr: L.ArrayDecl) -> str:
        """Format an array declaration."""
        dtype = arr.symbol.dtype
        typename = self._dtype_to_name(dtype)

        symbol = self.format(arr.symbol)
        if arr.values is None:
            return f"{symbol} = np.empty({arr.sizes}, dtype={typename})\n"
        elif arr.values.size == 1:
            return f"{symbol} = np.full({arr.sizes}, {arr.values[0]}, dtype={typename})\n"
        av = build_initializer_lists(arr.values)
        av = f"np.array({av}, dtype={typename})"
        return f"{symbol} = {av}\n"

    def format_array_access(self, arr: L.ArrayAccess) -> str:
        """Format array access."""
        array = self.format(arr.array)
        idx = ", ".join(self.format(ix) for ix in arr.indices)
        return f"{array}[{idx}]"

    def format_multi_index(self, index: L.MultiIndex) -> str:
        """Format a multi-index."""
        return self.format(index.global_index)

    def format_variable_decl(self, v: L.VariableDecl) -> str:
        """Format a variable declaration."""
        sym = self.format(v.symbol)
        val = self.format(v.value)
        return f"{sym} = {val}\n"

    def format_nary_op(self, oper: L.NaryOp) -> str:
        """Format a n argument operation."""
        # Format children
        args = [self.format(arg) for arg in oper.args]

        # Apply parentheses
        for i in range(len(args)):
            if oper.args[i].precedence >= oper.precedence:
                args[i] = f"({args[i]})"

        # Return combined string
        return f" {oper.op} ".join(args)

    def format_binary_op(self, oper: L.BinOp) -> str:
        """Format a binary operation."""
        # Format children
        lhs = self.format(oper.lhs)
        rhs = self.format(oper.rhs)

        # Apply parentheses
        if oper.lhs.precedence >= oper.precedence:
            lhs = f"({lhs})"
        if oper.rhs.precedence >= oper.precedence:
            rhs = f"({rhs})"

        # Return combined string
        return f"{lhs} {oper.op} {rhs}"

    def format_neg(self, val: L.Neg) -> str:
        """Format unary negation."""
        arg = self.format(val.arg)
        return f"-{arg}"

    def format_not(self, val: L.Not) -> str:
        """Format not operation."""
        arg = self.format(val.arg)
        return f"not({arg})"

    def format_andor(self, oper: L.And | L.Or) -> str:
        """Format and or or operation."""
        # Format children
        lhs = self.format(oper.lhs)
        rhs = self.format(oper.rhs)

        # Apply parentheses
        if oper.lhs.precedence >= oper.precedence:
            lhs = f"({lhs})"
        if oper.rhs.precedence >= oper.precedence:
            rhs = f"({rhs})"

        opstr = {"||": "or", "&&": "and"}[oper.op]

        # Return combined string
        return f"{lhs} {opstr} {rhs}"

    def format_literal_float(self, val: L.LiteralFloat) -> str:
        """Format a literal float."""
        return f"{val.value}"

    def format_literal_int(self, val: L.LiteralInt) -> str:
        """Format a literal int."""
        return f"{val.value}"

    def format_for_range(self, r: L.ForRange) -> str:
        """Format a loop over a range."""
        begin = self.format(r.begin)
        end = self.format(r.end)
        index = self.format(r.index)
        output = f"for {index} in range({begin}, {end}):\n"
        b = self.format(r.body).split("\n")
        for line in b:
            output += f"    {line}\n"
        return output

    def format_statement(self, s: L.Statement) -> str:
        """Format a statement."""
        return self.format(s.expr)

    def format_assign(self, expr: L.Assign) -> str:
        """Format assignment."""
        rhs = self.format(expr.rhs)
        lhs = self.format(expr.lhs)
        return f"{lhs} {expr.op} {rhs}\n"

    def format_conditional(self, s: L.Conditional) -> str:
        """Format a conditional."""
        # Format children
        c = self.format(s.condition)
        t = self.format(s.true)
        f = self.format(s.false)

        # Apply parentheses
        if s.condition.precedence >= s.precedence:
            c = f"({c})"
        if s.true.precedence >= s.precedence:
            t = f"({t})"
        if s.false.precedence >= s.precedence:
            f = f"({f})"

        # Return combined string
        return f"({t} if {c} else {f})"

    def format_symbol(self, s: L.Symbol) -> str:
        """Format a symbol."""
        return f"{s.name}"

    def format_mathfunction(self, f: L.MathFunction) -> str:
        """Format a math function."""
        function_map = {
            "ln": "log",
            "acos": "arccos",
            "asin": "arcsin",
            "atan": "arctan",
            "atan2": "arctan2",
            "acosh": "arccosh",
            "asinh": "arcsinh",
            "atanh": "arctanh",
        }
        function = function_map.get(f.function, f.function)
        args = [self.format(arg) for arg in f.args]
        if "bessel" in function:
            # TODO: fix with https://github.com/numba/numba/issues/2312.
            raise NotImplementedError("Bessel function not supported.")
        if function == "erf":
            return f"math.erf({args[0]})"
        argstr = ", ".join(args)
        return f"np.{function}({argstr})"

    impl = {
        "StatementList": format_statement_list,
        "Comment": format_comment,
        "Section": format_section,
        "ArrayDecl": format_array_decl,
        "ArrayAccess": format_array_access,
        "MultiIndex": format_multi_index,
        "VariableDecl": format_variable_decl,
        "ForRange": format_for_range,
        "Statement": format_statement,
        "Assign": format_assign,
        "AssignAdd": format_assign,
        "Product": format_nary_op,
        "Sum": format_nary_op,
        "Add": format_binary_op,
        "Sub": format_binary_op,
        "Mul": format_binary_op,
        "Div": format_binary_op,
        "Neg": format_neg,
        "Not": format_not,
        "LiteralFloat": format_literal_float,
        "LiteralInt": format_literal_int,
        "Symbol": format_symbol,
        "Conditional": format_conditional,
        "MathFunction": format_mathfunction,
        "And": format_andor,
        "Or": format_andor,
        "NE": format_binary_op,
        "EQ": format_binary_op,
        "GE": format_binary_op,
        "LE": format_binary_op,
        "GT": format_binary_op,
        "LT": format_binary_op,
    }

    def format(self, s: L.LNode) -> str:
        """Format output."""
        name = s.__class__.__name__
        try:
            return self.impl[name](self, s)  # type: ignore
        except KeyError:
            raise RuntimeError("Unknown statement: ", name)

1	# Copyright (C) 2025 Chris Richardson and Paul T. Kühner
2	#
3	# This file is part of FFCx. (https://www.fenicsproject.org)
4	#
5	# SPDX-License-Identifier: LGPL-3.0-or-later
6	"""Numba implementation for output."""
7
8	import numpy as np	1✔
9	from numpy import typing as npt	1✔
10
11	import ffcx.codegeneration.lnodes as L	1✔
12	from ffcx.codegeneration.utils import dtype_to_scalar_dtype	1✔
13
14
15	def build_initializer_lists(values: npt.NDArray) -> str:	1✔
16	"""Build list of values."""
17	arr = "["	1✔
18	if len(values.shape) == 1:	1✔
19	return "[" + ", ".join(str(v) for v in values) + "]"	1✔
20	elif len(values.shape) > 1:	1✔
21	arr += ",\n".join(build_initializer_lists(v) for v in values)	1✔
22	arr += "]"	1✔
23	return arr	1✔
24
25
26	class Formatter:	1✔
27	"""Implementation for numba output backend."""
28
29	scalar_type: np.dtype	1✔
30	real_type: np.dtype	1✔
31
32	def __init__(self, dtype: npt.DTypeLike) -> None:	1✔
33	"""Initialise."""
34	self.scalar_type = np.dtype(dtype)	1✔
35	self.real_type = dtype_to_scalar_dtype(dtype)	1✔
36
37	def _dtype_to_name(self, dtype: L.DataType) -> str:	1✔
38	"""Convert dtype to Python name."""
39	if dtype == L.DataType.SCALAR:	1✔
40	return f"np.{self.scalar_type}"	1✔
41	if dtype == L.DataType.REAL:	1✔
42	return f"np.{self.real_type}"	1✔
NEW 43	if dtype == L.DataType.INT:	×
NEW 44	return f"np.{np.int32}"	×
NEW 45	if dtype == L.DataType.BOOL:	×
NEW 46	return f"np.{np.bool}"	×
NEW 47	raise ValueError(f"Invalid dtype: {dtype}")	×
48
49	def format_section(self, section: L.Section) -> str:	1✔
50	"""Format a section."""
51	# add new line before section
52	comments = self._format_comment_str("------------------------")	1✔
53	comments += self._format_comment_str(f"Section: {section.name}")	1✔
54	comments += self._format_comment_str(f"Inputs: {', '.join(w.name for w in section.input)}")	1✔
55	comments += self._format_comment_str(	1✔
56	f"Outputs: {', '.join(w.name for w in section.output)}"
57	)
58	declarations = "".join(self.format(s) for s in section.declarations)	1✔
59
60	body = ""	1✔
61	if len(section.statements) > 0:	1✔
62	body = "".join(self.format(s) for s in section.statements)	1✔
63
64	body += self._format_comment_str("------------------------")	1✔
65	return comments + declarations + body	1✔
66
67	def format_statement_list(self, slist: L.StatementList) -> str:	1✔
68	"""Format a list of statements."""
69	output = ""	1✔
70	for s in slist.statements:	1✔
71	output += self.format(s)	1✔
72	return output	1✔
73
74	def _format_comment_str(self, comment: str) -> str:	1✔
75	"""Format str to comment string."""
76	return f"# {comment} \n"	1✔
77
78	def format_comment(self, c: L.Comment) -> str:	1✔
79	"""Format a comment."""
80	return self._format_comment_str(c.comment)	1✔
81
82	def format_array_decl(self, arr: L.ArrayDecl) -> str:	1✔
83	"""Format an array declaration."""
84	dtype = arr.symbol.dtype	1✔
85	typename = self._dtype_to_name(dtype)	1✔
86
87	symbol = self.format(arr.symbol)	1✔
88	if arr.values is None:	1✔
NEW 89	return f"{symbol} = np.empty({arr.sizes}, dtype={typename})\n"	×
90	elif arr.values.size == 1:	1✔
91	return f"{symbol} = np.full({arr.sizes}, {arr.values[0]}, dtype={typename})\n"	1✔
92	av = build_initializer_lists(arr.values)	1✔
93	av = f"np.array({av}, dtype={typename})"	1✔
94	return f"{symbol} = {av}\n"	1✔
95
96	def format_array_access(self, arr: L.ArrayAccess) -> str:	1✔
97	"""Format array access."""
98	array = self.format(arr.array)	1✔
99	idx = ", ".join(self.format(ix) for ix in arr.indices)	1✔
100	return f"{array}[{idx}]"	1✔
101
102	def format_multi_index(self, index: L.MultiIndex) -> str:	1✔
103	"""Format a multi-index."""
104	return self.format(index.global_index)	1✔
105
106	def format_variable_decl(self, v: L.VariableDecl) -> str:	1✔
107	"""Format a variable declaration."""
108	sym = self.format(v.symbol)	1✔
109	val = self.format(v.value)	1✔
110	return f"{sym} = {val}\n"	1✔
111
112	def format_nary_op(self, oper: L.NaryOp) -> str:	1✔
113	"""Format a n argument operation."""
114	# Format children
115	args = [self.format(arg) for arg in oper.args]	1✔
116
117	# Apply parentheses
118	for i in range(len(args)):	1✔
119	if oper.args[i].precedence >= oper.precedence:	1✔
120	args[i] = f"({args[i]})"	1✔
121
122	# Return combined string
123	return f" {oper.op} ".join(args)	1✔
124
125	def format_binary_op(self, oper: L.BinOp) -> str:	1✔
126	"""Format a binary operation."""
127	# Format children
128	lhs = self.format(oper.lhs)	1✔
129	rhs = self.format(oper.rhs)	1✔
130
131	# Apply parentheses
132	if oper.lhs.precedence >= oper.precedence:	1✔
133	lhs = f"({lhs})"	1✔
134	if oper.rhs.precedence >= oper.precedence:	1✔
135	rhs = f"({rhs})"	1✔
136
137	# Return combined string
138	return f"{lhs} {oper.op} {rhs}"	1✔
139
140	def format_neg(self, val: L.Neg) -> str:	1✔
141	"""Format unary negation."""
142	arg = self.format(val.arg)	1✔
143	return f"-{arg}"	1✔
144
145	def format_not(self, val: L.Not) -> str:	1✔
146	"""Format not operation."""
NEW 147	arg = self.format(val.arg)	×
NEW 148	return f"not({arg})"	×
149
150	def format_andor(self, oper: L.And \| L.Or) -> str:	1✔
151	"""Format and or or operation."""
152	# Format children
NEW 153	lhs = self.format(oper.lhs)	×
NEW 154	rhs = self.format(oper.rhs)	×
155
156	# Apply parentheses
NEW 157	if oper.lhs.precedence >= oper.precedence:	×
NEW 158	lhs = f"({lhs})"	×
NEW 159	if oper.rhs.precedence >= oper.precedence:	×
NEW 160	rhs = f"({rhs})"	×
161
NEW 162	opstr = {"\|\|": "or", "&&": "and"}[oper.op]	×
163
164	# Return combined string
NEW 165	return f"{lhs} {opstr} {rhs}"	×
166
167	def format_literal_float(self, val: L.LiteralFloat) -> str:	1✔
168	"""Format a literal float."""
169	return f"{val.value}"	1✔
170
171	def format_literal_int(self, val: L.LiteralInt) -> str:	1✔
172	"""Format a literal int."""
173	return f"{val.value}"	1✔
174
175	def format_for_range(self, r: L.ForRange) -> str:	1✔
176	"""Format a loop over a range."""
177	begin = self.format(r.begin)	1✔
178	end = self.format(r.end)	1✔
179	index = self.format(r.index)	1✔
180	output = f"for {index} in range({begin}, {end}):\n"	1✔
181	b = self.format(r.body).split("\n")	1✔
182	for line in b:	1✔
183	output += f" {line}\n"	1✔
184	return output	1✔
185
186	def format_statement(self, s: L.Statement) -> str:	1✔
187	"""Format a statement."""
188	return self.format(s.expr)	1✔
189
190	def format_assign(self, expr: L.Assign) -> str:	1✔
191	"""Format assignment."""
192	rhs = self.format(expr.rhs)	1✔
193	lhs = self.format(expr.lhs)	1✔
194	return f"{lhs} {expr.op} {rhs}\n"	1✔
195
196	def format_conditional(self, s: L.Conditional) -> str:	1✔
197	"""Format a conditional."""
198	# Format children
NEW 199	c = self.format(s.condition)	×
NEW 200	t = self.format(s.true)	×
NEW 201	f = self.format(s.false)	×
202
203	# Apply parentheses
NEW 204	if s.condition.precedence >= s.precedence:	×
NEW 205	c = f"({c})"	×
NEW 206	if s.true.precedence >= s.precedence:	×
NEW 207	t = f"({t})"	×
NEW 208	if s.false.precedence >= s.precedence:	×
NEW 209	f = f"({f})"	×
210
211	# Return combined string
NEW 212	return f"({t} if {c} else {f})"	×
213
214	def format_symbol(self, s: L.Symbol) -> str:	1✔
215	"""Format a symbol."""
216	return f"{s.name}"	1✔
217
218	def format_mathfunction(self, f: L.MathFunction) -> str:	1✔
219	"""Format a math function."""
220	function_map = {	1✔
221	"ln": "log",
222	"acos": "arccos",
223	"asin": "arcsin",
224	"atan": "arctan",
225	"atan2": "arctan2",
226	"acosh": "arccosh",
227	"asinh": "arcsinh",
228	"atanh": "arctanh",
229	}
230	function = function_map.get(f.function, f.function)	1✔
231	args = [self.format(arg) for arg in f.args]	1✔
232	if "bessel" in function:	1✔
233	# TODO: fix with https://github.com/numba/numba/issues/2312.
NEW 234	raise NotImplementedError("Bessel function not supported.")	×
235	if function == "erf":	1✔
NEW 236	return f"math.erf({args[0]})"	×
237	argstr = ", ".join(args)	1✔
238	return f"np.{function}({argstr})"	1✔
239
240	impl = {	1✔
241	"StatementList": format_statement_list,
242	"Comment": format_comment,
243	"Section": format_section,
244	"ArrayDecl": format_array_decl,
245	"ArrayAccess": format_array_access,
246	"MultiIndex": format_multi_index,
247	"VariableDecl": format_variable_decl,
248	"ForRange": format_for_range,
249	"Statement": format_statement,
250	"Assign": format_assign,
251	"AssignAdd": format_assign,
252	"Product": format_nary_op,
253	"Sum": format_nary_op,
254	"Add": format_binary_op,
255	"Sub": format_binary_op,
256	"Mul": format_binary_op,
257	"Div": format_binary_op,
258	"Neg": format_neg,
259	"Not": format_not,
260	"LiteralFloat": format_literal_float,
261	"LiteralInt": format_literal_int,
262	"Symbol": format_symbol,
263	"Conditional": format_conditional,
264	"MathFunction": format_mathfunction,
265	"And": format_andor,
266	"Or": format_andor,
267	"NE": format_binary_op,
268	"EQ": format_binary_op,
269	"GE": format_binary_op,
270	"LE": format_binary_op,
271	"GT": format_binary_op,
272	"LT": format_binary_op,
273	}
274
275	def format(self, s: L.LNode) -> str:	1✔
276	"""Format output."""
277	name = s.__class__.__name__	1✔
278	try:	1✔
279	return self.impl[name](self, s) # type: ignore	1✔
NEW 280	except KeyError:	×
NEW 281	raise RuntimeError("Unknown statement: ", name)	×

FEniCS / ffcx / 20030291596

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous