pockit.base.fastfunc

class FastFunc:

JIT-compiled, vectorized functions for values and symbolic derivatives.

FastFunc takes a SymPy expression and a list of SymPy symbols as function arguments and generates JIT-compiled, vectorized functions for values, gradients, and Hessians. The derivatives use a sparse format, so only nonzero elements are computed.

FastFunc( function: int | float | sympy.core.expr.Expr, args: list[sympy.core.symbol.Symbol], simplify: bool = False, fastmath: bool = False, *, cache: Optional[str] = None)

Suppose the expression of the input function is f(a_1, a_2, ..., a_n), with n arguments (a_1, a_2, ..., a_n) as the second argument args. The generated functions F, G, and H take two arguments x and k, where x is a 1D array of length n * k, and k is an integer. The first k elements of x are the values of a_1 at k different points, the next k elements are the values of a_2, and so on. The return value of F is a 1D array of length k, where the i-th element is the value of f(a_1, a_2, ..., a_n) at the i-th point. The return value of G is a 2D array of shape (len(G_index), k), where G_index contains the indices of non-zero elements in the gradient matrix. The return value of H is a 2D array of shape (len(H_index_row), k), where H_index_row, H_index_col contain the indices of nonzero elements in the lower triangular part of the Hessian matrix.

If simplify is True, every symbolic expression will be simplified (by sympy.simplify()) before being compiled. This will slow down the compilation speed.

If fastmath is True, the fastmath flag will be passed to the Numba JIT compiler. See Numba and LLVM documentation for details.

If cache is a path to a file, the FastFunc object will do the following: 1. If the file does not exist, the generated functions will be written to the file so they can be loaded later. 2. If the file exists and there is a hash string at the beginning of the file (auto-generated by FastFunc), the hash will be compared with the hash of the current function to determine whether to load the file directly or overwrite it. 3. If the file exists and there is no hash string at the beginning of the file, the file is considered a user-provided file and will be loaded directly.

Arguments:
  • function: sympy.Expr of the function.
  • args: sympy.Symbol objects used as the function arguments.
  • simplify: Whether to use sympy.simplify() to simplify expressions before compilation.
  • fastmath: Whether to use Numba fastmath mode.
  • cache: Path to a file to cache the generated functions.
F: Callable[[numpy.ndarray[Any, numpy.dtype[numpy.float64]], numpy.int32], numpy.ndarray[Any, numpy.dtype[numpy.float64]]]

Vectorized function to compute value.

G: Callable[[numpy.ndarray[Any, numpy.dtype[numpy.float64]], numpy.int32], numpy.ndarray[Any, numpy.dtype[numpy.float64]]]

Vectorized function to compute gradient.

H: Callable[[numpy.ndarray[Any, numpy.dtype[numpy.float64]], numpy.int32], numpy.ndarray[Any, numpy.dtype[numpy.float64]]]

Vectorized function to compute the Hessian.

G_index: numpy.ndarray[typing.Any, numpy.dtype[numpy.int32]]

Indices of nonzero gradient elements.

H_index_row: numpy.ndarray[typing.Any, numpy.dtype[numpy.int32]]

Row indices of nonzero elements in the lower triangle of the Hessian.

H_index_col: numpy.ndarray[typing.Any, numpy.dtype[numpy.int32]]

Column indices of nonzero elements in the lower triangle of the Hessian.

def ensure_directory(path: str | pathlib.Path) -> None:

Ensure that a path exists and is a directory.

Arguments:
  • path: Directory to create if it does not already exist.
Raises:
  • NotADirectoryError: If path exists and is not a directory.