sympy

sympy

Tools that facilitate in building sympy expressions.

class UnevaluatedExpression(*args, name: str | None = None, **hints)[source]

Bases: Expr

Base class for expression classes with an evaluate() method.

Deriving from Expr allows us to keep expression trees condense before unfolding them with their doit method. This allows us to:

  1. condense the LaTeX representation of an expression tree by providing a custom _latex() method.

  2. overwrite its printer methods (see NumPyPrintable and e.g. [TR-001] Custom lambdification).

The UnevaluatedExpression base class makes implementations of its derived classes more secure by enforcing the developer to provide implementations for these methods, so that SymPy mechanisms work correctly. Decorators like implement_expr() and implement_doit_method() provide convenient means to implement the missing methods.

static __new__(cls: type[DecoratedClass], *args, name: str | None = None, **hints) DecoratedClass[source]

Constructor for a class derived from UnevaluatedExpression.

This __new__() method correctly sets the args, assumptions etc. Overwrite it in order to further specify its signature. The function create_expression() can be used in its implementation, like so:

>>> class MyExpression(UnevaluatedExpression):
...    def __new__(
...        cls, x: sp.Symbol, y: sp.Symbol, n: int, **hints
...    ) -> "MyExpression":
...        return create_expression(cls, x, y, n, **hints)
...
...    def evaluate(self) -> sp.Expr:
...        x, y, n = self.args
...        return (x + y)**n
...
>>> x, y = sp.symbols("x y")
>>> expr = MyExpression(x, y, n=3)
>>> expr
MyExpression(x, y, 3)
>>> expr.evaluate()
(x + y)**3
abstract evaluate() Expr[source]

Evaluate and ‘unfold’ this UnevaluatedExpression by one level.

>>> from ampform.dynamics import BreakupMomentumSquared
>>> issubclass(BreakupMomentumSquared, UnevaluatedExpression)
True
>>> s, m1, m2 = sp.symbols("s m1 m2")
>>> expr = BreakupMomentumSquared(s, m1, m2)
>>> expr
BreakupMomentumSquared(s, m1, m2)
>>> expr.evaluate()
(s - (m1 - m2)**2)*(s - (m1 + m2)**2)/(4*s)
>>> expr.doit(deep=False)
(s - (m1 - m2)**2)*(s - (m1 + m2)**2)/(4*s)

Note

When decorating this class with implement_doit_method(), its evaluate() method is equivalent to doit() with deep=False.

_latex(printer: LatexPrinter, *args) str[source]

Provide a mathematical Latex representation for pretty printing.

>>> from ampform.dynamics import BreakupMomentumSquared
>>> issubclass(BreakupMomentumSquared, UnevaluatedExpression)
True
>>> s, m1 = sp.symbols("s m1")
>>> expr = BreakupMomentumSquared(s, m1, m1)
>>> print(sp.latex(expr))
q^2\left(s\right)
>>> print(sp.latex(expr.doit()))
- m_{1}^{2} + \frac{s}{4}
class NumPyPrintable(*args)[source]

Bases: Expr

Expr class that can lambdify to NumPy code.

This interface for classes that derive from sympy.Expr enforce the implementation of a _numpycode() method in case the class does not correctly lambdify() to NumPy code. For more info on SymPy printers, see Printing.

Several computational frameworks try to converge their interface to that of NumPy. See for instance TensorFlow’s NumPy API and jax.numpy. This fact is used in TensorWaves to lambdify() SymPy expressions to these different backends with the same lambdification code.

Note

This interface differs from UnevaluatedExpression in that it should not implement an evaluate() (and therefore a doit()) method.

Warning

The implemented _numpycode() method should countain as little SymPy computations as possible. Instead, it should get most information from its construction args, so that SymPy can use printer tricks like cse(), prior expanding with doit(), and other simplifications that can make the generated code shorter. An example is the BoostZMatrix class, which takes \(\beta\) as input instead of the FourMomentumSymbol from which \(\beta\) is computed.

abstract _numpycode(printer: NumPyPrinter, *args) str[source]

Lambdify this NumPyPrintable class to NumPy code.

DecoratedClass

TypeVar for decorators like implement_doit_method().

alias of TypeVar(‘DecoratedClass’, bound=UnevaluatedExpression)

implement_expr(n_args: int) Callable[[type[DecoratedClass]], type[DecoratedClass]][source]

Decorator for classes that derive from UnevaluatedExpression.

Implement a __new__() and doit() method for a class that derives from Expr (via UnevaluatedExpression).

implement_new_method(n_args: int) Callable[[type[DecoratedClass]], type[DecoratedClass]][source]

Implement UnevaluatedExpression.__new__() on a derived class.

Implement a __new__() method for a class that derives from Expr (via UnevaluatedExpression).

implement_doit_method(decorated_class: type[DecoratedClass]) type[DecoratedClass][source]

Implement doit() method for an UnevaluatedExpression class.

Implement a doit() method for a class that derives from Expr (via UnevaluatedExpression). A doit() method is an extension of an evaluate() method in the sense that it can work recursively on deeper expression trees.

DecoratedExpr

TypeVar for decorators like make_commutative().

alias of TypeVar(‘DecoratedExpr’, bound=Expr)

make_commutative(decorated_class: type[DecoratedExpr]) type[DecoratedExpr][source]

Set commutative and ‘extended real’ assumptions on expression class.

See also

Assumptions

create_expression(cls: type[DecoratedExpr], *args, evaluate: bool = False, name: str | None = None, **kwargs) DecoratedExpr[source]

Helper function for implementing UnevaluatedExpression.__new__.

create_symbol_matrix(name: str, m: int, n: int) MutableDenseMatrix[source]

Create a Matrix with symbols as elements.

The MatrixSymbol has some issues when one is interested in the elements of the matrix. This function instead creates a Matrix where the elements are Indexed instances.

To convert these Indexed instances to a Symbol, use symplot.substitute_indexed_symbols().

>>> create_symbol_matrix("A", m=2, n=3)
Matrix([
[A[0, 0], A[0, 1], A[0, 2]],
[A[1, 0], A[1, 1], A[1, 2]]])
class PoolSum(expression, *indices: tuple[Symbol, Iterable[Basic]], **hints)[source]

Bases: UnevaluatedExpression

Sum over indices where the values are taken from a domain set.

>>> i, j, m, n = sp.symbols("i j m n")
>>> expr = PoolSum(i**m + j**n, (i, (-1, 0, +1)), (j, (2, 4, 5)))
>>> expr
PoolSum(i**m + j**n, (i, (-1, 0, 1)), (j, (2, 4, 5)))
>>> print(sp.latex(expr))
\sum_{i=-1}^{1} \sum_{j\in\left\{2,4,5\right\}}{i^{m} + j^{n}}
>>> expr.doit()
3*(-1)**m + 3*0**m + 3*2**n + 3*4**n + 3*5**n + 3
property expression: Expr
property indices: list[tuple[Symbol, tuple[Float, ...]]]
property free_symbols: set[Basic]

Return from the atoms of self those which are free symbols.

For most expressions, all symbols are free symbols. For some classes this is not true. e.g. Integrals use Symbols for the dummy variables which are bound variables, so Integral has a method to return all symbols except those. Derivative keeps track of symbols with respect to which it will perform a derivative; those are bound variables, too, so it has its own free_symbols method.

Any other method that uses bound variables should implement a free_symbols method.

cleanup() Expr | PoolSum[source]

Remove redundant summations, like indices with one or no value.

>>> x, i = sp.symbols("x i")
>>> PoolSum(x**i, (i, [0, 1, 2])).cleanup().doit()
x**2 + x + 1
>>> PoolSum(x, (i, [0, 1, 2])).cleanup()
x
>>> PoolSum(x).cleanup()
x
>>> PoolSum(x**i, (i, [0])).cleanup()
1

Submodules and Subpackages