feat!: implement faster lambdification

[redeboer]

Closes #273

Implementation of TR-002. See main result of this PR here:
https://tensorwaves--281.org.readthedocs.build/en/281/usage/faster-lambdify.html#specifying-complexity

[codecov]

Codecov Report

Merging #281 (6e1a860) into main (9be08e4) will decrease coverage by 3.19%.
The diff coverage is 29.78%.

@@            Coverage Diff             @@
##             main     #281      +/-   ##
==========================================
- Coverage   84.36%   81.17%   -3.20%     
==========================================
  Files          13       13              
  Lines         761      802      +41     
==========================================
+ Hits          642      651       +9     
- Misses        119      151      +32     

Flag	Coverage Δ
unittests	81.17% <29.78%> (-3.20%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
src/tensorwaves/model.py	71.72% <29.78%> (-13.61%)	⬇️

[redeboer]

Also gave multithreading a try. Unfortunately, lambdified expressions can't be pickled, so multiprocessing won't work (also not with pathos/dill).

Original snippet

tensorwaves/src/tensorwaves/model.py

Lines 106 to 139 in a409726

	def optimized_lambdify(
	args: Sequence[sp.Symbol],
	expression: sp.Expr,
	modules: Optional[Union[str, tuple, dict]] = None,
	*,
	min_complexity: int = 0,
	max_complexity: int,
	) -> Callable:
	"""Speed up `~sympy.utilities.lambdify.lambdify` with `.split_expression`.
	.. seealso:: :doc:`/usage/faster-lambdify`
	"""
	top_expression, definitions = split_expression(
	expression,
	min_complexity=min_complexity,
	max_complexity=max_complexity,
	)
	top_symbols = sorted(definitions, key=lambda s: s.name)
	top_lambdified = sp.lambdify(top_symbols, top_expression, modules)
	sub_lambdified = [ # same order as positional arguments in top_lambdified
	sp.lambdify(args, definitions[symbol], modules)
	for symbol in tqdm(
	iterable=top_symbols,
	desc="Lambdifying sub-expressions",
	unit="expr",
	disable=logging.getLogger().level > logging.WARNING,
	)
	]
	def recombined_function(*args): # type: ignore
	new_args = [sub_expr(*args) for sub_expr in sub_lambdified]
	return top_lambdified(*new_args)
	return recombined_function

Rewritten with multi-threading

# import multiprocessing
from pathos import multiprocessing

...

def optimized_lambdify(
    args: Sequence[sp.Symbol],
    expression: sp.Expr,
    modules: Optional[Union[str, tuple, dict]] = None,
    *,
    min_complexity: int = 0,
    max_complexity: int,
    n_threads: int = 1,
) -> Callable:
    """Speed up `~sympy.utilities.lambdify.lambdify` with `.split_expression`.

    .. seealso:: :doc:`/usage/faster-lambdify`
    """
    top_expression, definitions = split_expression(
        expression,
        min_complexity=min_complexity,
        max_complexity=max_complexity,
    )
    top_symbols = sorted(definitions, key=lambda s: s.name)
    top_lambdified = sp.lambdify(top_symbols, top_expression, modules)

    progress_bar = tqdm(
        total=len(top_symbols),
        desc="Lambdifying sub-expressions",
        unit="expr",
        disable=logging.getLogger().level > logging.WARNING,
    )

    def local_lambdify(symbol: sp.Symbol) -> Tuple[sp.Symbol, Callable]:
        lambdified_expression = sp.lambdify(args, definitions[symbol], modules)
        progress_bar.update()
        return symbol, lambdified_expression

    sub_expr_mapping: Dict[sp.Symbol, Callable] = {}
    if n_threads > 1:
        with multiprocessing.Pool(n_threads) as pool:
            for symbol, lambdified_sub_expr in pool.imap_unordered(
                local_lambdify, top_symbols
            ):
                sub_expr_mapping[symbol] = lambdified_sub_expr
    else:
        for symbol in top_symbols:
            _, lambdified_sub_expr = local_lambdify(symbol)
            sub_expr_mapping[symbol] = lambdified_sub_expr

    # retrieve order of positional arguments in top_lambdified
    sub_lambdified = [sub_expr_mapping[s] for s in top_symbols]

    def recombined_function(*args):  # type: ignore
        new_args = [sub_expr(*args) for sub_expr in sub_lambdified]
        return top_lambdified(*new_args)

    return recombined_function