Remove numba backend (#55)

* Remove numba backend

* Manually type numba functions

* Ensure linearized system is contiguous and float64

* Remove njit from cycle reduction functions (windows CI fails mysteriously)

Author: jessegrabowski

gEconpy/model/compile.py

@@ -10,9 +10,8 @@
 from sympytensor import as_tensor
 
 from gEconpy.classes.containers import SteadyStateResults
-from gEconpy.numbaf.utilities import numba_lambdify
 
-BACKENDS = Literal["numpy", "numba", "pytensor"]
+BACKENDS = Literal["numpy", "pytensor"]
 
 
 def output_to_tensor(x, cache):
@@ -151,7 +150,7 @@ def compile_function(
     **kwargs,
 ) -> tuple[Callable, dict]:
     """
-    Dispatch compilation of a sympy function to one of three possible backends: numpy, numba, or pytensor.
+    Dispatch compilation of a sympy function to one of three possible backends: numpy, or pytensor.
 
     Parameters
     ----------
@@ -161,7 +160,7 @@ def compile_function(
     outputs: list[Union[sp.Symbol, sp.Expr]]
         The outputs of the function.
 
-    backend: str, one of "numpy", "numba", "pytensor"
+    backend: str, one of "numpy" or "pytensor"
         The backend to use for the compiled function.
 
     cache: dict, optional
@@ -195,10 +194,6 @@ def compile_function(
         f, cache = compile_to_numpy(
             inputs, outputs, cache, stack_return, pop_return, **kwargs
         )
-    elif backend == "numba":
-        f, cache = compile_to_numba(
-            inputs, outputs, cache, stack_return, pop_return, **kwargs
-        )
     elif backend == "pytensor":
         f, cache = compile_to_pytensor_function(
             inputs, outputs, cache, stack_return, pop_return, return_symbolic, **kwargs
@@ -256,48 +251,6 @@ def compile_to_numpy(
     return f, cache
 
 
-def compile_to_numba(
-    inputs: list[sp.Symbol],
-    outputs: list[sp.Symbol | sp.Expr],
-    cache: dict,
-    stack_return: bool,
-    pop_return: bool,
-    **kwargs,
-):
-    """
-    Convert a sympy function to a numba njit function using :func:`numba_lambdify`.
-
-    Parameters
-    ----------
-    inputs: list[sp.Symbol]
-        The inputs to the function.
-    outputs: list[Union[sp.Symbol, sp.Expr]]
-        The outputs of the function.
-    cache: dict
-        Mapping between sympy variables and pytensor variables. Ignored by this function; included for compatibility
-        with other compile functions.
-    stack_return: bool
-        If True, the function will return a single numpy array with all outputs. Otherwise it will return a list
-        of numpy arrays.
-    pop_return: bool
-        If True, the function will return only the 0th element of the output. Used to remove the list wrapper around
-        single-output functions.
-    kwargs: dict
-        Ignored, included for compatibility with other compile functions
-
-    Returns
-    -------
-    f: Callable
-        The compiled function.
-    cache: dict
-        Pytensor caching information.
-    """
-    f = numba_lambdify(inputs, outputs, stack_outputs=stack_return)
-    if pop_return:
-        f = pop_return_wrapper(f)
-    return f, cache
-
-
 def compile_to_pytensor_function(
     inputs: list[sp.Symbol],
     outputs: list[sp.Symbol | sp.Expr],

gEconpy/model/model.py

@@ -6,7 +6,6 @@
 from copy import deepcopy
 from typing import Literal, cast
 
-import numba as nb
 import numpy as np
 import pandas as pd
 import sympy as sp
@@ -42,6 +41,7 @@
 )
 from gEconpy.utilities import get_name, postprocess_optimizer_res, safe_to_ss
 
+
 VariableType = sp.Symbol | TimeAwareSymbol
 _log = logging.getLogger(__name__)
 
@@ -1133,7 +1133,10 @@ def linearize_model(
             **param_dict, **steady_state, not_loglin_variable=not_loglin_flags
         )
 
-        return A, B, C, D
+        # Using A.dtype to avoid hard-coding float64 (we might be using float32)
+        # The reason for casting is mostly D, which sometimes comes out as an int32/64 array
+
+        return list(map(lambda x: np.ascontiguousarray(x, dtype=A.dtype), [A, B, C, D]))
 
     def solve_model(
         self,
@@ -1337,6 +1340,8 @@ def solve_model(
             **parameter_updates,
         )
 
+        assert all(x.flags["C_CONTIGUOUS"] for x in [A, B, C, D])
+
         if solver == "gensys":
             gensys_results = solve_policy_function_with_gensys(A, B, C, D, tol)
             G_1, constant, impact, f_mat, f_wt, y_wt, gev, eu, loose = gensys_results

gEconpy/model/parameters.py

@@ -29,7 +29,7 @@ def compile_param_dict_func(
     deterministic_dict: SymbolDictionary
         A dictionary of deterministic parameters, with the keys being the parameters and the values being the
         expressions to compute them.
-    backend: str, one of "numpy", "numba", "pytensor"
+    backend: str, one of "numpy", "pytensor"
         The backend to use for the compiled function.
     cache: dict, optional
         A dictionary mapping from pytensor symbols to sympy expressions. Used to prevent duplicate mappings from

gEconpy/model/steady_state.py

@@ -158,11 +158,9 @@ def compile_ss_resid_and_sq_err(
         backend=backend,
         cache=cache,
         return_symbolic=return_symbolic,
-        # for pytensor/numba, the return is a single object; don't stack into a (1,n,n) array
+        # for pytensor, the return is a single object; don't stack into a (1,n,n) array
         stack_return=backend == "numpy",
-        # Numba directly returns the jacobian as an array, don't pop
-        # pytensor and lambdify return a list of one item, so we have to extract it.
-        pop_return=backend != "numba",
+        pop_return=True,
         **kwargs,
     )
 
@@ -209,8 +207,7 @@ def compile_ss_resid_and_sq_err(
         return_symbolic=return_symbolic,
         # error_hess is a list of one element; don't stack into a (1,n,n) array
         stack_return=backend != "pytensor",
-        # Numba directly returns the hessian as an array, don't pop
-        pop_return=backend != "numba",
+        pop_return=True,
         **kwargs,
     )
 

gEconpy/numbaf/utilities.py

@@ -166,159 +166,3 @@ def _print_exp(self, expr):
             self._module_format(self._module + ".exp"),
             ",".join(self._print(i) for i in expr.args),
         )
-
-
-def numba_lambdify(
-    inputs: list[sp.Symbol],
-    expr: list[sp.Expr] | sp.Matrix | list[sp.Matrix],
-    func_signature: str | None = None,
-    ravel_outputs=False,
-    stack_outputs=False,
-) -> Callable:
-    """
-    Convert a sympy expression into a Numba-compiled function.  Unlike sp.lambdify, the resulting function can be
-    pickled. In addition, common sub-expressions are gathered using sp.cse and assigned to local variables,
-    giving a (very) modest performance boost. A signature can optionally be provided for numba.njit.
-
-    Finally, the resulting function always returns a numpy array, rather than a list.
-
-    Parameters
-    ----------
-    inputs: list of sympy.Symbol
-        A list of "exogenous" variables. The distinction between "exogenous" and "enodgenous" is
-        useful when passing the resulting function to a scipy.optimize optimizer. In this context, exogenous
-        variables should be the choice varibles used to minimize the function.
-    expr : list of sympy.Expr or sp.Matrix
-        The sympy expression(s) to be converted. Expects a list of expressions (in the case that we're compiling a
-        system to be stacked into a single output vector), a single matrix (which is returned as a single nd-array)
-        or a list of matrices (which are returned as a list of nd-arrays)
-    func_signature: str
-        A numba function signature, passed to the numba.njit decorator on the generated function.
-    ravel_outputs: bool, default False
-        If true, all outputs of the jitted function will be raveled before they are returned. This is useful for
-        removing size-1 dimensions from sympy vectors.
-    stack_outputs: bool, default False
-        If true, stack all return values into a single vector. Otherwise they are returned as a tuple as usual.
-
-    Returns
-    -------
-    numba.types.function
-        A Numba-compiled function equivalent to the input expression.
-
-    Notes
-    -----
-    The function returned by this function is pickleable.
-    """
-    ZERO_PATTERN = re.compile(r"(?<![\.\w])0([ ,\]])")
-    ZERO_ONE_INDEX_PATTERN = re.compile(r"((?<=\[)(([0,1])\.0)(?=\]))")
-    FLOAT_SUBS = {
-        sp.core.numbers.One(): sp.Float(1),
-        sp.core.numbers.NegativeOne(): sp.Float(-1),
-    }
-    printer = NumbaFriendlyNumPyPrinter()
-
-    if func_signature is None:
-        decorator = "@nb.njit"
-    else:
-        decorator = f"@nb.njit({func_signature})"
-
-    # Special case: expr is [[]]. This can occur if no user-defined steady-state values were provided.
-    # It shouldn't happen otherwise.
-    if expr == [[]]:
-        sub_dict = ()
-        code = ""
-        retvals = ["[None]"]
-
-    else:
-        # Need to make the float substitutions so that numba can correctly interpret everything, but we have to handle
-        # several cases:
-        # Case 1: expr is just a single Sympy thing
-        if isinstance(expr, sp.Matrix | sp.Expr):
-            expr = expr.subs(FLOAT_SUBS)
-
-        # Case 2: expr is a list. Items in the list are either lists of expressions (systems of equations),
-        # single equations, or matrices.
-        elif isinstance(expr, list):
-            new_expr = []
-            for item in expr:
-                # Case 2a: It's a simple list of sympy things
-                if isinstance(item, sp.Matrix | sp.Expr):
-                    new_expr.append(item.subs(FLOAT_SUBS))
-                # Case 2b: It's a system of equations, List[List[sp.Expr]]
-                elif isinstance(item, list):
-                    if all([isinstance(x, sp.Matrix | sp.Expr) for x in item]):
-                        new_expr.append([x.subs(FLOAT_SUBS) for x in item])
-                    else:
-                        raise ValueError("Unexpected input type for expr")
-
-                # Case 2c: It's a constant -- just pass it along unchanged.
-                elif isinstance(item, int | float):
-                    new_expr.append(item)
-                else:
-                    raise ValueError(f"Unexpected input type for expr: {expr}")
-
-            expr = new_expr
-        else:
-            raise ValueError("Unexpected input type for expr")
-        sub_dict, expr = sp.cse(expr)
-
-        # Converting matrices to a list of lists is convenient because NumPyPrinter() won't wrap them in np.array
-        exprs = []
-        for ex in expr:
-            if hasattr(ex, "tolist"):
-                exprs.append(ex.tolist())
-            else:
-                exprs.append(ex)
-
-        codes = []
-        retvals = []
-        for i, expr in enumerate(exprs):
-            code = printer.doprint(expr)
-
-            delimiter = "]," if "]," in code else ","
-            delimiter = ","
-            code = code.split(delimiter)
-            code = [" " * 8 + eq.strip() for eq in code]
-            code = f"{delimiter}\n".join(code)
-            code = code.replace("numpy.", "np.")
-
-            # Handle conversion of 0 to 0.0
-            code = re.sub(ZERO_PATTERN, r"0.0\g<1>", code)
-
-            # Repair indexing -- we might have converted x[0] to x[0.0] or x[1] to x[1.0]
-            code = re.sub(ZERO_ONE_INDEX_PATTERN, r"\g<3>", code)
-
-            code_name = f"retval_{i}"
-            retvals.append(code_name)
-            code = f"    {code_name} = np.array(\n{code}\n    )"
-            if ravel_outputs:
-                code += ".ravel()"
-
-            codes.append(code)
-        code = "\n".join(codes)
-
-    input_signature = ", ".join([getattr(x, "safe_name", x.name) for x in inputs])
-
-    assignments = "\n".join(
-        [
-            f"    {x} = {printer.doprint(y).replace('numpy.', 'np.')}"
-            for x, y in sub_dict
-        ]
-    )
-    assignments = re.sub(ZERO_ONE_INDEX_PATTERN, r"\g<3>", assignments)
-
-    if len(retvals) > 1:
-        returns = f"({','.join(retvals)})"
-        if stack_outputs:
-            returns = f"np.stack({returns})"
-    else:
-        returns = retvals[0]
-    # returns = f'[{",".join(retvals)}]' if len(retvals) > 1 else retvals[0]
-    full_code = f"{decorator}\ndef f({input_signature}):\n\n{assignments}\n\n{code}\n\n    return {returns}"
-
-    docstring = f"'''Automatically generated code:\n{full_code}'''"
-    code = f"{decorator}\ndef f({input_signature}):\n    {docstring}\n\n{assignments}\n\n{code}\n\n    return {returns}"
-
-    exec_namespace = {}
-    exec(code, globals(), exec_namespace)
-    return exec_namespace["f"]

gEconpy/solvers/cycle_reduction.py

@@ -15,7 +15,11 @@
 )
 
 
-# @nb.njit(cache=True)
+# TODO: These njit decorators cause the CI to fail on Windows only -- no idea why. Disabling it for now.
+# @nb.njit(
+#     (nb.float64[:, ::1], nb.float64[:, ::1], nb.float64[:, ::1], nb.int64, nb.float64),
+#     cache=True,
+# )
 def nb_cycle_reduction(
     A0: np.ndarray,
     A1: np.ndarray,
@@ -113,8 +117,11 @@ def nb_cycle_reduction(
     return X, res, result, log_norm
 
 
-# @nb.njit(cache=True)
-def nb_solve_shock_matrix(B, C, D, G_1):
+# @nb.njit(
+#     (nb.float64[:, ::1], nb.float64[:, ::1], nb.float64[:, ::1], nb.float64[:, ::1]),
+#     cache=True,
+# )
+def nb_solve_shock_matrix(B: np.ndarray, C: np.ndarray, D: np.ndarray, G_1: np.ndarray):
     """
     Given the partial solution to the linear approximate policy function G_1, solve for the remaining component of the
     policy function, R.

tests/model/test_compile.py

@@ -32,8 +32,8 @@ def f():
     assert result["b"] == 2.0
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba", "pytensor"])
-def test_scalar_function(backend: Literal["numpy", "numba", "pytensor"]):
+@pytest.mark.parametrize("backend", ["numpy", "pytensor"])
+def test_scalar_function(backend: Literal["numpy", "pytensor"]):
     x = sp.symbols("x")
     f = x**2
     f_func, _ = compile_function(
@@ -42,11 +42,9 @@ def test_scalar_function(backend: Literal["numpy", "numba", "pytensor"]):
     assert f_func(x=2) == 4
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba", "pytensor"])
+@pytest.mark.parametrize("backend", ["numpy", "pytensor"])
 @pytest.mark.parametrize("stack_return", [True, False])
-def test_multiple_outputs(
-    backend: Literal["numpy", "numba", "pytensor"], stack_return: bool
-):
+def test_multiple_outputs(backend: Literal["numpy", "pytensor"], stack_return: bool):
     x, y, z = sp.symbols("x y z ")
     x2 = x**2
     y2 = y**2
@@ -68,8 +66,8 @@ def test_multiple_outputs(
     )
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba", "pytensor"])
-def test_matrix_function(backend: Literal["numpy", "numba", "pytensor"]):
+@pytest.mark.parametrize("backend", ["numpy", "pytensor"])
+def test_matrix_function(backend: Literal["numpy", "pytensor"]):
     x, y, z = sp.symbols("x y z")
     f = sp.Matrix([x, y, z]).reshape(1, 3)
 
@@ -87,7 +85,7 @@ def test_matrix_function(backend: Literal["numpy", "numba", "pytensor"]):
     np.testing.assert_allclose(res, np.array([[2.0, 3.0, 4.0]]))
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba", "pytensor"])
+@pytest.mark.parametrize("backend", ["numpy", "pytensor"])
 def test_compile_gradient(backend: BACKENDS):
     x, y, z = sp.symbols("x y z")
     f = x**2 + y**2 + z**2