From 933b643b62414c089417899448a17e2bffad5654 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Sun, 8 Feb 2026 22:13:02 +0300 Subject: [PATCH 01/11] libsympy.py added torchify function --- src/libsympy.py | 319 +++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 315 insertions(+), 4 deletions(-) diff --git a/src/libsympy.py b/src/libsympy.py index f1ec648..870dffd 100644 --- a/src/libsympy.py +++ b/src/libsympy.py @@ -47,8 +47,22 @@ from sympy.parsing.latex import parse_latex from sympy.parsing.sympy_parser import parse_expr from sympy.interactive import printing +import torch printing.init_printing() +from loguru import logger +logger.remove() + +import torch +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +print('torch:', torch.__version__, '| cuda:', torch.version.cuda, '| cuda available:', torch.cuda.is_available(), '| device:', device) +print('torch file:', getattr(torch, '__file__', None)) + +import numpy as np +import matplotlib.pyplot as plt +from torchquad import Simpson, set_up_backend + +set_up_backend("torch", data_type="float64") # Sets global defaults for all plots plt.rcParams.update({ @@ -297,7 +311,7 @@ def get_iterated_functions(f, fixed_vals={C1:0, C2:0}, prm=alpha, funcs = list(map(ifunc, param_vals)) return(funcs) -def get_piecewise(): +def get_piecewise(fV, xs): """ todo Returns a piecewise function of the function. @@ -318,7 +332,6 @@ def getfV(self, fV, xs): return(res) """ xintervals = [] - [fV, xs] = [self.V, self.xs] for i in range(len(xs)): if i == 0: xintervals.append((fV[i], x < xs[i])) @@ -406,13 +419,311 @@ def tensor_to_product(expr): expr = TensorProduct(a, b) + TensorProduct(a, b, c) converted = tensor_to_product(expr) """ + from sympy.physics.quantum import TensorProduct as _TensorProduct return expr.replace( - lambda x: isinstance(x, TensorProduct), + lambda x: isinstance(x, _TensorProduct), lambda x: Mul(*x.args) ) +def torchify(expr, + variables, + method=None, + dim=None, + N=21, + *limits, + params=None, + params_values=None, + eps=1e-8, + modules=None, + verbose=False): + """Torch + torchquad numeric integration helper for SymPy expressions. + + This function: + - lambdifies a SymPy expression into Torch ops + - integrates its real/imag parts using torchquad + - supports finite, infinite, and semi-infinite limits via change of variables + Parameters + ---------- + expr : sympy.Expr + Expression to integrate. May be complex. + variables : sequence[sympy.Symbol] + Integration variables (order matters). + method : torchquad integrator instance or None + Defaults to Simpson(). + dim : int or None + Integration dimension. Defaults to len(variables). + N : int + Integrator resolution (torchquad's N parameter). + limits : tuples + Integration limits as (var, lower, upper). You may pass either: + - multiple tuples: torchify(expr, [x,y], ..., (x,-oo,oo), (y,0,1)) + - a single list of tuples: torchify(expr, [x,y], ..., [(x,-oo,oo),(y,0,1)]) + params : sequence[sympy.Symbol] or None + Extra symbolic parameters (not integrated over) that appear in expr. + If None, inferred from expr.free_symbols excluding `variables`. + params_values : dict or sequence or None + Numerical values for params. If dict, keys are symbols. + If params is non-empty and params_values is None, raises ValueError. + eps : float + Small cutoff used when mapping infinite limits to (-pi/2, pi/2). + modules : list[dict] or None + Custom lambdify modules. If None, uses a reasonable torch mapping. + verbose : bool + Print basic debug info. + + Returns + ------- + (re, im) : torch.Tensor, torch.Tensor + Real and imaginary parts of the integral. + """ -#----Plotting + if method is None: + method = Simpson() + + import sympy as _sp + + def _is_limit_tuple(obj): + return ( + isinstance(obj, (list, tuple)) + and len(obj) == 3 + and getattr(obj[0], "is_Symbol", False) + ) + + def _normalize_limits(lims): + # Accept either a flat list/tuple of (v,a,b) or a single nested list. + if len(lims) == 1 and isinstance(lims[0], (list, tuple)): + if len(lims[0]) == 0: + return tuple() + if isinstance(lims[0][0], (list, tuple)) and len(lims[0][0]) == 3: + return tuple(lims[0]) + return tuple(lims) + + # --- Backward/alternate call styles --- + # 1) SymPy-like: torchify(expr, (x,a,b), (y,c,d), ...) + # In our signature, that arrives as variables=(x,a,b), method=(y,c,d). + inferred_limits = None + provided_symbols = None + + if _is_limit_tuple(variables): + # shift positional args: treat variables/method/(dim if also a limit tuple) as limits + pos_limits = [variables] + if _is_limit_tuple(method): + pos_limits.append(method) + method = Simpson() + if _is_limit_tuple(dim): + pos_limits.append(dim) + dim = None + if _is_limit_tuple(N): + pos_limits.append(N) + N = 21 + pos_limits.extend(list(limits)) + inferred_limits = tuple(pos_limits) + variables = [lim[0] for lim in inferred_limits] + limits = tuple() + else: + # Capture if user provided a list of symbols. + if isinstance(variables, (list, tuple)) and all(getattr(v, "is_Symbol", False) for v in variables): + provided_symbols = list(variables) + + # 2) Mixed positional: torchify(expr, [x,y], (x,a,b), (y,c,d)) + # In our signature, (x,a,b) is method and (y,c,d) is dim. + pos_limits = [] + if _is_limit_tuple(method): + pos_limits.append(method) + method = Simpson() + if _is_limit_tuple(dim): + pos_limits.append(dim) + dim = None + if _is_limit_tuple(N): + pos_limits.append(N) + N = 21 + if pos_limits: + inferred_limits = tuple(pos_limits) + tuple(limits) + limits = tuple() + + # Normalize the remaining *limits input (nested list supported) + limits = _normalize_limits(limits) + if inferred_limits is not None: + inferred_limits = _normalize_limits(inferred_limits) + if limits: + inferred_limits = tuple(inferred_limits) + tuple(limits) + limits = inferred_limits + + # If a single (var, lower, upper) tuple slipped through, wrap it. + if _is_limit_tuple(limits): + limits = (limits,) + + # If limits were provided, infer integration variables from them. + if limits: + invalid_limits = [lim for lim in limits if not _is_limit_tuple(lim)] + if invalid_limits: + raise TypeError(f"limits must be (var, lower, upper) tuples; got: {limits}") + variables_from_limits = [lim[0] for lim in limits] + # If caller gave a superset of symbols, interpret extras as params (when not explicitly provided) + if provided_symbols is not None and set(variables_from_limits).issubset(set(provided_symbols)): + if params is None: + extras = [s for s in provided_symbols if s not in set(variables_from_limits)] + params = [s for s in extras if s in expr.free_symbols] + variables = variables_from_limits + + if variables is None: + raise ValueError("variables must be provided (or inferable from limits)") + variables = list(variables) + bad_vars = [v for v in variables if not getattr(v, "is_Symbol", False)] + if bad_vars: + raise TypeError(f"variables must contain only SymPy Symbols; got: {bad_vars}") + + if dim is None: + dim = len(variables) + if dim != len(variables): + raise ValueError(f"dim={dim} must match len(variables)={len(variables)}") + + if not limits: + raise ValueError("limits are required: provide (var, lower, upper) for each variable") + + if len(limits) != len(variables): + raise ValueError("Number of limits must match number of variables.") + + limit_map = {lim[0]: (lim[1], lim[2]) for lim in limits} + for v in variables: + if v not in limit_map: + raise ValueError(f"No limits provided for variable {v}.") + + # --- Change of variables for infinite / semi-infinite limits --- + # We always map to finite domains because torchquad needs finite boxes. + # For each integrated variable v, we may replace it with a new symbol t_v. + new_vars = [] + domain = [] + subs_map = {} + jacobian = Integer(1) + + def _is_pos_inf(x): + return x == oo + + def _is_neg_inf(x): + return x == -oo + + for v in variables: + lower, upper = limit_map[v] + + if (_is_neg_inf(lower) and _is_pos_inf(upper)): + # (-inf, inf): v = tan(t) + t_v = Symbol(f"t_{v.name}", real=True) + subs_map[v] = tan(t_v) + jacobian *= 1 / cos(t_v) ** 2 + new_vars.append(t_v) + domain.append([float(-pi / 2 + eps), float(pi / 2 - eps)]) + elif (not _is_neg_inf(lower) and _is_pos_inf(upper)): + # (a, inf): v = a + tan(t)^2 + t_v = Symbol(f"t_{v.name}", real=True) + subs_map[v] = lower + tan(t_v) ** 2 + jacobian *= 2 * tan(t_v) / cos(t_v) ** 2 + new_vars.append(t_v) + domain.append([0.0, float(pi / 2 - eps)]) + elif (_is_neg_inf(lower) and not _is_pos_inf(upper)): + # (-inf, b): v = b - tan(t)^2 + t_v = Symbol(f"t_{v.name}", real=True) + subs_map[v] = upper - tan(t_v) ** 2 + jacobian *= -2 * tan(t_v) / cos(t_v) ** 2 + new_vars.append(t_v) + domain.append([0.0, float(pi / 2 - eps)]) + else: + # finite (a, b) + new_vars.append(v) + domain.append([float(lower), float(upper)]) + + expr_work = (expr.subs(subs_map) * jacobian) + + # Infer params if not provided. + if params is None: + params = sorted(list(expr_work.free_symbols - set(new_vars)), key=lambda s: s.name) + else: + params = list(params) + + if params: + if params_values is None: + raise ValueError("params_values must be provided when params are present") + if isinstance(params_values, dict): + param_vals = [params_values[p] for p in params] + else: + param_vals = list(params_values) + if len(param_vals) != len(params): + raise ValueError("params_values length must match params length") + else: + param_vals = [] + + if verbose: + print(f"torchify: dim={dim}, N={N}") + print(f" integrated vars: {variables} -> {new_vars}") + print(f" params: {params}") + + # Default torch-aware lambdify modules + if modules is None: + modules = [{ + # Trigonometric + "sin": torch.sin, + "cos": torch.cos, + "tan": torch.tan, + "asin": torch.asin, + "acos": torch.acos, + "atan": torch.atan, + "atan2": torch.atan2, + + # Hyperbolic + "sinh": torch.sinh, + "cosh": torch.cosh, + "tanh": torch.tanh, + "asinh": torch.asinh, + "acosh": torch.acosh, + "atanh": torch.atanh, + + # Exponentials / logs + "exp": torch.exp, + "log": torch.log, + "ln": torch.log, + "log10": torch.log10, + "log2": torch.log2, + + # Roots / powers + "sqrt": torch.sqrt, + "Pow": torch.pow, + + # Misc + "Abs": torch.abs, + "sign": torch.sign, + "floor": torch.floor, + "ceiling": torch.ceil, + "Min": torch.minimum, + "Max": torch.maximum, + + # Piecewise / heaviside + "Heaviside": lambda x: torch.heaviside(x, torch.zeros_like(x)), + + # Constants + "pi": getattr(torch, "pi", float(np.pi)), + "E": float(np.e), + }] + + # Lambdify expects a flat argument list. + arglist = tuple(new_vars + params) + expr_torch = lambdify(arglist, expr_work, modules=modules) + + def f_re(d): + vals = expr_torch(*[d[:, i] for i in range(d.shape[1])], *param_vals) + vals = torch.as_tensor(vals, device=d.device) + return vals.real if torch.is_complex(vals) else vals + + def f_im(d): + vals = expr_torch(*[d[:, i] for i in range(d.shape[1])], *param_vals) + vals = torch.as_tensor(vals, device=d.device) + if torch.is_complex(vals): + return vals.imag + return torch.zeros_like(vals) + + re = method.integrate(f_re, dim=dim, N=N, integration_domain=domain) + im = method.integrate(f_im, dim=dim, N=N, integration_domain=domain) + return re, im + #----Plotting def plot_energy_levels(data, constYs=None, title="", line_width=0.025): """ Plots horizontal lines at y-points grouped by x-coordinates. From e1973ed0bdd928549226469f4d67850c650c8b6a Mon Sep 17 00:00:00 2001 From: ibeuler Date: Mon, 9 Feb 2026 19:55:55 +0300 Subject: [PATCH 02/11] gitignore --- .gitignore | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..05d3a8d --- /dev/null +++ b/.gitignore @@ -0,0 +1,29 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# Virtual environments +.venv/ +venv/ +env/ + +# Distribution / packaging +build/ +dist/ +*.egg-info/ +*.egg + +# Jupyter Notebook checkpoints +.ipynb_checkpoints/ + +# OS files +Thumbs.db +.DS_Store + +# IDE +.vscode/ +.idea/ + +# Logs +*.log From d0d63d76e920e8d6dc8ad8a8a1536d6b5d256451 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Mon, 9 Feb 2026 19:56:15 +0300 Subject: [PATCH 03/11] libtorch --- src/libsympy.py | 299 ------------------------------------------------ src/libtorch.py | 226 ++++++++++++++++++++++++++++++++++++ 2 files changed, 226 insertions(+), 299 deletions(-) create mode 100644 src/libtorch.py diff --git a/src/libsympy.py b/src/libsympy.py index 870dffd..f9332cc 100644 --- a/src/libsympy.py +++ b/src/libsympy.py @@ -424,305 +424,6 @@ def tensor_to_product(expr): lambda x: isinstance(x, _TensorProduct), lambda x: Mul(*x.args) ) -def torchify(expr, - variables, - method=None, - dim=None, - N=21, - *limits, - params=None, - params_values=None, - eps=1e-8, - modules=None, - verbose=False): - """Torch + torchquad numeric integration helper for SymPy expressions. - - This function: - - lambdifies a SymPy expression into Torch ops - - integrates its real/imag parts using torchquad - - supports finite, infinite, and semi-infinite limits via change of variables - - Parameters - ---------- - expr : sympy.Expr - Expression to integrate. May be complex. - variables : sequence[sympy.Symbol] - Integration variables (order matters). - method : torchquad integrator instance or None - Defaults to Simpson(). - dim : int or None - Integration dimension. Defaults to len(variables). - N : int - Integrator resolution (torchquad's N parameter). - limits : tuples - Integration limits as (var, lower, upper). You may pass either: - - multiple tuples: torchify(expr, [x,y], ..., (x,-oo,oo), (y,0,1)) - - a single list of tuples: torchify(expr, [x,y], ..., [(x,-oo,oo),(y,0,1)]) - params : sequence[sympy.Symbol] or None - Extra symbolic parameters (not integrated over) that appear in expr. - If None, inferred from expr.free_symbols excluding `variables`. - params_values : dict or sequence or None - Numerical values for params. If dict, keys are symbols. - If params is non-empty and params_values is None, raises ValueError. - eps : float - Small cutoff used when mapping infinite limits to (-pi/2, pi/2). - modules : list[dict] or None - Custom lambdify modules. If None, uses a reasonable torch mapping. - verbose : bool - Print basic debug info. - - Returns - ------- - (re, im) : torch.Tensor, torch.Tensor - Real and imaginary parts of the integral. - """ - - if method is None: - method = Simpson() - - import sympy as _sp - - def _is_limit_tuple(obj): - return ( - isinstance(obj, (list, tuple)) - and len(obj) == 3 - and getattr(obj[0], "is_Symbol", False) - ) - - def _normalize_limits(lims): - # Accept either a flat list/tuple of (v,a,b) or a single nested list. - if len(lims) == 1 and isinstance(lims[0], (list, tuple)): - if len(lims[0]) == 0: - return tuple() - if isinstance(lims[0][0], (list, tuple)) and len(lims[0][0]) == 3: - return tuple(lims[0]) - return tuple(lims) - - # --- Backward/alternate call styles --- - # 1) SymPy-like: torchify(expr, (x,a,b), (y,c,d), ...) - # In our signature, that arrives as variables=(x,a,b), method=(y,c,d). - inferred_limits = None - provided_symbols = None - - if _is_limit_tuple(variables): - # shift positional args: treat variables/method/(dim if also a limit tuple) as limits - pos_limits = [variables] - if _is_limit_tuple(method): - pos_limits.append(method) - method = Simpson() - if _is_limit_tuple(dim): - pos_limits.append(dim) - dim = None - if _is_limit_tuple(N): - pos_limits.append(N) - N = 21 - pos_limits.extend(list(limits)) - inferred_limits = tuple(pos_limits) - variables = [lim[0] for lim in inferred_limits] - limits = tuple() - else: - # Capture if user provided a list of symbols. - if isinstance(variables, (list, tuple)) and all(getattr(v, "is_Symbol", False) for v in variables): - provided_symbols = list(variables) - - # 2) Mixed positional: torchify(expr, [x,y], (x,a,b), (y,c,d)) - # In our signature, (x,a,b) is method and (y,c,d) is dim. - pos_limits = [] - if _is_limit_tuple(method): - pos_limits.append(method) - method = Simpson() - if _is_limit_tuple(dim): - pos_limits.append(dim) - dim = None - if _is_limit_tuple(N): - pos_limits.append(N) - N = 21 - if pos_limits: - inferred_limits = tuple(pos_limits) + tuple(limits) - limits = tuple() - - # Normalize the remaining *limits input (nested list supported) - limits = _normalize_limits(limits) - if inferred_limits is not None: - inferred_limits = _normalize_limits(inferred_limits) - if limits: - inferred_limits = tuple(inferred_limits) + tuple(limits) - limits = inferred_limits - - # If a single (var, lower, upper) tuple slipped through, wrap it. - if _is_limit_tuple(limits): - limits = (limits,) - - # If limits were provided, infer integration variables from them. - if limits: - invalid_limits = [lim for lim in limits if not _is_limit_tuple(lim)] - if invalid_limits: - raise TypeError(f"limits must be (var, lower, upper) tuples; got: {limits}") - variables_from_limits = [lim[0] for lim in limits] - # If caller gave a superset of symbols, interpret extras as params (when not explicitly provided) - if provided_symbols is not None and set(variables_from_limits).issubset(set(provided_symbols)): - if params is None: - extras = [s for s in provided_symbols if s not in set(variables_from_limits)] - params = [s for s in extras if s in expr.free_symbols] - variables = variables_from_limits - - if variables is None: - raise ValueError("variables must be provided (or inferable from limits)") - variables = list(variables) - bad_vars = [v for v in variables if not getattr(v, "is_Symbol", False)] - if bad_vars: - raise TypeError(f"variables must contain only SymPy Symbols; got: {bad_vars}") - - if dim is None: - dim = len(variables) - if dim != len(variables): - raise ValueError(f"dim={dim} must match len(variables)={len(variables)}") - - if not limits: - raise ValueError("limits are required: provide (var, lower, upper) for each variable") - - if len(limits) != len(variables): - raise ValueError("Number of limits must match number of variables.") - - limit_map = {lim[0]: (lim[1], lim[2]) for lim in limits} - for v in variables: - if v not in limit_map: - raise ValueError(f"No limits provided for variable {v}.") - - # --- Change of variables for infinite / semi-infinite limits --- - # We always map to finite domains because torchquad needs finite boxes. - # For each integrated variable v, we may replace it with a new symbol t_v. - new_vars = [] - domain = [] - subs_map = {} - jacobian = Integer(1) - - def _is_pos_inf(x): - return x == oo - - def _is_neg_inf(x): - return x == -oo - - for v in variables: - lower, upper = limit_map[v] - - if (_is_neg_inf(lower) and _is_pos_inf(upper)): - # (-inf, inf): v = tan(t) - t_v = Symbol(f"t_{v.name}", real=True) - subs_map[v] = tan(t_v) - jacobian *= 1 / cos(t_v) ** 2 - new_vars.append(t_v) - domain.append([float(-pi / 2 + eps), float(pi / 2 - eps)]) - elif (not _is_neg_inf(lower) and _is_pos_inf(upper)): - # (a, inf): v = a + tan(t)^2 - t_v = Symbol(f"t_{v.name}", real=True) - subs_map[v] = lower + tan(t_v) ** 2 - jacobian *= 2 * tan(t_v) / cos(t_v) ** 2 - new_vars.append(t_v) - domain.append([0.0, float(pi / 2 - eps)]) - elif (_is_neg_inf(lower) and not _is_pos_inf(upper)): - # (-inf, b): v = b - tan(t)^2 - t_v = Symbol(f"t_{v.name}", real=True) - subs_map[v] = upper - tan(t_v) ** 2 - jacobian *= -2 * tan(t_v) / cos(t_v) ** 2 - new_vars.append(t_v) - domain.append([0.0, float(pi / 2 - eps)]) - else: - # finite (a, b) - new_vars.append(v) - domain.append([float(lower), float(upper)]) - - expr_work = (expr.subs(subs_map) * jacobian) - - # Infer params if not provided. - if params is None: - params = sorted(list(expr_work.free_symbols - set(new_vars)), key=lambda s: s.name) - else: - params = list(params) - - if params: - if params_values is None: - raise ValueError("params_values must be provided when params are present") - if isinstance(params_values, dict): - param_vals = [params_values[p] for p in params] - else: - param_vals = list(params_values) - if len(param_vals) != len(params): - raise ValueError("params_values length must match params length") - else: - param_vals = [] - - if verbose: - print(f"torchify: dim={dim}, N={N}") - print(f" integrated vars: {variables} -> {new_vars}") - print(f" params: {params}") - - # Default torch-aware lambdify modules - if modules is None: - modules = [{ - # Trigonometric - "sin": torch.sin, - "cos": torch.cos, - "tan": torch.tan, - "asin": torch.asin, - "acos": torch.acos, - "atan": torch.atan, - "atan2": torch.atan2, - - # Hyperbolic - "sinh": torch.sinh, - "cosh": torch.cosh, - "tanh": torch.tanh, - "asinh": torch.asinh, - "acosh": torch.acosh, - "atanh": torch.atanh, - - # Exponentials / logs - "exp": torch.exp, - "log": torch.log, - "ln": torch.log, - "log10": torch.log10, - "log2": torch.log2, - - # Roots / powers - "sqrt": torch.sqrt, - "Pow": torch.pow, - - # Misc - "Abs": torch.abs, - "sign": torch.sign, - "floor": torch.floor, - "ceiling": torch.ceil, - "Min": torch.minimum, - "Max": torch.maximum, - - # Piecewise / heaviside - "Heaviside": lambda x: torch.heaviside(x, torch.zeros_like(x)), - - # Constants - "pi": getattr(torch, "pi", float(np.pi)), - "E": float(np.e), - }] - - # Lambdify expects a flat argument list. - arglist = tuple(new_vars + params) - expr_torch = lambdify(arglist, expr_work, modules=modules) - - def f_re(d): - vals = expr_torch(*[d[:, i] for i in range(d.shape[1])], *param_vals) - vals = torch.as_tensor(vals, device=d.device) - return vals.real if torch.is_complex(vals) else vals - - def f_im(d): - vals = expr_torch(*[d[:, i] for i in range(d.shape[1])], *param_vals) - vals = torch.as_tensor(vals, device=d.device) - if torch.is_complex(vals): - return vals.imag - return torch.zeros_like(vals) - - re = method.integrate(f_re, dim=dim, N=N, integration_domain=domain) - im = method.integrate(f_im, dim=dim, N=N, integration_domain=domain) - return re, im #----Plotting def plot_energy_levels(data, constYs=None, title="", line_width=0.025): """ diff --git a/src/libtorch.py b/src/libtorch.py new file mode 100644 index 0000000..890c30a --- /dev/null +++ b/src/libtorch.py @@ -0,0 +1,226 @@ +# -*- coding: utf-8 -*- +""" +libtorch.py — SymPy-to-Torch conversion and torchquad integration. + +Part of libphysics. + +Two-function design for performance: + 1. torchify() — expensive SymPy work done ONCE (change-of-vars + lambdify) + 2. torchquad_integrate() — cheap torch-only integration, called MANY times + +Usage +===== + from libphysics.libtorch import torchify, torchquad_integrate + + # Step 1: build once (slow: SymPy subs + lambdify, ~200-500 ms) + wigner = torchify( + integrand, + variables=[y_A, y_B], + limits=[(y_A, -oo, oo), (y_B, -oo, oo)], + params=[x_A, p_A, x_B, p_B], + ) + + # Step 2: integrate many times (fast: pure torch, ~3-8 ms each) + re, im = torchquad_integrate(wigner, params_values=[1, 1, 1, 1], N=121**2) +""" +from dataclasses import dataclass +from typing import List, Callable + +import numpy as np +import torch +from sympy import lambdify, Symbol, tan, cos, pi, oo, Integer + + +# --------------------------------------------------------------------------- +# Default torch module mapping for lambdify +# --------------------------------------------------------------------------- +def _default_torch_modules(): + return { + # Trigonometric + "sin": torch.sin, "cos": torch.cos, "tan": torch.tan, + "asin": torch.asin, "acos": torch.acos, "atan": torch.atan, + "atan2": torch.atan2, + # Hyperbolic + "sinh": torch.sinh, "cosh": torch.cosh, "tanh": torch.tanh, + "asinh": torch.asinh, "acosh": torch.acosh, "atanh": torch.atanh, + # Exponentials / logs + "exp": torch.exp, "log": torch.log, "ln": torch.log, + "log10": torch.log10, "log2": torch.log2, + # Roots / powers + "sqrt": torch.sqrt, "Pow": torch.pow, + # Misc + "Abs": torch.abs, "sign": torch.sign, + "floor": torch.floor, "ceiling": torch.ceil, + "Min": torch.minimum, "Max": torch.maximum, + # Piecewise / heaviside + "Heaviside": lambda x: torch.heaviside(x, torch.zeros_like(x)), + # Constants + "pi": getattr(torch, "pi", float(np.pi)), + "E": float(np.e), + } + + +# --------------------------------------------------------------------------- +# TorchExpr — container returned by torchify when limits are provided +# --------------------------------------------------------------------------- +@dataclass +class TorchExpr: + """Pre-built torch function + finite integration domain.""" + func: Callable # f(new_var_0, …, new_var_n, param_0, …, param_m) + domain: List[List[float]] # finite box for torchquad + dim: int # number of integration variables + n_params: int # number of extra parameters + + +# --------------------------------------------------------------------------- +# torchify — expensive work done ONCE +# --------------------------------------------------------------------------- +def torchify(expr, variables, limits=None, params=None, eps=1e-8, modules=None): + """ + Convert a SymPy expression into a torch-compatible function via lambdify. + + If *limits* are provided, performs change-of-variables for infinite / + semi-infinite limits at the **SymPy level** (symbolic, done once) and + multiplies by the analytic Jacobian. Returns a ``TorchExpr`` ready + for ``torchquad_integrate``. + + If *limits* are **not** provided, returns a plain callable (simple + lambdify with torch modules). + + Parameters + ---------- + expr : sympy.Expr + Symbolic expression (may be complex-valued). + variables : list[sympy.Symbol] + Integration variables (order matters). + limits : list[tuple] or None + ``(var, lower, upper)`` for each variable. ``sympy.oo`` / + ``-sympy.oo`` trigger automatic change of variables. + params : list[sympy.Symbol] or None + Extra symbolic parameters that appear in *expr* but are **not** + integrated over. If ``None``, inferred automatically. + eps : float + Small cutoff for mapping infinite limits (``±π/2 ∓ eps``). + modules : list[dict] or None + Custom lambdify module list. ``None`` → default torch mapping. + + Returns + ------- + TorchExpr if *limits* were given (use with ``torchquad_integrate``) + callable if *limits* were ``None`` (plain torch function) + """ + if modules is None: + modules = [_default_torch_modules()] + + variables = list(variables) + + # ------------------------------------------------------------------ + # Simple case: no limits → plain lambdify + # ------------------------------------------------------------------ + if limits is None: + return lambdify(tuple(variables), expr, modules=modules) + + # ------------------------------------------------------------------ + # With limits: symbolic change of variables (done once, fast forever) + # ------------------------------------------------------------------ + limit_map = {lim[0]: (lim[1], lim[2]) for lim in limits} + + new_vars = [] + domain = [] + subs_map = {} + jacobian = Integer(1) + + for v in variables: + lower, upper = limit_map[v] + + if lower == -oo and upper == oo: + # (-∞, ∞) → v = tan(t), dv = sec²(t) dt + t_v = Symbol(f"t_{v.name}", real=True) + subs_map[v] = tan(t_v) + jacobian *= 1 / cos(t_v) ** 2 + new_vars.append(t_v) + domain.append([float(-pi / 2 + eps), float(pi / 2 - eps)]) + + elif lower != -oo and upper == oo: + # (a, ∞) → v = a + tan²(t), dv = 2 tan(t) sec²(t) dt + t_v = Symbol(f"t_{v.name}", real=True) + subs_map[v] = lower + tan(t_v) ** 2 + jacobian *= 2 * tan(t_v) / cos(t_v) ** 2 + new_vars.append(t_v) + domain.append([0.0, float(pi / 2 - eps)]) + + elif lower == -oo and upper != oo: + # (-∞, b) → v = b - tan²(t), dv = -2 tan(t) sec²(t) dt + t_v = Symbol(f"t_{v.name}", real=True) + subs_map[v] = upper - tan(t_v) ** 2 + jacobian *= -2 * tan(t_v) / cos(t_v) ** 2 + new_vars.append(t_v) + domain.append([0.0, float(pi / 2 - eps)]) + + else: + # finite [a, b] + new_vars.append(v) + domain.append([float(lower), float(upper)]) + + # Substitute + bake in the Jacobian — all symbolic, done ONCE + expr_work = expr.subs(subs_map) * jacobian + + # Infer params from remaining free symbols + if params is None: + params = sorted( + list(expr_work.free_symbols - set(new_vars)), + key=lambda s: s.name, + ) + else: + params = list(params) + + # Lambdify: new_vars first, then params + arglist = tuple(new_vars + params) + func = lambdify(arglist, expr_work, modules=modules) + + return TorchExpr(func=func, domain=domain, dim=len(new_vars), n_params=len(params)) + + +# --------------------------------------------------------------------------- +# torchquad_integrate — cheap work, called MANY times +# --------------------------------------------------------------------------- +def torchquad_integrate(texpr, params_values=None, method=None, N=21): + """ + Numerically integrate a ``TorchExpr`` using torchquad. + + Parameters + ---------- + texpr : TorchExpr + Object returned by ``torchify(..., limits=...)``. + params_values : list or None + Numerical values for the parameters (same order as *params* in torchify). + method : torchquad integrator or None + Defaults to ``Simpson()``. + N : int + Integrator resolution (torchquad's *N* parameter). + + Returns + ------- + re, im : torch.Tensor + Real and imaginary parts of the integral. + """ + from torchquad import Simpson + + if method is None: + method = Simpson() + + param_vals = list(params_values) if params_values else [] + + def f_re(d): + vals = texpr.func(*[d[:, i] for i in range(d.shape[1])], *param_vals) + vals = torch.as_tensor(vals, device=d.device) + return vals.real if torch.is_complex(vals) else vals + + def f_im(d): + vals = texpr.func(*[d[:, i] for i in range(d.shape[1])], *param_vals) + vals = torch.as_tensor(vals, device=d.device) + return vals.imag if torch.is_complex(vals) else torch.zeros_like(vals) + + re = method.integrate(f_re, dim=texpr.dim, N=N, integration_domain=texpr.domain) + im = method.integrate(f_im, dim=texpr.dim, N=N, integration_domain=texpr.domain) + return re, im \ No newline at end of file From b6fbe61f48ab8b5e286194c40c67b4e84d98e544 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Mon, 9 Feb 2026 20:08:58 +0300 Subject: [PATCH 04/11] unnessesary imports removed --- src/libsympy.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/src/libsympy.py b/src/libsympy.py index f9332cc..3e3e9b6 100644 --- a/src/libsympy.py +++ b/src/libsympy.py @@ -47,7 +47,7 @@ from sympy.parsing.latex import parse_latex from sympy.parsing.sympy_parser import parse_expr from sympy.interactive import printing -import torch + printing.init_printing() from loguru import logger @@ -58,12 +58,6 @@ print('torch:', torch.__version__, '| cuda:', torch.version.cuda, '| cuda available:', torch.cuda.is_available(), '| device:', device) print('torch file:', getattr(torch, '__file__', None)) -import numpy as np -import matplotlib.pyplot as plt -from torchquad import Simpson, set_up_backend - -set_up_backend("torch", data_type="float64") - # Sets global defaults for all plots plt.rcParams.update({ 'font.family': "serif", From 4d29ada53739f9e440e1af2ad60d116e975e73e7 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Fri, 13 Feb 2026 17:14:23 +0300 Subject: [PATCH 05/11] updated dependencies --- requirements.txt | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/requirements.txt b/requirements.txt index 1854b59..24f017f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,18 @@ -ipython>=8.5.0 -jupytext>=1.16.2 -matplotlib>=3.6.1 -mpmath>=1.2.1 -numpy>=1.23.3 -scipy>=1.9.0 -sympy>=1.11.1 +ipython>=8.18.1 +jupytext>=1.16.4 +matplotlib>=3.8.2 +matplotlib-inline>=0.1.6 +mpmath>=1.3.0 +numpy>=1.26.4 +scipy>=1.11.4 +sympy>=1.12.1 + +# Used by sympy.parsing.latex.parse_latex (see libphysics/src/libsympy.py) +antlr4-python3-runtime==4.11 + +# Logging +loguru>=0.7.0 + +# Torch utilities (see libphysics/src/libsympy.py, libphysics/src/libtorch.py) +torch>=2.0.0 +torchquad>=0.5.0 From 0140e91de1fd7d99d6039d4e8610ccdb036fbc81 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Fri, 13 Feb 2026 17:15:05 +0300 Subject: [PATCH 06/11] removed unnessesary imports --- src/libsympy.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/libsympy.py b/src/libsympy.py index 3e3e9b6..03d37f2 100644 --- a/src/libsympy.py +++ b/src/libsympy.py @@ -50,14 +50,6 @@ printing.init_printing() -from loguru import logger -logger.remove() - -import torch -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') -print('torch:', torch.__version__, '| cuda:', torch.version.cuda, '| cuda available:', torch.cuda.is_available(), '| device:', device) -print('torch file:', getattr(torch, '__file__', None)) - # Sets global defaults for all plots plt.rcParams.update({ 'font.family': "serif", From 39ae591c826553e532dbb60314255a666bb01bc4 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Sat, 14 Feb 2026 03:17:19 +0300 Subject: [PATCH 07/11] torch_integrate_batched_simpson function --- src/libtorch.py | 254 ++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 246 insertions(+), 8 deletions(-) diff --git a/src/libtorch.py b/src/libtorch.py index 890c30a..0addb25 100644 --- a/src/libtorch.py +++ b/src/libtorch.py @@ -21,15 +21,30 @@ ) # Step 2: integrate many times (fast: pure torch, ~3-8 ms each) - re, im = torchquad_integrate(wigner, params_values=[1, 1, 1, 1], N=121**2) + re, im = torchquad_integrate(wigner, params_values=[1, 1, 1, 1], N=121) """ from dataclasses import dataclass +import sys from typing import List, Callable import numpy as np import torch from sympy import lambdify, Symbol, tan, cos, pi, oo, Integer +from loguru import logger + +# Remove loguru's default handler only if user requests it via env var. +# Set LIBPHYSICS_REMOVE_LOGURU=1 or "true"/"yes" to disable. +def setup_logging(enable: bool = False, level: str = "INFO"): + """ + By default, we remove all handlers (silence). + If enable=True, we add a standard output handler. + """ + logger.remove() + if enable: + logger.add(sys.stderr, level=level) + +setup_logging(enable=False) # --------------------------------------------------------------------------- # Default torch module mapping for lambdify @@ -75,7 +90,7 @@ class TorchExpr: # --------------------------------------------------------------------------- # torchify — expensive work done ONCE # --------------------------------------------------------------------------- -def torchify(expr, variables, limits=None, params=None, eps=1e-8, modules=None): +def torchify(expr, variables, limits=None, params=None, modules=None): """ Convert a SymPy expression into a torch-compatible function via lambdify. @@ -99,8 +114,6 @@ def torchify(expr, variables, limits=None, params=None, eps=1e-8, modules=None): params : list[sympy.Symbol] or None Extra symbolic parameters that appear in *expr* but are **not** integrated over. If ``None``, inferred automatically. - eps : float - Small cutoff for mapping infinite limits (``±π/2 ∓ eps``). modules : list[dict] or None Custom lambdify module list. ``None`` → default torch mapping. @@ -139,7 +152,7 @@ def torchify(expr, variables, limits=None, params=None, eps=1e-8, modules=None): subs_map[v] = tan(t_v) jacobian *= 1 / cos(t_v) ** 2 new_vars.append(t_v) - domain.append([float(-pi / 2 + eps), float(pi / 2 - eps)]) + domain.append([float(-pi / 2), float(pi / 2)]) elif lower != -oo and upper == oo: # (a, ∞) → v = a + tan²(t), dv = 2 tan(t) sec²(t) dt @@ -147,7 +160,7 @@ def torchify(expr, variables, limits=None, params=None, eps=1e-8, modules=None): subs_map[v] = lower + tan(t_v) ** 2 jacobian *= 2 * tan(t_v) / cos(t_v) ** 2 new_vars.append(t_v) - domain.append([0.0, float(pi / 2 - eps)]) + domain.append([0.0, float(pi / 2)]) elif lower == -oo and upper != oo: # (-∞, b) → v = b - tan²(t), dv = -2 tan(t) sec²(t) dt @@ -155,7 +168,7 @@ def torchify(expr, variables, limits=None, params=None, eps=1e-8, modules=None): subs_map[v] = upper - tan(t_v) ** 2 jacobian *= -2 * tan(t_v) / cos(t_v) ** 2 new_vars.append(t_v) - domain.append([0.0, float(pi / 2 - eps)]) + domain.append([0.0, float(pi / 2)]) else: # finite [a, b] @@ -223,4 +236,229 @@ def f_im(d): re = method.integrate(f_re, dim=texpr.dim, N=N, integration_domain=texpr.domain) im = method.integrate(f_im, dim=texpr.dim, N=N, integration_domain=texpr.domain) - return re, im \ No newline at end of file + return re, im + + +def _simpson_weights_1d(a: float, b: float, N: int, *, device=None, dtype=None) -> torch.Tensor: + """Return Simpson weights including the dx/3 scaling (shape: [N]).""" + if N < 3 or (N % 2) == 0: + raise ValueError(f"Simpson rule requires odd N>=3; got N={N}") + dx = (b - a) / (N - 1) + w = torch.ones(N, device=device, dtype=dtype) + # 1, 4, 2, 4, ..., 2, 4, 1 + w[1:-1:2] = 4 + w[2:-1:2] = 2 + w = w * (dx / 3.0) + return w + + +def _is_scalar_like(x) -> bool: + """True for Python numbers / 0-d tensors.""" + if torch.is_tensor(x): + return x.ndim == 0 + return isinstance(x, (int, float, complex, np.number)) + + +def _normalize_params_values(params_values, n_params: int, *, device, dtype): + """Normalize params into (B, n_params) plus batch_shape. + + Accepts: + - tensor of shape (..., n_params) + - list/tuple length n_params containing tensors and/or scalars + + Scalars are broadcast to the inferred batch_shape. + """ + if params_values is None: + if n_params != 0: + raise ValueError(f"Expected {n_params} params; got None") + return torch.empty((1, 0), device=device, dtype=dtype), tuple() + + # Case A: tensor (..., n_params) + if torch.is_tensor(params_values): + params_tensor = params_values + if params_tensor.shape[-1] != n_params: + raise ValueError( + f"params_values last dimension must be n_params={n_params}; got shape={tuple(params_tensor.shape)}" + ) + batch_shape = tuple(params_tensor.shape[:-1]) + B = int(np.prod(batch_shape)) if batch_shape else 1 + params_flat = params_tensor.reshape(B, n_params).to(device=device, dtype=dtype) + return params_flat, batch_shape + + # Case B: list/tuple of params (scalars and/or tensors) + if not isinstance(params_values, (list, tuple)): + # single scalar treated as 1 param + params_values = [params_values] + + if len(params_values) != n_params: + raise ValueError(f"Expected {n_params} params; got {len(params_values)}") + + # Infer batch_shape from the first non-scalar param (or scalars-only => ()) + batch_shape = None + for p in params_values: + if torch.is_tensor(p) and p.ndim > 0: + batch_shape = tuple(p.shape) + break + if not _is_scalar_like(p): + pt = torch.as_tensor(p) + if pt.ndim > 0: + batch_shape = tuple(pt.shape) + break + if batch_shape is None: + batch_shape = tuple() + + # Broadcast all params to batch_shape, then stack + params_list = [] + for p in params_values: + pt = torch.as_tensor(p, device=device) + if pt.ndim == 0: + if batch_shape: + pt = pt.expand(batch_shape) + else: + if tuple(pt.shape) != batch_shape: + raise ValueError( + f"All parameter tensors must have the same shape; got {tuple(pt.shape)} vs {batch_shape}" + ) + params_list.append(pt) + + params_tensor = torch.stack(params_list, dim=-1).to(dtype=dtype) + B = int(np.prod(batch_shape)) if batch_shape else 1 + params_flat = params_tensor.reshape(B, n_params) + return params_flat, batch_shape + + +def torch_integrate_batched_simpson( + texpr: TorchExpr, + params_values, + *, + N: int = 121, + chunk_size_params: int = 256, + chunk_size_points: int | None = None, + device=None, + dtype=None, +): + """Batched tensor-product Simpson integration for any dimension. + + This integrates a ``TorchExpr`` over its finite domain using a tensor-product + Simpson rule with *N* points per dimension. + + Supports a batch of parameter points (e.g. a 250x250 mesh) by keeping the + batch dimension and summing only over the sample grid. + + Notes + ----- + - The number of sample points grows as $N^{dim}$. For dim>3 this can become + very large quickly. Use smaller N and/or set ``chunk_size_points``. + - ``chunk_size_points`` trades more Python overhead for lower peak memory. + + Parameters + ---------- + texpr : TorchExpr + Output of ``torchify(..., limits=...)``. + params_values : tensor | list/tuple + Either a tensor of shape ``(..., n_params)`` or a list/tuple of length + ``n_params`` of tensors/scalars with broadcastable batch shapes. + N : int + Odd Simpson points per dimension. + chunk_size_params : int + Number of parameter points processed per chunk. + chunk_size_points : int | None + Optional number of sample points processed per chunk. + + Returns + ------- + re, im : torch.Tensor + Shapes match the parameter batch shape (e.g. 250x250). + """ + if texpr.dim <= 0: + raise ValueError(f"texpr.dim must be positive; got dim={texpr.dim}") + if texpr.n_params < 0: + raise ValueError(f"texpr.n_params must be non-negative; got n_params={texpr.n_params}") + if len(texpr.domain) != texpr.dim: + raise ValueError(f"texpr.domain must have length dim={texpr.dim}; got {len(texpr.domain)}") + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if dtype is None: + dtype = torch.float32 if device.type == "cuda" else torch.float64 + + if chunk_size_params <= 0: + raise ValueError(f"chunk_size_params must be positive; got {chunk_size_params}") + if chunk_size_points is not None and chunk_size_points <= 0: + raise ValueError(f"chunk_size_points must be positive; got {chunk_size_points}") + + params_flat, batch_shape = _normalize_params_values( + params_values, + texpr.n_params, + device=device, + dtype=dtype, + ) + B = int(params_flat.shape[0]) + + # Build per-dimension coordinates and weights + coords_1d = [] + weights_1d = [] + for (a, b) in texpr.domain: + a = float(a) + b = float(b) + coords_1d.append(torch.linspace(a, b, N, device=device, dtype=dtype)) + weights_1d.append(_simpson_weights_1d(a, b, N, device=device, dtype=dtype)) + + # Build full grid points (P, dim) and weights (P,) + # This is the most straightforward approach; chunk_size_points can reduce peak memory. + grid = torch.cartesian_prod(*coords_1d) # (P, dim) + w = weights_1d[0] + for wi in weights_1d[1:]: + w = torch.kron(w, wi) + P = int(grid.shape[0]) + + if chunk_size_points is None: + chunk_size_points = P + + re_out = torch.empty(B, device=device, dtype=dtype) + im_out = torch.empty(B, device=device, dtype=dtype) + + # Evaluate in chunks over params and points + for start_p in range(0, B, chunk_size_params): + stop_p = min(B, start_p + chunk_size_params) + pc = params_flat[start_p:stop_p, :] # (Bc, n_params) + param_args = [pc[:, i].unsqueeze(0) for i in range(texpr.n_params)] + + re_acc = torch.zeros(stop_p - start_p, device=device, dtype=dtype) + im_acc = torch.zeros(stop_p - start_p, device=device, dtype=dtype) + + for start_x in range(0, P, chunk_size_points): + stop_x = min(P, start_x + chunk_size_points) + g = grid[start_x:stop_x, :] # (Px, dim) + ww = w[start_x:stop_x].unsqueeze(1) # (Px, 1) + + var_args = [g[:, i].unsqueeze(1) for i in range(texpr.dim)] + vals = texpr.func(*var_args, *param_args) + vals = torch.as_tensor(vals, device=device) + + # Expect (Px, Bc) after broadcasting. + if vals.ndim == 1: + vals = vals.unsqueeze(1) + elif vals.ndim == 2 and vals.shape[0] == (stop_p - start_p) and vals.shape[1] == (stop_x - start_x): + # Sometimes lambdify might return (Bc, Px) + vals = vals.t().contiguous() + + if vals.shape[0] != (stop_x - start_x) or vals.shape[1] != (stop_p - start_p): + raise RuntimeError( + f"Unexpected integrand output shape {tuple(vals.shape)}; expected ({stop_x - start_x}, {stop_p - start_p})" + ) + + if torch.is_complex(vals): + re_acc += (ww * vals.real).sum(dim=0) + im_acc += (ww * vals.imag).sum(dim=0) + else: + re_acc += (ww * vals).sum(dim=0) + # im stays zero + + re_out[start_p:stop_p] = re_acc + im_out[start_p:stop_p] = im_acc + + re_out = re_out.reshape(batch_shape) + im_out = im_out.reshape(batch_shape) + return re_out, im_out + From a39f8c766938ffa64e3a5693dc4152033d94ec88 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Sat, 14 Feb 2026 03:18:27 +0300 Subject: [PATCH 08/11] test_libtorch added on a simple example with well known results for compareson --- tests/test_libtorch.ipynb | 214 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 214 insertions(+) create mode 100644 tests/test_libtorch.ipynb diff --git a/tests/test_libtorch.ipynb b/tests/test_libtorch.ipynb new file mode 100644 index 0000000..b0ac1b7 --- /dev/null +++ b/tests/test_libtorch.ipynb @@ -0,0 +1,214 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "7e9699ae", + "metadata": {}, + "source": [ + "# libtorch integration test\n", + "\n", + "This notebook validates that `libtorch.py` integrates SymPy expressions using Torch + torchquad.\n", + "\n", + "It is intentionally minimal and should run on CPU-only setups; CUDA is optional." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "7ce926ba", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "cwd: c:\\Users\\Pc\\Desktop\\LOCALS\\F_Nutku_Studies\\libphysics\\tests\n", + "src: C:\\Users\\Pc\\Desktop\\LOCALS\\F_Nutku_Studies\\libphysics\\src\n", + "torch: 2.5.1+cu121\n", + "cuda available: True\n", + "cuda: 12.1\n" + ] + } + ], + "source": [ + "\n", + "from __future__ import annotations\n", + "\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import torch\n", + "\n", + "\n", + "def _find_repo_root_with_src(start: Path) -> Path:\n", + " for p in [start, *start.parents]:\n", + " if (p / \"src\" / \"libtorch.py\").exists():\n", + " return p\n", + " raise FileNotFoundError(\"Could not find 'src/libtorch.py' above current working directory\")\n", + "\n", + "\n", + "ROOT = _find_repo_root_with_src(Path.cwd().resolve())\n", + "SRC = ROOT / \"src\"\n", + "if str(SRC) not in sys.path:\n", + " sys.path.insert(0, str(SRC))\n", + "\n", + "print(\"cwd:\", Path.cwd())\n", + "print(\"src:\", SRC)\n", + "print(\"torch:\", torch.__version__)\n", + "print(\"cuda available:\", torch.cuda.is_available())\n", + "if torch.cuda.is_available():\n", + " print(\"cuda:\", torch.version.cuda)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "0a58163c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "sympy: 1.13.1\n" + ] + } + ], + "source": [ + "import math\n", + "\n", + "import sympy as sp\n", + "\n", + "from libtorch import torchify, torchquad_integrate, setup_logging\n", + "\n", + "torch.set_default_dtype(torch.float64)\n", + "\n", + "\n", + "print(\"sympy:\", sp.__version__)" + ] + }, + { + "cell_type": "markdown", + "id": "3eb4b20c", + "metadata": {}, + "source": [ + "### Integral expressions for both tests\n", + "\n", + "Test 1 (complex Gaussian, over (-∞, ∞)) \n", + "Let x, k ∈ ℝ. \n", + "$$\n", + "\\int_{-\\infty}^{\\infty} e^{-x^{2}} e^{i k x}\\,dx = \\sqrt{\\pi}\\, e^{-k^{2}/4}.\n", + "$$\n", + "\n", + "Test 2 (semi-infinite exponential, over (0, ∞)) \n", + "Let x ∈ ℝ. \n", + "$$\n", + "\\int_{0}^{\\infty} e^{-x}\\,dx = 1.\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "17510ebc", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "re: 1.380388447018694 expected: 1.380388447043143 abs err: 2.4448887359085347e-11\n", + "im: 7.806255641895632e-18\n" + ] + } + ], + "source": [ + "# Test 1: complex Gaussian integral over (-oo, oo)\n", + "# ∫ exp(-x^2) * exp(I*k*x) dx = sqrt(pi) * exp(-k^2/4)\n", + "\n", + "x, k = sp.symbols(\"x k\", real=True)\n", + "expr = sp.exp(-x**2) * sp.exp(sp.I * k * x)\n", + "\n", + "texpr = torchify(\n", + " expr,\n", + " variables=[x],\n", + " limits=[(x, -sp.oo, sp.oo)],\n", + " params=[k],\n", + ")\n", + "\n", + "re, im = torchquad_integrate(texpr, params_values=[1.0], N=151)\n", + "\n", + "expected = math.sqrt(math.pi) * math.exp(-(1.0**2) / 4.0)\n", + "re_f = float(re.detach().cpu())\n", + "im_f = float(im.detach().cpu())\n", + "\n", + "print(\"re:\", re_f, \"expected:\", expected, \"abs err:\", abs(re_f - expected))\n", + "print(\"im:\", im_f)\n", + "\n", + "assert abs(re_f - expected) < 5e-3\n", + "assert abs(im_f) < 5e-3" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "26b67910", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "re: 0.9999999997326827 expected: 1.0 abs err: 2.6731727942319594e-10\n", + "im: 0.0\n" + ] + } + ], + "source": [ + "# Test 2: semi-infinite integral over (0, oo)\n", + "# ∫_0^∞ exp(-x) dx = 1\n", + "\n", + "x = sp.Symbol(\"x\", real=True)\n", + "expr = sp.exp(-x)\n", + "\n", + "texpr = torchify(\n", + " expr,\n", + " variables=[x],\n", + " limits=[(x, 0, sp.oo)],\n", + " params=[],\n", + ")\n", + "\n", + "re, im = torchquad_integrate(texpr, params_values=[], N=151)\n", + "re_f = float(re.detach().cpu())\n", + "im_f = float(im.detach().cpu())\n", + "\n", + "print(\"re:\", re_f, \"expected:\", 1.0, \"abs err:\", abs(re_f - 1.0))\n", + "print(\"im:\", im_f)\n", + "\n", + "assert abs(re_f - 1.0) < 5e-3\n", + "assert abs(im_f) < 5e-3" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 7680dcd4e04976fdd800db85c341f0b84342ef93 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Wed, 4 Mar 2026 01:34:00 +0300 Subject: [PATCH 09/11] libtorch updated 04.03.2026 --- src/libtorch.py | 133 +++++++++++++++++++++++++++++++----------------- 1 file changed, 86 insertions(+), 47 deletions(-) diff --git a/src/libtorch.py b/src/libtorch.py index 0addb25..6147db6 100644 --- a/src/libtorch.py +++ b/src/libtorch.py @@ -39,6 +39,14 @@ def setup_logging(enable: bool = False, level: str = "INFO"): """ By default, we remove all handlers (silence). If enable=True, we add a standard output handler. + ------ + Example: + # Disable logging (default) + setup_logging() + # Enable logging to stderr at INFO level + setup_logging(enable=True, level="INFO") + # Enable logging to stderr at DEBUG level + setup_logging(enable=True, level="DEBUG") """ logger.remove() if enable: @@ -46,35 +54,6 @@ def setup_logging(enable: bool = False, level: str = "INFO"): setup_logging(enable=False) -# --------------------------------------------------------------------------- -# Default torch module mapping for lambdify -# --------------------------------------------------------------------------- -def _default_torch_modules(): - return { - # Trigonometric - "sin": torch.sin, "cos": torch.cos, "tan": torch.tan, - "asin": torch.asin, "acos": torch.acos, "atan": torch.atan, - "atan2": torch.atan2, - # Hyperbolic - "sinh": torch.sinh, "cosh": torch.cosh, "tanh": torch.tanh, - "asinh": torch.asinh, "acosh": torch.acosh, "atanh": torch.atanh, - # Exponentials / logs - "exp": torch.exp, "log": torch.log, "ln": torch.log, - "log10": torch.log10, "log2": torch.log2, - # Roots / powers - "sqrt": torch.sqrt, "Pow": torch.pow, - # Misc - "Abs": torch.abs, "sign": torch.sign, - "floor": torch.floor, "ceiling": torch.ceil, - "Min": torch.minimum, "Max": torch.maximum, - # Piecewise / heaviside - "Heaviside": lambda x: torch.heaviside(x, torch.zeros_like(x)), - # Constants - "pi": getattr(torch, "pi", float(np.pi)), - "E": float(np.e), - } - - # --------------------------------------------------------------------------- # TorchExpr — container returned by torchify when limits are provided # --------------------------------------------------------------------------- @@ -121,9 +100,66 @@ def torchify(expr, variables, limits=None, params=None, modules=None): ------- TorchExpr if *limits* were given (use with ``torchquad_integrate``) callable if *limits* were ``None`` (plain torch function) + ------- + Usage: + x, k = sp.symbols("x k", real=True) + integrand = sp.exp(-x**2) * sp.exp(sp.I * k * x) + + texpr = torchify( + integrand, + variables=[x], + limits=[(x, -sp.oo, sp.oo)], + params=[k], + ) + + # if the expr is given as an sp integral: + inntegrand = integral_expression.function + limits = list(integral_expression.limits) + + texpr = torchify( + inntegrand, + variables=[x], + limits=limits, + params=[k], + ) + + # if the expr is given as an sp equation integral: + inntegrand = integral_equation.rhs.function + limits = list(integral_equation.rhs.limits) + + texpr = torchify( + inntegrand, + variables=[x], + limits=limits, + params=[k], + ) """ + def safe_sqrt(x): + return torch.sqrt(torch.as_tensor(x, dtype=torch.float64)) if modules is None: - modules = [_default_torch_modules()] + modules = [{ + # Trigonometric + "sin": torch.sin, "cos": torch.cos, "tan": torch.tan, + "asin": torch.asin, "acos": torch.acos, "atan": torch.atan, + "atan2": torch.atan2, + # Hyperbolic + "sinh": torch.sinh, "cosh": torch.cosh, "tanh": torch.tanh, + "asinh": torch.asinh, "acosh": torch.acosh, "atanh": torch.atanh, + # Exponentials / logs + "exp": torch.exp, "log": torch.log, "ln": torch.log, + "log10": torch.log10, "log2": torch.log2, + # Roots / powers + "sqrt": safe_sqrt, "Pow": torch.pow, + # Misc + "Abs": torch.abs, "sign": torch.sign, + "floor": torch.floor, "ceiling": torch.ceil, + "Min": torch.minimum, "Max": torch.maximum, + # Piecewise / heaviside + "Heaviside": lambda x: torch.heaviside(x, torch.zeros_like(x)), + # Constants + "pi": getattr(torch, "pi", float(np.pi)), + "E": float(np.e), + }] variables = list(variables) @@ -147,7 +183,6 @@ def torchify(expr, variables, limits=None, params=None, modules=None): lower, upper = limit_map[v] if lower == -oo and upper == oo: - # (-∞, ∞) → v = tan(t), dv = sec²(t) dt t_v = Symbol(f"t_{v.name}", real=True) subs_map[v] = tan(t_v) jacobian *= 1 / cos(t_v) ** 2 @@ -155,7 +190,6 @@ def torchify(expr, variables, limits=None, params=None, modules=None): domain.append([float(-pi / 2), float(pi / 2)]) elif lower != -oo and upper == oo: - # (a, ∞) → v = a + tan²(t), dv = 2 tan(t) sec²(t) dt t_v = Symbol(f"t_{v.name}", real=True) subs_map[v] = lower + tan(t_v) ** 2 jacobian *= 2 * tan(t_v) / cos(t_v) ** 2 @@ -163,7 +197,6 @@ def torchify(expr, variables, limits=None, params=None, modules=None): domain.append([0.0, float(pi / 2)]) elif lower == -oo and upper != oo: - # (-∞, b) → v = b - tan²(t), dv = -2 tan(t) sec²(t) dt t_v = Symbol(f"t_{v.name}", real=True) subs_map[v] = upper - tan(t_v) ** 2 jacobian *= -2 * tan(t_v) / cos(t_v) ** 2 @@ -175,10 +208,8 @@ def torchify(expr, variables, limits=None, params=None, modules=None): new_vars.append(v) domain.append([float(lower), float(upper)]) - # Substitute + bake in the Jacobian — all symbolic, done ONCE expr_work = expr.subs(subs_map) * jacobian - # Infer params from remaining free symbols if params is None: params = sorted( list(expr_work.free_symbols - set(new_vars)), @@ -216,6 +247,9 @@ def torchquad_integrate(texpr, params_values=None, method=None, N=21): ------- re, im : torch.Tensor Real and imaginary parts of the integral. + ------- + Usage: + re, im = torchquad_integrate(texpr, params_values=[1, 1, 1, 1], N=121) """ from torchquad import Simpson @@ -224,28 +258,25 @@ def torchquad_integrate(texpr, params_values=None, method=None, N=21): param_vals = list(params_values) if params_values else [] - def f_re(d): - vals = texpr.func(*[d[:, i] for i in range(d.shape[1])], *param_vals) - vals = torch.as_tensor(vals, device=d.device) - return vals.real if torch.is_complex(vals) else vals - def f_im(d): + def f(d): vals = texpr.func(*[d[:, i] for i in range(d.shape[1])], *param_vals) - vals = torch.as_tensor(vals, device=d.device) - return vals.imag if torch.is_complex(vals) else torch.zeros_like(vals) + return torch.as_tensor(vals, device=d.device) + + res = method.integrate(f, dim=texpr.dim, N=N, integration_domain= texpr.domain) + return (res.real if torch.is_complex(res) else res), (res.imag if torch.is_complex(res) else torch.zeros_like(res)) - re = method.integrate(f_re, dim=texpr.dim, N=N, integration_domain=texpr.domain) - im = method.integrate(f_im, dim=texpr.dim, N=N, integration_domain=texpr.domain) - return re, im def _simpson_weights_1d(a: float, b: float, N: int, *, device=None, dtype=None) -> torch.Tensor: - """Return Simpson weights including the dx/3 scaling (shape: [N]).""" + """Return Simpson weights including the dx/3 scaling (shape: [N]). + this is used to build the full tensor-product weights for torchquad_integrate. + """ if N < 3 or (N % 2) == 0: raise ValueError(f"Simpson rule requires odd N>=3; got N={N}") dx = (b - a) / (N - 1) w = torch.ones(N, device=device, dtype=dtype) - # 1, 4, 2, 4, ..., 2, 4, 1 + # 1, 4, 2, 4, ..., 2, 4, 1 this is the simppson pattern w[1:-1:2] = 4 w[2:-1:2] = 2 w = w * (dx / 3.0) @@ -369,6 +400,14 @@ def torch_integrate_batched_simpson( ------- re, im : torch.Tensor Shapes match the parameter batch shape (e.g. 250x250). + ------- + Usage: + >>> re, im = torch_integrate_batched_simpson( + ... texpr, + ... params_values=params_grid, # shape (250, 250, n_params) + ... N=21, + ... chunk_size_params=256, + ... chunk_size_points=10000) """ if texpr.dim <= 0: raise ValueError(f"texpr.dim must be positive; got dim={texpr.dim}") From 72ee8258e478a8706c16e3b4cf8610197d618b36 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Sat, 7 Mar 2026 01:50:53 +0300 Subject: [PATCH 10/11] added test --- tests/test_libtorch.ipynb | 110 ++++++++++++++++++++++++++++++-------- 1 file changed, 88 insertions(+), 22 deletions(-) diff --git a/tests/test_libtorch.ipynb b/tests/test_libtorch.ipynb index b0ac1b7..25a25e5 100644 --- a/tests/test_libtorch.ipynb +++ b/tests/test_libtorch.ipynb @@ -14,7 +14,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 1, "id": "7ce926ba", "metadata": {}, "outputs": [ @@ -31,9 +31,6 @@ } ], "source": [ - "\n", - "from __future__ import annotations\n", - "\n", "import sys\n", "from pathlib import Path\n", "\n", @@ -62,7 +59,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 2, "id": "0a58163c", "metadata": {}, "outputs": [ @@ -109,34 +106,103 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 3, "id": "17510ebc", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "re: 1.380388447018694 expected: 1.380388447043143 abs err: 2.4448887359085347e-11\n", - "im: 7.806255641895632e-18\n" - ] - } - ], + "outputs": [], "source": [ "# Test 1: complex Gaussian integral over (-oo, oo)\n", "# ∫ exp(-x^2) * exp(I*k*x) dx = sqrt(pi) * exp(-k^2/4)\n", "\n", "x, k = sp.symbols(\"x k\", real=True)\n", - "expr = sp.exp(-x**2) * sp.exp(sp.I * k * x)\n", + "integrand = sp.exp(-x**2) * sp.exp(sp.I * k * x)\n", "\n", "texpr = torchify(\n", - " expr,\n", + " integrand,\n", " variables=[x],\n", " limits=[(x, -sp.oo, sp.oo)],\n", " params=[k],\n", - ")\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "71dde1a0", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[[-1.5707963267948966, 1.5707963267948966]]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "texpr.domain" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "c4c85e18", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "1" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "texpr.dim" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "8a7c9139", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "c:\\Users\\Pc\\Desktop\\LOCALS\\F_Nutku_Studies\\.venv\\Lib\\site-packages\\torchquad\\integration\\utils.py:262: UserWarning: DEPRECATION WARNING: In future versions of torchquad, an array-like object will be returned.\n", + " warnings.warn(\n" + ] + } + ], + "source": [ "\n", - "re, im = torchquad_integrate(texpr, params_values=[1.0], N=151)\n", + "re, im = torchquad_integrate(texpr, params_values=[1.0], N=151)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "a6307112", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "re: 1.380388447018694 expected: 1.380388447043143 abs err: 2.4448887359085347e-11\n", + "im: 3.621235256101585e-17\n" + ] + } + ], + "source": [ "\n", "expected = math.sqrt(math.pi) * math.exp(-(1.0**2) / 4.0)\n", "re_f = float(re.detach().cpu())\n", @@ -151,7 +217,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 8, "id": "26b67910", "metadata": {}, "outputs": [ @@ -192,7 +258,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": ".venv (3.12.0)", "language": "python", "name": "python3" }, From d4deeff3c9982ae3918668c412dd74d125ba4fd5 Mon Sep 17 00:00:00 2001 From: ibeuler Date: Thu, 12 Mar 2026 05:30:08 +0300 Subject: [PATCH 11/11] oop libtorch --- src/libtorch.py | 1465 +++++++++++++++++++++++++++++++++++------------ 1 file changed, 1099 insertions(+), 366 deletions(-) diff --git a/src/libtorch.py b/src/libtorch.py index 6147db6..01b5789 100644 --- a/src/libtorch.py +++ b/src/libtorch.py @@ -1,40 +1,63 @@ # -*- coding: utf-8 -*- -""" -libtorch.py — SymPy-to-Torch conversion and torchquad integration. +"""libtorch.py — SymPy-to-Torch conversion and numerical integration helpers. Part of libphysics. -Two-function design for performance: - 1. torchify() — expensive SymPy work done ONCE (change-of-vars + lambdify) - 2. torchquad_integrate() — cheap torch-only integration, called MANY times +Two-stage design for performance: + 1. libtorch.torchify() — expensive SymPy work done ONCE + (change-of-vars + ``lambdify``) + 2. TorchExpr.torchquad_integrate() — cheap torch-only integration, + called MANY times -Usage +USAGE ===== - from libphysics.libtorch import torchify, torchquad_integrate - - # Step 1: build once (slow: SymPy subs + lambdify, ~200-500 ms) - wigner = torchify( - integrand, - variables=[y_A, y_B], - limits=[(y_A, -oo, oo), (y_B, -oo, oo)], - params=[x_A, p_A, x_B, p_B], - ) + from libphysics.libtorch import libtorch + + # Step 1: build once (slow: SymPy subs + lambdify, ~200–500 ms) + lt = libtorch() + texpr = lt.torchify( + sp.Integral(integrand, (y_A, -sp.oo, sp.oo), (y_B, -sp.oo, sp.oo)) + ) - # Step 2: integrate many times (fast: pure torch, ~3-8 ms each) - re, im = torchquad_integrate(wigner, params_values=[1, 1, 1, 1], N=121) + # Step 2: integrate many times (fast: pure torch, ~3–8 ms each) + re, im = texpr.torchquad_integrate(N=121) + + # Batched integration over a parameter grid using Simpson rule + re_b, im_b = texpr.torch_integrate_batched_simpson( + params_values=params_grid, + N=21, + chunk_size_params=256, + chunk_size_points=10_000, + ) """ from dataclasses import dataclass import sys -from typing import List, Callable +from typing import List, Callable, Any import numpy as np import torch -from sympy import lambdify, Symbol, tan, cos, pi, oo, Integer +from scipy import special as scipy_special +from sympy import ( + Eq, + Function, + Integral, + Integer, + Symbol, + atan, + atanh, + cos, + cosh, + diff, + lambdify, + oo, + pi, + sinh, + tan, + tanh, +) from loguru import logger -# Remove loguru's default handler only if user requests it via env var. -# Set LIBPHYSICS_REMOVE_LOGURU=1 or "true"/"yes" to disable. def setup_logging(enable: bool = False, level: str = "INFO"): """ By default, we remove all handlers (silence). @@ -54,90 +77,719 @@ def setup_logging(enable: bool = False, level: str = "INFO"): setup_logging(enable=False) -# --------------------------------------------------------------------------- -# TorchExpr — container returned by torchify when limits are provided -# --------------------------------------------------------------------------- + @dataclass class TorchExpr: - """Pre-built torch function + finite integration domain.""" + """Pre-built torch function together with its finite integration domain. + + This is the object returned by ``libtorch.torchify`` when integration + limits are supplied (directly or via a SymPy ``Integral``). It stores + both the numerical callable and all metadata needed by the integration + helpers in this module. + + Attributes + ---------- + func : Callable + The torch-compatible function produced by ``lambdify``. + Its signature is:: + + func(new_var_0, ..., new_var_{dim-1}, param_0, ..., param_{n_params-1}) + + where the ``new_var_*`` correspond to the transformed integration + variables after any change-of-variables and the ``param_*`` are any + additional symbolic parameters. + domain : list[list[float]] + Finite box for the integration domain. Each entry is ``[a, b]`` + for one dimension after change-of-variables. This is what the + integrators use as their integration range. + dim : int + Number of integration variables (i.e. the dimensionality of the + domain). + n_params : int + Number of additional symbolic parameters expected by ``func``. + sympy_expr : Any + The SymPy expression *after* all symbolic substitutions and + Jacobians have been applied. Kept for debugging / inspection. + variables : list[sympy.Symbol] | None + Integration variables used by the numerical object after any + change-of-variables has been applied. For infinite limits these + are the transformed symbols such as ``t_x`` rather than the + original symbolic variable ``x``. + """ func: Callable # f(new_var_0, …, new_var_n, param_0, …, param_m) domain: List[List[float]] # finite box for torchquad dim: int # number of integration variables n_params: int # number of extra parameters + sympy_expr: Any # sympy expression reduced if + variables: List[Symbol] = None # sympy variables after change of variables + + @staticmethod + def _rule_simpson(a: float, b: float, N: int, *, device=None, dtype=None): + if N < 3 or (N % 2) == 0: + raise ValueError(f"Simpson rule requires odd N>=3; got N={N}") + coords = torch.linspace(a, b, N, device=device, dtype=dtype) + dx = (b - a) / (N - 1) + w = torch.ones(N, device=device, dtype=dtype) + w[1:-1:2] = 4 + w[2:-1:2] = 2 + weights = w * (dx / 3.0) + return coords, weights + + @staticmethod + def _rule_trapezoid(a: float, b: float, N: int, *, device=None, dtype=None): + if N < 2: + raise ValueError(f"Trapezoid rule requires N>=2; got N={N}") + coords = torch.linspace(a, b, N, device=device, dtype=dtype) + dx = (b - a) / (N - 1) + w = torch.ones(N, device=device, dtype=dtype) + w[0] = 0.5 + w[-1] = 0.5 + weights = w * dx + return coords, weights + + @staticmethod + def _rule_gauss_legendre(a: float, b: float, N: int, *, device=None, dtype=None): + if N < 1: + raise ValueError(f"Gauss-Legendre rule requires N>=1; got N={N}") + nodes, weights = np.polynomial.legendre.leggauss(N) + nodes_t = torch.as_tensor(nodes, device=device, dtype=dtype) + weights_t = torch.as_tensor(weights, device=device, dtype=dtype) + # Affine map from [-1, 1] to [a, b] + coords = 0.5 * (b - a) * nodes_t + 0.5 * (b + a) + weights = 0.5 * (b - a) * weights_t + return coords, weights + + @staticmethod + def _resolve_quadrature_rule(method): + if callable(method): + return method + + if method is None: + method = "simpson" + + method_name = str(method).strip().lower() + if method_name == "simpson": + return TorchExpr._rule_simpson + if method_name == "trapezoid": + return TorchExpr._rule_trapezoid + if method_name in {"gauss-legendre", "gauss_legendre", "legendre"}: + return TorchExpr._rule_gauss_legendre + + raise ValueError( + "Unknown integration method. Expected one of {'simpson', 'trapezoid', 'gauss-legendre'} or a callable. " + f"Got: {method}" + ) + @staticmethod + def _is_scalar_like(x) -> bool: + """Return ``True`` for Python scalars or 0-dim tensors.""" + if torch.is_tensor(x): + return x.ndim == 0 + return isinstance(x, (int, float, complex, np.number)) + + @staticmethod + def _normalize_params_values(params_values, n_params: int, *, device, dtype): + """Normalize parameter inputs into a 2D tensor and batch shape. + + Accepted input formats are: + + - ``None`` when ``n_params == 0``, + - a tensor of shape ``(..., n_params)``, or + - a list/tuple of length ``n_params`` whose elements are scalars or + tensors with the same shape. + """ + if params_values is None: + if n_params != 0: + raise ValueError(f"Expected {n_params} params; got None") + return torch.empty((1, 0), device=device, dtype=dtype), tuple() + + if torch.is_tensor(params_values): + params_tensor = params_values + if params_tensor.shape[-1] != n_params: + raise ValueError( + f"params_values last dimension must be n_params={n_params}; got shape={tuple(params_tensor.shape)}" + ) + batch_shape = tuple(params_tensor.shape[:-1]) + batch_size = int(np.prod(batch_shape)) if batch_shape else 1 + params_flat = params_tensor.reshape(batch_size, n_params).to(device=device, dtype=dtype) + return params_flat, batch_shape -# --------------------------------------------------------------------------- -# torchify — expensive work done ONCE -# --------------------------------------------------------------------------- -def torchify(expr, variables, limits=None, params=None, modules=None): - """ - Convert a SymPy expression into a torch-compatible function via lambdify. - - If *limits* are provided, performs change-of-variables for infinite / - semi-infinite limits at the **SymPy level** (symbolic, done once) and - multiplies by the analytic Jacobian. Returns a ``TorchExpr`` ready - for ``torchquad_integrate``. - - If *limits* are **not** provided, returns a plain callable (simple - lambdify with torch modules). + if not isinstance(params_values, (list, tuple)): + params_values = [params_values] - Parameters - ---------- - expr : sympy.Expr - Symbolic expression (may be complex-valued). - variables : list[sympy.Symbol] - Integration variables (order matters). - limits : list[tuple] or None - ``(var, lower, upper)`` for each variable. ``sympy.oo`` / - ``-sympy.oo`` trigger automatic change of variables. - params : list[sympy.Symbol] or None - Extra symbolic parameters that appear in *expr* but are **not** - integrated over. If ``None``, inferred automatically. - modules : list[dict] or None - Custom lambdify module list. ``None`` → default torch mapping. + if len(params_values) != n_params: + raise ValueError(f"Expected {n_params} params; got {len(params_values)}") - Returns - ------- - TorchExpr if *limits* were given (use with ``torchquad_integrate``) - callable if *limits* were ``None`` (plain torch function) - ------- - Usage: - x, k = sp.symbols("x k", real=True) - integrand = sp.exp(-x**2) * sp.exp(sp.I * k * x) + batch_shape = None + for param_value in params_values: + if torch.is_tensor(param_value) and param_value.ndim > 0: + batch_shape = tuple(param_value.shape) + break + if not TorchExpr._is_scalar_like(param_value): + param_tensor = torch.as_tensor(param_value) + if param_tensor.ndim > 0: + batch_shape = tuple(param_tensor.shape) + break + if batch_shape is None: + batch_shape = tuple() + + params_list = [] + for param_value in params_values: + param_tensor = torch.as_tensor(param_value, device=device) + if param_tensor.ndim == 0: + if batch_shape: + param_tensor = param_tensor.expand(batch_shape) + else: + if tuple(param_tensor.shape) != batch_shape: + raise ValueError( + f"All parameter tensors must have the same shape; got {tuple(param_tensor.shape)} vs {batch_shape}" + ) + params_list.append(param_tensor) + + params_tensor = torch.stack(params_list, dim=-1).to(dtype=dtype) + batch_size = int(np.prod(batch_shape)) if batch_shape else 1 + params_flat = params_tensor.reshape(batch_size, n_params) + return params_flat, batch_shape + + # ------------------------------------------------------------------ + # Convenience methods: numeric integration on this TorchExpr + # ------------------------------------------------------------------ - texpr = torchify( + def torchquad_integrate(self, params_values=None, method=None, N: int = 21): + """Integrate this ``TorchExpr`` using the torchquad-based integrator. + + This is the simplest integration entry point. It is intended for + low-dimensional problems where a single torchquad integral is + sufficient. + + Parameters + ---------- + params_values : list | tuple | torch.Tensor | None, optional + Numerical values for the symbolic parameters. The accepted + formats match those of ``torchquad_integrate``: + + - ``None`` if there are no parameters, + - a Python sequence (list/tuple) of numbers/tensors, or + - a tensor whose last dimension is ``n_params``. + method : torchquad integrator instance or None, optional + If ``None``, a default ``Simpson`` integrator from torchquad is + constructed internally. + N : int, optional + Resolution parameter passed through to torchquad. Larger values + increase accuracy but also the number of function evaluations. + + Returns + ------- + re, im : torch.Tensor + Real and imaginary parts of the integral. If the integrand is + real-valued, ``im`` will be a zero tensor. + + USAGE + ===== + 1. Simple 1D integral without parameters:: + + ````python + from libphysics.libtorch import libtorch + from sympy import symbols, Integral, exp, oo + + x = symbols("x", real=True) + integral_expr = Integral(exp(-x**2), (x, -oo, oo)) + + lt = libtorch() + texpr = lt.torchify(integral_expr) + re, im = texpr.torchquad_integrate(N=121) + ```` + + 2. Integral with parameters (e.g. Fourier transform):: + + ````python + from libphysics.libtorch import libtorch + from sympy import symbols, Integral, exp, I, oo + import torch + + x, k = symbols("x k", real=True) + integral_expr = Integral(exp(-x**2) * exp(I * k * x), (x, -oo, oo)) + + lt = libtorch() + texpr = lt.torchify(integral_expr) + + # Evaluate at k = 0.5 + re, im = texpr.torchquad_integrate(params_values=[0.5], N=121) + ```` + """ + from torchquad import Simpson + + if method is None: + method = Simpson() + + param_vals = list(params_values) if params_values else [] + + def integrand(domain_points): + param_vals_local = [torch.as_tensor(param_value, device=domain_points.device) for param_value in param_vals] + values = self.func(*[domain_points[:, i] for i in range(domain_points.shape[1])], *param_vals_local) + return torch.as_tensor(values, device=domain_points.device) + + result = method.integrate( integrand, - variables=[x], - limits=[(x, -sp.oo, sp.oo)], - params=[k], + dim=self.dim, + N=N, + integration_domain=self.domain, + ) + return ( + result.real if torch.is_complex(result) else result, + result.imag if torch.is_complex(result) else torch.zeros_like(result), + ) + + def torch_integrate_batched( + self, + *, + params_values = None, + method: str | Callable = "simpson", + N: int = 121, + chunk_size_params: int = 256, + chunk_size_points: int | None = None, + device=None, + dtype=None, + ): + """Batched tensor-product quadrature integration for this ``TorchExpr``. + + This method is intended for situations where you want to evaluate + the same integral for many different parameter values (for example, + on a 2D or 3D mesh). + + Parameters + ---------- + params_values : tensor | list | tuple + Numerical parameters. See ``torch_integrate_batched`` + for a full description of the accepted shapes and types. + method : str or callable, optional + Quadrature rule name (``"simpson"``, ``"trapezoid"``, + ``"gauss-legendre"``) or a callable returning + ``(coords_1d, weights_1d)`` for ``(a, b, N)``. + N : int, optional + Odd number of Simpson points per dimension. + chunk_size_params : int, optional + Number of parameter points processed in one chunk. Reducing + this lowers peak memory usage at the cost of more Python loops. + chunk_size_points : int or None, optional + Maximum number of grid points processed per chunk. ``None`` + means "all at once". + device : torch.device or None, optional + Device used for all internal tensors. If ``None``, CUDA is + used when available, otherwise CPU. + dtype : torch.dtype or None, optional + Floating dtype for internal computations. Defaults to + ``float32`` on CUDA and ``float64`` on CPU. + + Returns + ------- + re, im : torch.Tensor + Tensors with the same batch shape as the input parameters, + containing the real and imaginary parts of the integral. + + USAGE + ===== + 1. 1D integral without parameters (single value):: + + ````python + from libphysics.libtorch import libtorch + from sympy import symbols, Integral, exp, oo + + x = symbols("x", real=True) + integral_expr = Integral(exp(-x**2), (x, -oo, oo)) + + lt = libtorch() + texpr = lt.torchify(integral_expr) + + # No parameters → pass ``None`` for ``params_values`` + re, im = texpr.torch_integrate_batched( + params_values=None, + N=121, ) + ```` + + 2. 1D integral evaluated on a grid of parameter values:: + + ````python + from libphysics.libtorch import libtorch + from sympy import symbols, Integral, exp, I, oo + import torch - # if the expr is given as an sp integral: - inntegrand = integral_expression.function - limits = list(integral_expression.limits) + x, k = symbols("x k", real=True) + integral_expr = Integral(exp(-x**2) * exp(I * k * x), (x, -oo, oo)) + + lt = libtorch() + texpr = lt.torchify(integral_expr) + + # Parameter grid: 250 points between -5 and 5 + k_grid = torch.linspace(-5.0, 5.0, 250).unsqueeze(-1) # shape (250, 1) + + re, im = texpr.torch_integrate_batched( + params_values=k_grid, + N=51, + chunk_size_params=64, + chunk_size_points=10_000, + ) + ```` + """ + if self.dim <= 0: + raise ValueError(f"self.dim must be positive; got dim={self.dim}") + if self.n_params < 0: + raise ValueError(f"self.n_params must be non-negative; got n_params={self.n_params}") + if len(self.domain) != self.dim: + raise ValueError(f"self.domain must have length dim={self.dim}; got {len(self.domain)}") + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if dtype is None: + dtype = torch.float32 if device.type == "cuda" else torch.float64 + + if chunk_size_params <= 0: + raise ValueError(f"chunk_size_params must be positive; got {chunk_size_params}") + if chunk_size_points is not None and chunk_size_points <= 0: + raise ValueError(f"chunk_size_points must be positive; got {chunk_size_points}") + + params_flat, batch_shape = self._normalize_params_values( + params_values, + self.n_params, + device=device, + dtype=dtype, + ) + batch_size = int(params_flat.shape[0]) + + rule = self._resolve_quadrature_rule(method) + + coords_1d = [] + weights_1d = [] + for (lower_bound, upper_bound) in self.domain: + lower_bound = float(lower_bound) + upper_bound = float(upper_bound) + coords, weights = rule(lower_bound, upper_bound, N, device=device, dtype=dtype) + coords = torch.as_tensor(coords, device=device, dtype=dtype).reshape(-1) + weights = torch.as_tensor(weights, device=device, dtype=dtype).reshape(-1) + if coords.numel() != weights.numel(): + raise ValueError( + f"Quadrature rule returned mismatched coordinates/weights lengths: {coords.numel()} vs {weights.numel()}" + ) + coords_1d.append(coords) + weights_1d.append(weights) - texpr = torchify( - inntegrand, - variables=[x], - limits=limits, - params=[k], + if self.dim == 1: + grid = coords_1d[0].unsqueeze(1) + weights = weights_1d[0] + else: + grid = torch.cartesian_prod(*coords_1d) + weights = weights_1d[0] + for weight_1d in weights_1d[1:]: + weights = torch.kron(weights, weight_1d) + n_points = int(grid.shape[0]) + + if chunk_size_points is None: + chunk_size_points = n_points + + re_out = torch.empty(batch_size, device=device, dtype=dtype) + im_out = torch.empty(batch_size, device=device, dtype=dtype) + + for start_param in range(0, batch_size, chunk_size_params): + stop_param = min(batch_size, start_param + chunk_size_params) + param_chunk = params_flat[start_param:stop_param, :] + param_args = [param_chunk[:, index].unsqueeze(0) for index in range(self.n_params)] + + re_acc = torch.zeros(stop_param - start_param, device=device, dtype=dtype) + im_acc = torch.zeros(stop_param - start_param, device=device, dtype=dtype) + + for start_point in range(0, n_points, chunk_size_points): + stop_point = min(n_points, start_point + chunk_size_points) + grid_chunk = grid[start_point:stop_point, :] + weight_chunk = weights[start_point:stop_point].unsqueeze(1) + + var_args = [grid_chunk[:, index].unsqueeze(1) for index in range(self.dim)] + values = self.func(*var_args, *param_args) + values = torch.as_tensor(values, device=device) + + if values.ndim == 1: + values = values.unsqueeze(1) + elif ( + values.ndim == 2 + and values.shape[0] == (stop_param - start_param) + and values.shape[1] == (stop_point - start_point) + ): + values = values.t().contiguous() + + if values.shape[0] != (stop_point - start_point) or values.shape[1] != (stop_param - start_param): + raise RuntimeError( + f"Unexpected integrand output shape {tuple(values.shape)}; expected ({stop_point - start_point}, {stop_param - start_param})" + ) + + if torch.is_complex(values): + re_acc += (weight_chunk * values.real).sum(dim=0) + im_acc += (weight_chunk * values.imag).sum(dim=0) + else: + re_acc += (weight_chunk * values).sum(dim=0) + + re_out[start_param:stop_param] = re_acc + im_out[start_param:stop_param] = im_acc + + re_out = re_out.reshape(batch_shape) + im_out = im_out.reshape(batch_shape) + return re_out, im_out + + def torch_integrate_batched_simpson( + self, + *, + params_values = None, + N: int = 121, + chunk_size_params: int = 256, + chunk_size_points: int | None = None, + device=None, + dtype=None, + ): + """Backward-compatible wrapper for ``torch_integrate_batched(method='simpson')``.""" + return self.torch_integrate_batched( + params_values=params_values, + method="simpson", + N=N, + chunk_size_params=chunk_size_params, + chunk_size_points=chunk_size_points, + device=device, + dtype=dtype, + ) + + + +class libtorch: + def __init__(self): + pass + + @staticmethod + def _normalize_cov_method(change_of_variables_method: str) -> str: + method = str(change_of_variables_method).strip().lower() + aliases = { + "tangent": "tangent", + "tan": "tangent", + "algebraic": "algebraic", + "rational": "algebraic", + "tanh-sinh": "tanh-sinh", + "tanh_sinh": "tanh-sinh", + "tanhsinh": "tanh-sinh", + } + if method not in aliases: + raise ValueError( + "Unknown change_of_variables_method. Expected one of {'tangent', 'algebraic', 'tanh-sinh'}. " + f"Got: {change_of_variables_method}" ) + return aliases[method] + + @staticmethod + def _inset_open_interval(lower: float, upper: float, eps: float) -> list[float]: + return [float(lower + eps), float(upper - eps)] + + def _transform_limit_with_method(self, variable, lower, upper, *, change_of_variables_method: str, eps: float): + """Return transformed variable, substitution, jacobian and finite domain.""" + method = self._normalize_cov_method(change_of_variables_method) + + if lower != -oo and upper != oo: + return variable, variable, Integer(1), [float(lower), float(upper)] + + t_v = Symbol(f"t_{variable.name}", real=True) + open_interval = None + + if method == "tangent": + if lower == -oo and upper == oo: + mapped = tan(t_v) + open_interval = (float(-pi / 2), float(pi / 2)) + elif lower != -oo and upper == oo: + mapped = lower + tan(t_v) ** 2 + open_interval = (0.0, float(pi / 2)) + elif lower == -oo and upper != oo: + mapped = upper - tan(t_v) ** 2 + open_interval = (0.0, float(pi / 2)) + else: + raise ValueError("Unexpected limit pattern in tangent transform") + + elif method == "algebraic": + if lower == -oo and upper == oo: + mapped = t_v / (1 - t_v ** 2) + open_interval = (-1.0, 1.0) + elif lower != -oo and upper == oo: + mapped = lower + t_v / (1 - t_v) + open_interval = (0.0, 1.0) + elif lower == -oo and upper != oo: + mapped = upper - t_v / (1 - t_v) + open_interval = (0.0, 1.0) + else: + raise ValueError("Unexpected limit pattern in algebraic transform") + + elif method == "tanh-sinh": + # Finite-interval parameter t in (-1,1) or (0,1) mapped through sinh(atanh(t)). + core = sinh(atanh(t_v)) + if lower == -oo and upper == oo: + mapped = core + open_interval = (-1.0, 1.0) + elif lower != -oo and upper == oo: + mapped = lower + core ** 2 + open_interval = (0.0, 1.0) + elif lower == -oo and upper != oo: + mapped = upper - core ** 2 + open_interval = (0.0, 1.0) + else: + raise ValueError("Unexpected limit pattern in tanh-sinh transform") + + else: + raise ValueError(f"Unsupported change_of_variables_method: {method}") + + jacobian = diff(mapped, t_v) + domain = self._inset_open_interval(open_interval[0], open_interval[1], eps) + return t_v, mapped, jacobian, domain + + def _merge_lambdify_modules(self, modules, extra_mapping): + """Merge lambdify modules with an extra mapping dict at highest priority.""" + if not extra_mapping: + return modules + + if modules is None: + modules = self._default_modules() + + if isinstance(modules, (list, tuple)): + merged = list(modules) + merged.insert(0, dict(extra_mapping)) + return merged + + return [dict(extra_mapping), modules] + + def _build_nested_definite_integral_callable(self, texpr: TorchExpr, n_params: int, *, inner_N: int = 61): + """Create a callable for lambdify that evaluates a definite nested integral numerically.""" + + def eval_nested(*args): + if len(args) != n_params: + raise ValueError(f"Expected {n_params} nested integral params; got {len(args)}") + + if n_params == 0: + re_val, im_val = texpr.torchquad_integrate(N=inner_N) + re_tensor = torch.as_tensor(re_val) + im_tensor = torch.as_tensor(im_val, device=re_tensor.device) + return re_tensor + 1j * im_tensor + + args_tensors = [torch.as_tensor(arg) for arg in args] + broadcasted = torch.broadcast_tensors(*args_tensors) + target_device = broadcasted[0].device + batch_shape = tuple(broadcasted[0].shape) + flat_params = [tensor.reshape(-1) for tensor in broadcasted] + batch_size = int(flat_params[0].shape[0]) + + out_re = [] + out_im = [] + for idx in range(batch_size): + params_i = [flat_params[param_index][idx] for param_index in range(n_params)] + re_val, im_val = texpr.torchquad_integrate(params_values=params_i, N=inner_N) + out_re.append(torch.as_tensor(re_val, device=target_device)) + out_im.append(torch.as_tensor(im_val, device=target_device)) + + re_tensor = torch.stack(out_re).reshape(batch_shape) + im_tensor = torch.stack(out_im).reshape(batch_shape).to(device=re_tensor.device) + return re_tensor + 1j * im_tensor + + return eval_nested + + def _replace_nested_definite_integrals( + self, + expr, + *, + inner_N: int = 61, + change_of_variables_method: str = "tangent", + cov_eps: float = 1e-7, + ): + """Replace nested definite Integrals with numeric callables. + + Any nested integral with non-definite limits triggers a ValueError. + """ + if not hasattr(expr, "has") or not expr.has(Integral): + return expr, {} + + expr_work = expr + nested_mapping = {} + counter = 0 + + while expr_work.has(Integral): + leaf_integrals = [itg for itg in expr_work.atoms(Integral) if not itg.function.has(Integral)] + if not leaf_integrals: + break + + for inner_integral in leaf_integrals: + for lim in inner_integral.limits: + if len(lim) != 3: + raise ValueError( + "The integrand must not contain unevaluated indefinite integrals. " + f"Found nested integral with non-definite limit: {inner_integral}" + ) + + integration_vars = [lim[0] for lim in inner_integral.limits] + param_symbols = sorted( + list(inner_integral.free_symbols - set(integration_vars)), + key=lambda symbol: symbol.name, + ) + + nested_texpr = self.torchify( + inner_integral, + change_of_variables_method=change_of_variables_method, + cov_eps=cov_eps, + ) + nested_name = f"_nested_definite_integral_{counter}" + counter += 1 - # if the expr is given as an sp equation integral: - inntegrand = integral_equation.rhs.function - limits = list(integral_equation.rhs.limits) + nested_mapping[nested_name] = self._build_nested_definite_integral_callable( + nested_texpr, + len(param_symbols), + inner_N=inner_N, + ) + + placeholder = Function(nested_name)(*param_symbols) + expr_work = expr_work.xreplace({inner_integral: placeholder}) - texpr = torchify( - inntegrand, - variables=[x], - limits=limits, - params=[k], + if expr_work.has(Integral): + raise ValueError( + "The integrand must not contain unevaluated integrals. " + "Only nested definite integrals are supported." ) - """ - def safe_sqrt(x): - return torch.sqrt(torch.as_tensor(x, dtype=torch.float64)) - if modules is None: - modules = [{ + + return expr_work, nested_mapping + + def _default_modules(self): + """Internal helper: default SymPy → torch mapping for ``lambdify``. + + The returned list is passed directly to SymPy's ``lambdify`` as the + ``modules`` argument. It exposes a minimal subset of functions + implemented with torch operations so that the resulting numerical + function is differentiable and GPU-friendly where possible. + + Keeping this mapping in a separate method avoids cluttering + :meth:`torchify` and makes it easy to customize or extend in user + code by subclassing ``libtorch``. + """ + + def safe_sqrt(x): + return torch.sqrt(torch.as_tensor(x, dtype=torch.float64)) + + def safe_erf(x): + """Error function that supports real torch tensors and complex inputs.""" + if torch.is_tensor(x): + if torch.is_complex(x): + # torch.erf is not implemented for complex tensors on CPU. + out = scipy_special.erf(x.detach().cpu().numpy()) + return torch.as_tensor(out, device=x.device, dtype=x.dtype) + return torch.erf(x) + return scipy_special.erf(x) + + def safe_erfc(x): + """Complementary error function with complex-input fallback.""" + if torch.is_tensor(x): + if torch.is_complex(x): + out = scipy_special.erfc(x.detach().cpu().numpy()) + return torch.as_tensor(out, device=x.device, dtype=x.dtype) + return torch.erfc(x) + return scipy_special.erfc(x) + + return [{ # Trigonometric "sin": torch.sin, "cos": torch.cos, "tan": torch.tan, "asin": torch.asin, "acos": torch.acos, "atan": torch.atan, @@ -148,10 +800,13 @@ def safe_sqrt(x): # Exponentials / logs "exp": torch.exp, "log": torch.log, "ln": torch.log, "log10": torch.log10, "log2": torch.log2, + # Special functions used by symbolic integral evaluation + "erf": safe_erf, "erfc": safe_erfc, # Roots / powers "sqrt": safe_sqrt, "Pow": torch.pow, # Misc "Abs": torch.abs, "sign": torch.sign, + "conjugate": torch.conj, "conj": torch.conj, "floor": torch.floor, "ceiling": torch.ceil, "Min": torch.minimum, "Max": torch.maximum, # Piecewise / heaviside @@ -161,343 +816,421 @@ def safe_sqrt(x): "E": float(np.e), }] - variables = list(variables) + def torchify( + self, + expr, + *, + variables=None, + limits=None, + params=None, + modules=None, + change_of_variables_method: str = "tangent", + cov_eps: float = 1e-7, + ): + """Convert a SymPy object into a torch-compatible callable or TorchExpr. + + This mirrors the original functional ``torchify`` but lives as a + method on ``libtorch``. It performs change-of-variables for + infinite / semi-infinite limits at the SymPy level (once), then + ``lambdify``'s the transformed expression. + + ``expr`` can be: + + - a plain SymPy expression (e.g. ``exp(-x**2)``), + - a SymPy ``Integral`` (``Integral(f(x), (x, -oo, oo))``), or + - a SymPy ``Eq`` whose right-hand side is an ``Integral``. + + If an ``Integral`` (or equation with an integral) is passed, the + integrand and limits are extracted internally, so you no longer + need to manually use ``.function`` / ``.limits``. + + Parameters + ---------- + expr : sympy.Expr | sympy.Integral | sympy.Eq + Symbolic expression (may be complex-valued) or an integral. + variables : list[sympy.Symbol] or None + Integration variables (order matters). If ``expr`` is an + ``Integral`` and ``variables`` is ``None``, they are inferred + from the integral limits. + limits : list[tuple] or None + ``(var, lower, upper)`` for each variable. ``oo`` / ``-oo`` + trigger automatic change of variables. If ``expr`` is an + ``Integral`` and ``limits`` is ``None``, they are taken from + ``expr.limits``. + params : list[sympy.Symbol] or None + Extra symbolic parameters that appear in the transformed + expression but are **not** integrated over. If ``None``, they + are inferred automatically from the free symbols. + modules : list[dict] or None + Custom ``lambdify`` module list. ``None`` → default torch + mapping from ``_default_modules``. + change_of_variables_method : str, optional + Method used to transform infinite/semi-infinite domains to finite + ones. Supported values: ``"tangent"``, ``"algebraic"``, + ``"tanh-sinh"``. + cov_eps : float, optional + Small inward shift for transformed open intervals to avoid + endpoint singular evaluations. + + Returns + ------- + TorchExpr + If limits are provided (directly or via an ``Integral``), ready + for numerical integration. + callable + If ``limits`` is ``None``, a plain torch-ready function + (no integration). + + USAGE + ===== + 1. Plain function (no integration):: + + ````python + from libphysics.libtorch import libtorch + from sympy import symbols, exp + + x = symbols("x", real=True) + + lt = libtorch() + f = lt.torchify(exp(-x**2), variables=[x]) + ```` + + 2. Integral with finite limits:: + + ````python + from libphysics.libtorch import libtorch + from sympy import symbols, Integral, exp + + x = symbols("x", real=True) + integral = Integral(exp(-x**2), (x, 0, 1)) + + lt = libtorch() + texpr = lt.torchify(integral) + ```` + + 3. Integral with infinite limits (automatic change of variables):: + + ````python + from libphysics.libtorch import libtorch + from sympy import symbols, Integral, exp, oo + + x = symbols("x", real=True) + integral = Integral(exp(-x**2), (x, -oo, oo)) + + lt = libtorch() + texpr = lt.torchify(integral) + ```` + + 4. Equation with integral on the right-hand side:: + + ````python + from libphysics.libtorch import libtorch + from sympy import symbols, Eq, Integral, exp, oo + + x, y = symbols("x y", real=True) + eq = Eq(y, Integral(exp(-x**2), (x, -oo, oo))) + + lt = libtorch() + texpr = lt.torchify(eq) # integral auto-extracted + ```` + """ + + # Detect integral / equation-with-integral and extract base expr/limits. + integral = None + if isinstance(expr, Integral): + integral = expr + elif isinstance(expr, Eq) and isinstance(expr.rhs, Integral): + integral = expr.rhs + + if integral is not None: + base_expr = integral.function + nested_modules = {} + if base_expr.has(Integral): + # Evaluate nested definite integrals numerically using torchquad. + base_expr, nested_modules = self._replace_nested_definite_integrals( + base_expr, + change_of_variables_method=change_of_variables_method, + cov_eps=cov_eps, + ) + if limits is None: + limits = list(integral.limits) + if variables is None and limits is not None: + variables = [lim[0] for lim in limits] + else: + base_expr = expr + nested_modules = {} + + if variables is None: + raise ValueError("'variables' must be provided when expr is not an Integral") + + if modules is None: + modules = self._default_modules() + modules = self._merge_lambdify_modules(modules, nested_modules) + + variables = list(variables) + + # ------------------------------------------------------------------ + # Simple case: no limits → plain lambdify + # ------------------------------------------------------------------ + if limits is None: + return lambdify(tuple(variables), base_expr, modules=modules) + + # ------------------------------------------------------------------ + # With limits: symbolic change of variables (done once, fast forever) + # ------------------------------------------------------------------ + limit_map = {lim[0]: (lim[1], lim[2]) for lim in limits} + + new_vars = [] + domain = [] + subs_map = {} + jacobian = Integer(1) + + for v in variables: + lower, upper = limit_map[v] + t_var, mapped_expr, jac_expr, domain_entry = self._transform_limit_with_method( + v, + lower, + upper, + change_of_variables_method=change_of_variables_method, + eps=cov_eps, + ) + new_vars.append(t_var) + domain.append(domain_entry) + if t_var != v: + subs_map[v] = mapped_expr + jacobian *= jac_expr + + expr_work = base_expr.subs(subs_map) * jacobian + + if params is None: + params = sorted( + list(expr_work.free_symbols - set(new_vars)), + key=lambda s: s.name, + ) + else: + params = list(params) + + # Lambdify: new_vars first, then params + arglist = tuple(new_vars + params) + func = lambdify(arglist, expr_work, modules=modules) + + return TorchExpr( + func=func, + domain=domain, + dim=len(new_vars), + n_params=len(params), + sympy_expr=expr_work, + variables=new_vars, + ) - # ------------------------------------------------------------------ - # Simple case: no limits → plain lambdify - # ------------------------------------------------------------------ - if limits is None: - return lambdify(tuple(variables), expr, modules=modules) - # ------------------------------------------------------------------ - # With limits: symbolic change of variables (done once, fast forever) - # ------------------------------------------------------------------ - limit_map = {lim[0]: (lim[1], lim[2]) for lim in limits} - - new_vars = [] - domain = [] - subs_map = {} - jacobian = Integer(1) - - for v in variables: - lower, upper = limit_map[v] - - if lower == -oo and upper == oo: - t_v = Symbol(f"t_{v.name}", real=True) - subs_map[v] = tan(t_v) - jacobian *= 1 / cos(t_v) ** 2 - new_vars.append(t_v) - domain.append([float(-pi / 2), float(pi / 2)]) - - elif lower != -oo and upper == oo: - t_v = Symbol(f"t_{v.name}", real=True) - subs_map[v] = lower + tan(t_v) ** 2 - jacobian *= 2 * tan(t_v) / cos(t_v) ** 2 - new_vars.append(t_v) - domain.append([0.0, float(pi / 2)]) - - elif lower == -oo and upper != oo: - t_v = Symbol(f"t_{v.name}", real=True) - subs_map[v] = upper - tan(t_v) ** 2 - jacobian *= -2 * tan(t_v) / cos(t_v) ** 2 - new_vars.append(t_v) - domain.append([0.0, float(pi / 2)]) +# --------------------------------------------------------------------------- +# Functional compatibility helpers +# --------------------------------------------------------------------------- +def torchify( + expr, + variables=None, + limits=None, + params=None, + modules=None, + change_of_variables_method: str = "tangent", + cov_eps: float = 1e-7, +): + """Convert a SymPy object into a torch-compatible callable or ``TorchExpr``. - else: - # finite [a, b] - new_vars.append(v) - domain.append([float(lower), float(upper)]) + This functional wrapper is kept for backward compatibility with the + original ``libtorch`` API. It delegates to ``libtorch().torchify(...)`` + so the functional and object-oriented styles share the same core logic. - expr_work = expr.subs(subs_map) * jacobian + Parameters + ---------- + expr : sympy.Expr | sympy.Integral | sympy.Eq + Symbolic expression or integral-like SymPy object. + variables : list[sympy.Symbol] or None, optional + Variables passed through to ``libtorch.torchify``. + limits : list[tuple] or None, optional + Integration limits passed through to ``libtorch.torchify``. + params : list[sympy.Symbol] or None, optional + Additional symbolic parameters. + modules : list[dict] or None, optional + Custom ``lambdify`` modules mapping. - if params is None: - params = sorted( - list(expr_work.free_symbols - set(new_vars)), - key=lambda s: s.name, - ) - else: - params = list(params) + Returns + ------- + TorchExpr | callable + Same return contract as ``libtorch().torchify(...)``. - # Lambdify: new_vars first, then params - arglist = tuple(new_vars + params) - func = lambdify(arglist, expr_work, modules=modules) + USAGE + ===== + 1. Functional style kept in parallel with the OOP API:: - return TorchExpr(func=func, domain=domain, dim=len(new_vars), n_params=len(params)) + ````python + from libphysics.libtorch import torchify + from sympy import symbols, Integral, exp, oo + x = symbols("x", real=True) + integral_expr = Integral(exp(-x**2), (x, -oo, oo)) -# --------------------------------------------------------------------------- -# torchquad_integrate — cheap work, called MANY times -# --------------------------------------------------------------------------- -def torchquad_integrate(texpr, params_values=None, method=None, N=21): + texpr = torchify(integral_expr) + ```` """ - Numerically integrate a ``TorchExpr`` using torchquad. + return libtorch().torchify( + expr, + variables=variables, + limits=limits, + params=params, + modules=modules, + change_of_variables_method=change_of_variables_method, + cov_eps=cov_eps, + ) + + +def torchquad_integrate(texpr: TorchExpr, params_values=None, method=None, N: int = 21): + """Numerically integrate a ``TorchExpr`` using torchquad. + + This functional wrapper is kept for backward compatibility with the + original ``libtorch`` API. New code can call + :meth:`TorchExpr.torchquad_integrate` directly. Parameters ---------- texpr : TorchExpr - Object returned by ``torchify(..., limits=...)``. - params_values : list or None - Numerical values for the parameters (same order as *params* in torchify). - method : torchquad integrator or None - Defaults to ``Simpson()``. - N : int - Integrator resolution (torchquad's *N* parameter). + Object returned by ``libtorch().torchify(..., limits=...)``. + params_values : list | tuple | torch.Tensor | None, optional + Numerical parameter values in the same order expected by + ``texpr.func``. + method : torchquad integrator instance or None, optional + If ``None``, a default ``Simpson`` integrator is created. + N : int, optional + Resolution parameter passed to torchquad. Returns ------- re, im : torch.Tensor Real and imaginary parts of the integral. - ------- - Usage: - re, im = torchquad_integrate(texpr, params_values=[1, 1, 1, 1], N=121) - """ - from torchquad import Simpson - if method is None: - method = Simpson() + USAGE + ===== + 1. Functional style kept in parallel with the OOP API:: - param_vals = list(params_values) if params_values else [] + ````python + from libphysics.libtorch import libtorch, torchquad_integrate + from sympy import symbols, Integral, exp, oo + x = symbols("x", real=True) + integral_expr = Integral(exp(-x**2), (x, -oo, oo)) - def f(d): - vals = texpr.func(*[d[:, i] for i in range(d.shape[1])], *param_vals) - return torch.as_tensor(vals, device=d.device) - - res = method.integrate(f, dim=texpr.dim, N=N, integration_domain= texpr.domain) - return (res.real if torch.is_complex(res) else res), (res.imag if torch.is_complex(res) else torch.zeros_like(res)) - + lt = libtorch() + texpr = lt.torchify(integral_expr) + re, im = torchquad_integrate(texpr, N=121) + ```` + """ + return texpr.torchquad_integrate(params_values=params_values, method=method, N=N) def _simpson_weights_1d(a: float, b: float, N: int, *, device=None, dtype=None) -> torch.Tensor: - """Return Simpson weights including the dx/3 scaling (shape: [N]). - this is used to build the full tensor-product weights for torchquad_integrate. - """ - if N < 3 or (N % 2) == 0: - raise ValueError(f"Simpson rule requires odd N>=3; got N={N}") - dx = (b - a) / (N - 1) - w = torch.ones(N, device=device, dtype=dtype) - # 1, 4, 2, 4, ..., 2, 4, 1 this is the simppson pattern - w[1:-1:2] = 4 - w[2:-1:2] = 2 - w = w * (dx / 3.0) - return w + """Backward-compatible Simpson weights wrapper.""" + _coords, weights = TorchExpr._rule_simpson(a, b, N, device=device, dtype=dtype) + return weights def _is_scalar_like(x) -> bool: - """True for Python numbers / 0-d tensors.""" - if torch.is_tensor(x): - return x.ndim == 0 - return isinstance(x, (int, float, complex, np.number)) + """Functional wrapper around ``TorchExpr._is_scalar_like``.""" + return TorchExpr._is_scalar_like(x) def _normalize_params_values(params_values, n_params: int, *, device, dtype): - """Normalize params into (B, n_params) plus batch_shape. - - Accepts: - - tensor of shape (..., n_params) - - list/tuple length n_params containing tensors and/or scalars + """Functional wrapper around ``TorchExpr._normalize_params_values``.""" + return TorchExpr._normalize_params_values(params_values, n_params, device=device, dtype=dtype) - Scalars are broadcast to the inferred batch_shape. - """ - if params_values is None: - if n_params != 0: - raise ValueError(f"Expected {n_params} params; got None") - return torch.empty((1, 0), device=device, dtype=dtype), tuple() - - # Case A: tensor (..., n_params) - if torch.is_tensor(params_values): - params_tensor = params_values - if params_tensor.shape[-1] != n_params: - raise ValueError( - f"params_values last dimension must be n_params={n_params}; got shape={tuple(params_tensor.shape)}" - ) - batch_shape = tuple(params_tensor.shape[:-1]) - B = int(np.prod(batch_shape)) if batch_shape else 1 - params_flat = params_tensor.reshape(B, n_params).to(device=device, dtype=dtype) - return params_flat, batch_shape - # Case B: list/tuple of params (scalars and/or tensors) - if not isinstance(params_values, (list, tuple)): - # single scalar treated as 1 param - params_values = [params_values] - - if len(params_values) != n_params: - raise ValueError(f"Expected {n_params} params; got {len(params_values)}") - - # Infer batch_shape from the first non-scalar param (or scalars-only => ()) - batch_shape = None - for p in params_values: - if torch.is_tensor(p) and p.ndim > 0: - batch_shape = tuple(p.shape) - break - if not _is_scalar_like(p): - pt = torch.as_tensor(p) - if pt.ndim > 0: - batch_shape = tuple(pt.shape) - break - if batch_shape is None: - batch_shape = tuple() - - # Broadcast all params to batch_shape, then stack - params_list = [] - for p in params_values: - pt = torch.as_tensor(p, device=device) - if pt.ndim == 0: - if batch_shape: - pt = pt.expand(batch_shape) - else: - if tuple(pt.shape) != batch_shape: - raise ValueError( - f"All parameter tensors must have the same shape; got {tuple(pt.shape)} vs {batch_shape}" - ) - params_list.append(pt) - - params_tensor = torch.stack(params_list, dim=-1).to(dtype=dtype) - B = int(np.prod(batch_shape)) if batch_shape else 1 - params_flat = params_tensor.reshape(B, n_params) - return params_flat, batch_shape - - -def torch_integrate_batched_simpson( +def torch_integrate_batched( texpr: TorchExpr, - params_values, + params_values=None, *, + method: str | Callable = "simpson", N: int = 121, chunk_size_params: int = 256, chunk_size_points: int | None = None, device=None, dtype=None, ): - """Batched tensor-product Simpson integration for any dimension. + """Batched tensor-product quadrature integration for any dimension. - This integrates a ``TorchExpr`` over its finite domain using a tensor-product - Simpson rule with *N* points per dimension. - - Supports a batch of parameter points (e.g. a 250x250 mesh) by keeping the - batch dimension and summing only over the sample grid. - - Notes - ----- - - The number of sample points grows as $N^{dim}$. For dim>3 this can become - very large quickly. Use smaller N and/or set ``chunk_size_points``. - - ``chunk_size_points`` trades more Python overhead for lower peak memory. + This functional wrapper is kept so existing code that used the + original function-based API continues to work. New code can call + :meth:`TorchExpr.torch_integrate_batched` directly. Parameters ---------- texpr : TorchExpr - Output of ``torchify(..., limits=...)``. - params_values : tensor | list/tuple - Either a tensor of shape ``(..., n_params)`` or a list/tuple of length - ``n_params`` of tensors/scalars with broadcastable batch shapes. - N : int - Odd Simpson points per dimension. - chunk_size_params : int + Object returned by ``libtorch().torchify(..., limits=...)``. + params_values : tensor | list | tuple | None + Numerical parameter values. + N : int, optional + Odd number of Simpson points per dimension. + chunk_size_params : int, optional Number of parameter points processed per chunk. - chunk_size_points : int | None - Optional number of sample points processed per chunk. + chunk_size_points : int or None, optional + Number of sample points processed per chunk. + device : torch.device or None, optional + Device used for internal tensors. + dtype : torch.dtype or None, optional + Floating dtype used for internal tensors. Returns ------- re, im : torch.Tensor - Shapes match the parameter batch shape (e.g. 250x250). - ------- - Usage: - >>> re, im = torch_integrate_batched_simpson( - ... texpr, - ... params_values=params_grid, # shape (250, 250, n_params) - ... N=21, - ... chunk_size_params=256, - ... chunk_size_points=10000) + Real and imaginary parts of the integral. + + USAGE + ===== + 1. Functional style kept in parallel with the OOP API:: + + ````python + from libphysics.libtorch import libtorch, torch_integrate_batched_simpson + from sympy import symbols, Integral, exp, oo + + x = symbols("x", real=True) + integral_expr = Integral(exp(-x**2), (x, -oo, oo)) + + lt = libtorch() + texpr = lt.torchify(integral_expr) + re, im = torch_integrate_batched_simpson(texpr, params_values=None, N=121) + ```` """ - if texpr.dim <= 0: - raise ValueError(f"texpr.dim must be positive; got dim={texpr.dim}") - if texpr.n_params < 0: - raise ValueError(f"texpr.n_params must be non-negative; got n_params={texpr.n_params}") - if len(texpr.domain) != texpr.dim: - raise ValueError(f"texpr.domain must have length dim={texpr.dim}; got {len(texpr.domain)}") - - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - if dtype is None: - dtype = torch.float32 if device.type == "cuda" else torch.float64 - - if chunk_size_params <= 0: - raise ValueError(f"chunk_size_params must be positive; got {chunk_size_params}") - if chunk_size_points is not None and chunk_size_points <= 0: - raise ValueError(f"chunk_size_points must be positive; got {chunk_size_points}") - - params_flat, batch_shape = _normalize_params_values( - params_values, - texpr.n_params, + return texpr.torch_integrate_batched( + params_values=params_values, + method=method, + N=N, + chunk_size_params=chunk_size_params, + chunk_size_points=chunk_size_points, device=device, dtype=dtype, ) - B = int(params_flat.shape[0]) - - # Build per-dimension coordinates and weights - coords_1d = [] - weights_1d = [] - for (a, b) in texpr.domain: - a = float(a) - b = float(b) - coords_1d.append(torch.linspace(a, b, N, device=device, dtype=dtype)) - weights_1d.append(_simpson_weights_1d(a, b, N, device=device, dtype=dtype)) - - # Build full grid points (P, dim) and weights (P,) - # This is the most straightforward approach; chunk_size_points can reduce peak memory. - grid = torch.cartesian_prod(*coords_1d) # (P, dim) - w = weights_1d[0] - for wi in weights_1d[1:]: - w = torch.kron(w, wi) - P = int(grid.shape[0]) - - if chunk_size_points is None: - chunk_size_points = P - - re_out = torch.empty(B, device=device, dtype=dtype) - im_out = torch.empty(B, device=device, dtype=dtype) - - # Evaluate in chunks over params and points - for start_p in range(0, B, chunk_size_params): - stop_p = min(B, start_p + chunk_size_params) - pc = params_flat[start_p:stop_p, :] # (Bc, n_params) - param_args = [pc[:, i].unsqueeze(0) for i in range(texpr.n_params)] - - re_acc = torch.zeros(stop_p - start_p, device=device, dtype=dtype) - im_acc = torch.zeros(stop_p - start_p, device=device, dtype=dtype) - - for start_x in range(0, P, chunk_size_points): - stop_x = min(P, start_x + chunk_size_points) - g = grid[start_x:stop_x, :] # (Px, dim) - ww = w[start_x:stop_x].unsqueeze(1) # (Px, 1) - - var_args = [g[:, i].unsqueeze(1) for i in range(texpr.dim)] - vals = texpr.func(*var_args, *param_args) - vals = torch.as_tensor(vals, device=device) - - # Expect (Px, Bc) after broadcasting. - if vals.ndim == 1: - vals = vals.unsqueeze(1) - elif vals.ndim == 2 and vals.shape[0] == (stop_p - start_p) and vals.shape[1] == (stop_x - start_x): - # Sometimes lambdify might return (Bc, Px) - vals = vals.t().contiguous() - - if vals.shape[0] != (stop_x - start_x) or vals.shape[1] != (stop_p - start_p): - raise RuntimeError( - f"Unexpected integrand output shape {tuple(vals.shape)}; expected ({stop_x - start_x}, {stop_p - start_p})" - ) - if torch.is_complex(vals): - re_acc += (ww * vals.real).sum(dim=0) - im_acc += (ww * vals.imag).sum(dim=0) - else: - re_acc += (ww * vals).sum(dim=0) - # im stays zero - re_out[start_p:stop_p] = re_acc - im_out[start_p:stop_p] = im_acc +def torch_integrate_batched_simpson( + texpr: TorchExpr, + params_values=None, + *, + N: int = 121, + chunk_size_params: int = 256, + chunk_size_points: int | None = None, + device=None, + dtype=None, +): + """Backward-compatible wrapper for ``torch_integrate_batched(method='simpson')``.""" + return texpr.torch_integrate_batched( + params_values=params_values, + method="simpson", + N=N, + chunk_size_params=chunk_size_params, + chunk_size_points=chunk_size_points, + device=device, + dtype=dtype, + ) + - re_out = re_out.reshape(batch_shape) - im_out = im_out.reshape(batch_shape) - return re_out, im_out +lt = libtorch()