Einstein summation convention.
ebcc.util.einsumfunc.CONTRACTION_METHOD = 'backend'
module-attribute
The size of the contraction to fall back on the backend.
ebcc.util.einsumfunc.BACKEND_EINSUM_SIZE = 1000
module-attribute
The size of the contraction to let the backend optimize.
ebcc.util.einsumfunc.BACKEND_OPTIMIZE_SIZE = 100
module-attribute
Symbols used in einsum-like functions.
ebcc.util.einsumfunc.EinsumOperandError
Bases: ValueError
Exception for invalid inputs to einsum
.
ebcc.util.einsumfunc.einsum(*operands, alpha=1.0, beta=0.0, out=None, contract=None, optimize=True)
Evaluate an Einstein summation convention on the operands.
Using the Einstein summation convention, many common
multi-dimensional, linear algebraic array operations can be
represented in a simple fashion. In implicit mode einsum
computes these values.
In explicit mode, einsum
provides further flexibility to compute
other array operations that might not be considered classical
Einstein summation operations, by disabling, or forcing summation
over specified subscript labels.
See the numpy.einsum
documentation for clarification.
Parameters: |
|
---|
Returns: |
|
---|
Source code in ebcc/util/einsumfunc.py
def einsum(
*operands: OperandType,
alpha: T = 1.0, # type: ignore[assignment]
beta: T = 0.0, # type: ignore[assignment]
out: Optional[NDArray[T]] = None,
contract: Optional[Callable[..., NDArray[T]]] = None,
optimize: bool = True,
) -> NDArray[T]:
"""Evaluate an Einstein summation convention on the operands.
Using the Einstein summation convention, many common
multi-dimensional, linear algebraic array operations can be
represented in a simple fashion. In *implicit* mode `einsum`
computes these values.
In *explicit* mode, `einsum` provides further flexibility to compute
other array operations that might not be considered classical
Einstein summation operations, by disabling, or forcing summation
over specified subscript labels.
See the `numpy.einsum` documentation for clarification.
Args:
operands: Any valid input to `numpy.einsum`.
alpha: Scaling factor for the contraction.
beta: Scaling factor for the output.
out: If provided, the calculation is done into this array.
contract: The function to use for contraction.
optimize: If `True`, use the `numpy.einsum_path` to optimize the contraction.
Returns:
The calculation based on the Einstein summation convention.
"""
# Parse the kwargs
inp, outs, args = _parse_einsum_input(list(operands)) # type: ignore
subscript = "%s->%s" % (inp, outs)
# Get the contraction function
if contract is None:
contract = {
"backend": _contract_backend,
"ttgt": _contract_ttgt,
"tblis": _contract_tblis,
}[CONTRACTION_METHOD.lower()]
# Perform the contraction
if not len(args):
raise ValueError("No input operands")
elif len(args) == 1:
# If it's just a transpose, use numpy
res = _transpose_backend(subscript, args[0], alpha=alpha, beta=beta, out=out)
elif len(args) == 2:
# If it's a single contraction, call the backend directly
res = contract(subscript, args[0], args[1], alpha=alpha, beta=beta, out=out)
else:
# If it's a chain of contractions, use the path optimizer
args = list(args)
path_kwargs = dict(optimize=optimize, einsum_call=True)
contractions = np.einsum_path(subscript, *args, **path_kwargs)[1]
for contraction in contractions:
inds, idx_rm, einsum_str, remain = list(contraction[:4])
contraction_args = [args.pop(x) for x in inds] # type: ignore
if alpha != 1.0 or beta != 0.0:
raise NotImplementedError("Scaling factors not supported for >2 arguments")
if len(contraction_args) == 1:
a = contraction_args[0]
res = _transpose_backend(
einsum_str, a, alpha=types[float](1.0), beta=types[float](0.0), out=None
)
else:
a, b = contraction_args
res = contract(
einsum_str, a, b, alpha=types[float](1.0), beta=types[float](0.0), out=None
)
args.append(res)
return res
ebcc.util.einsumfunc.dirsum(*operands)
Direct sum of arrays.
Follows the numpy.einsum
input conventions.
Parameters: |
|
---|
Returns: |
|
---|
Source code in ebcc/util/einsumfunc.py
def dirsum(*operands: Union[str, tuple[int, ...], NDArray[T]]) -> NDArray[T]:
"""Direct sum of arrays.
Follows the `numpy.einsum` input conventions.
Args:
operands: Any valid input to `numpy.einsum`.
Returns:
The direct sum of the arrays.
"""
# Parse the input
input_str, output_str, input_arrays = _parse_einsum_input(list(operands))
input_chars = input_str.split(",")
for i, (chars, array) in enumerate(zip(input_chars, input_arrays)):
if len(chars) != array.ndim:
raise ValueError(f"Dimension mismatch for array {i}.")
if len(set(chars)) != len(chars):
unique_chars = "".join(set(chars))
array = einsum(f"{chars}->{unique_chars}", array)
input_chars[i] = unique_chars
if i == 0:
res = array
else:
shape = res.shape + (1,) * array.ndim
res = np.reshape(res, shape) + array
# Reshape the output
res = einsum(f"{''.join(input_chars)}->{output_str}", res)
return res