# This code is part of Qiskit.
#
# (C) Copyright IBM 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
"""
Array polynomial.
"""
from typing import List, Optional, Callable, Tuple, Union
from copy import copy
from itertools import product
from multiset import Multiset
import numpy as np
from numpy.typing import DTypeLike
from qiskit import QiskitError
from qiskit_dynamics import DYNAMICS_NUMPY as unp
from qiskit_dynamics import DYNAMICS_NUMPY_ALIAS as numpy_alias
from qiskit_dynamics.arraylias.alias import (
_preferred_lib,
_numpy_multi_dispatch,
ArrayLike,
_isArrayLike,
)
from qiskit_dynamics.perturbation.multiset_utils import (
_validate_non_negative_ints,
_get_all_submultisets,
_sorted_multisets,
_submultisets_and_complements,
_multiset_to_sorted_list,
)
from qiskit_dynamics.perturbation.custom_binary_op import _CustomBinaryOp
try:
import jax.numpy as jnp
except ImportError:
pass
[docs]
class ArrayPolynomial:
r"""A polynomial with array-valued coefficients.
This class represents a multi-variable function of the form:
.. math::
f(c_1, \dots, c_r) = A_\emptyset + \sum_{I \in S} c_I A_I,
where in the above:
- :math:`S` is a finite set of multisets indicating non-zero monomial terms,
- For a given multiset of non-negative integers :math:`I=(i_1, \dots, i_k)`,
:math:`c_I = c_{i_1} \times \dots \times c_{i_k}`, and
- The :math:`A_I` are arrays of the same shape, indexed by the first dimension.
See the :ref:`multiset and power series notation section <multiset power series>` of the
perturbation review for an explanation of the multiset notation.
An :class:`.ArrayPolynomial` is instantiated with the arguments:
- ``constant_term`` specifying the array :math:`A_\emptyset`.
- ``array_coefficients`` specifying a list of the arrays :math:`A_I`, or as a single array
whose first index lists the :math:`A_I`,
- ``monomial_labels`` specifying the set :math:`S` as a list of ``Multiset`` instances
ordered in correspondence with ``array_coefficients``.
For example, the :class:`.ArrayPolynomial` corresponding to the mathematical polynomial
.. math::
f(c_0, c_1) = A_\emptyset
+ c_{(0)} A_{(0)} + c_{(0, 1)}A_{(0, 1)} + c_{(1, 1)}A_{(1, 1)}
for arrays :math:`A_\emptyset, A_{(0)}, A_{(0, 1)}, A_{(1, 1)}` stored in variables ``A_c``,
``A0``, ``A01``, and ``A11`` can be instantiated with
.. code-block:: python
ap = ArrayPolynomial(
constant_term = A_c
array_coefficients=[A0, A01, A11],
monomial_labels=[Multiset({0: 1}), Multiset({0: 1, 1: 1}), Multiset({1: 2})]
)
Once instantiated, the polynomial can be evaluated on an array of variable values, e.g.
.. code-block:: python
c = np.array([c0, c1])
ap(c) # polynomial evaluated on variables
:class:`.ArrayPolynomial` supports some array properties, e.g. ``ap.shape`` and ``ap.ndim``
return the shape and number of dimensions of the output of the polynomial. Some array methods
are also supported, such as ``transpose`` and ``trace``, and their output produces a new
:class:`.ArrayPolynomial` which evaluates to the array one would obtain by first evaluating the
original, then calling the array method. E.g.
.. code-block:: python
ap2 = ap1.transpose()
ap2(c) == ap1(c).transpose()
Finally, :class:`.ArrayPolynomial` supports algebraic operations, e.g.
.. code-block:: python
ap3 = ap1 @ ap2
ap3(c) == ap1(c) @ ap2(c)
It also has specialized algebraic methods that perform algebraic operations while "ignoring"
terms. E.g., for two instances ``ap1`` and ``ap2``, the call
.. code-block:: python
ap1.matmul(ap2, monomial_filter=lambda x: len(x) <= 3)
is similar to ``ap1 @ ap2``, but will result in an :class:`.ArrayPolynomial` in which all terms
of degree larger than ``3`` will not be included in the results.
"""
__array_priority__ = 20
def __init__(
self,
constant_term: Optional[ArrayLike] = None,
array_coefficients: Optional[ArrayLike] = None,
monomial_labels: Optional[List[Multiset]] = None,
array_library: Optional[str] = None,
):
"""Construct a multivariable matrix polynomial.
Args:
constant_term: An array representing the constant term of the polynomial.
array_coefficients: A 3d array representing a list of array coefficients.
monomial_labels: A list of multisets with non-negative integer entries of the same
length as ``array_coefficients`` indicating the monomial coefficient for each
corresponding ``array_coefficients``.
array_library: Array library for stored array terms.
Raises:
QiskitError: If insufficient information is supplied to define an ArrayPolynomial,
or if monomial labels contain anything other than non-negative integers.
"""
if array_coefficients is None and constant_term is None:
raise QiskitError(
"At least one of array_coefficients and constant_term must be specified."
)
if monomial_labels is not None:
self._monomial_labels = [Multiset(m) for m in monomial_labels]
for m in self._monomial_labels:
_validate_non_negative_ints(m)
else:
self._monomial_labels = []
self._array_coefficients = None
if array_coefficients is not None:
self._array_coefficients = numpy_alias(like=array_library).asarray(array_coefficients)
self._constant_term = None
if constant_term is not None:
self._constant_term = numpy_alias(like=array_library).asarray(constant_term)
# there is redundancy here as both perform pre-computation on self._monomial_labels
self._compute_monomials = _get_monomial_compute_function(self._monomial_labels)
self._compute_monomials_jax = _get_monomial_compute_function_jax(self._monomial_labels)
@property
def monomial_labels(self) -> Union[List, None]:
"""The monomial labels corresponding to non-constant terms."""
return self._monomial_labels
@property
def array_coefficients(self) -> Union[ArrayLike, None]:
"""The array coefficients for non-constant terms."""
return self._array_coefficients
@property
def constant_term(self) -> Union[ArrayLike, None]:
"""The constant term."""
return self._constant_term
[docs]
def compute_monomials(self, c: ArrayLike) -> ArrayLike:
"""Vectorized computation of all scalar monomial terms in the polynomial as specified by
``self.monomial_labels``.
Args:
c: Array of variables.
Returns:
Array of all monomial terms ordered according to ``self.monomial_labels``.
"""
lib = _preferred_lib(c, self.constant_term, self.array_coefficients)
if "jax" in lib:
return self._compute_monomials_jax(c)
return self._compute_monomials(c)
@property
def shape(self) -> Tuple[int]:
"""Shape of the arrays in the polynomial."""
if self._constant_term is not None:
return self._constant_term.shape
else:
return self._array_coefficients.shape[1:]
@property
def ndim(self) -> int:
"""Number of dimensions of the coefficients of the polynomial."""
if self._constant_term is not None:
return self._constant_term.ndim
else:
return self._array_coefficients.ndim - 1
[docs]
def conj(self) -> "ArrayPolynomial":
"""Return an ArrayPolynomial that is the conjugate of self."""
constant_term = None
coefficients = None
if self._constant_term is not None:
constant_term = unp.conj(self._constant_term)
if self._array_coefficients is not None:
coefficients = unp.conj(self._array_coefficients)
return ArrayPolynomial(
array_coefficients=coefficients,
monomial_labels=copy(self._monomial_labels),
constant_term=constant_term,
)
[docs]
def transpose(self, axes: Optional[Tuple[int]] = None) -> "ArrayPolynomial":
"""Return the ArrayPolynomial when transposing all coefficients."""
constant_term = None
coefficients = None
if self._constant_term is not None:
constant_term = unp.transpose(self._constant_term, axes)
if self._array_coefficients is not None:
if axes is None:
axes = tuple(range(1, self.ndim + 1)[::-1])
else:
axes = tuple(ax + 1 for ax in axes)
axes = (0,) + axes
coefficients = unp.transpose(self._array_coefficients, axes)
return ArrayPolynomial(
array_coefficients=coefficients,
monomial_labels=copy(self._monomial_labels),
constant_term=constant_term,
)
[docs]
def trace(
self,
offset: Optional[int] = 0,
axis1: Optional[int] = 0,
axis2: Optional[int] = 1,
dtype: Optional[DTypeLike] = None,
) -> "ArrayPolynomial":
"""Take the trace of the coefficients."""
if self.ndim < 2:
raise QiskitError("ArrayPolynomial.trace() requires ArrayPolynomial.ndim at least 2.")
constant_term = None
coefficients = None
if self._constant_term is not None:
constant_term = unp.trace(
self._constant_term, offset=offset, axis1=axis1, axis2=axis2, dtype=dtype
)
if self._array_coefficients is not None:
coefficients = unp.trace(
self._array_coefficients,
offset=offset,
axis1=axis1 + 1,
axis2=axis2 + 1,
dtype=dtype,
)
return ArrayPolynomial(
array_coefficients=coefficients,
monomial_labels=copy(self._monomial_labels),
constant_term=constant_term,
)
[docs]
def sum(self, axis: Optional[Union[int, Tuple[int]]] = None, dtype: Optional[DTypeLike] = None):
"""Perform a sum on the coefficients."""
constant_term = None
coefficients = None
# constant term can be handled normally
if self.constant_term is not None:
constant_term = self.constant_term.sum(axis=axis, dtype=dtype)
# axis must be shifted for array coefficients
if self.array_coefficients is not None:
if self.ndim == 0 and axis is None:
coefficients = unp.array(self.array_coefficients, dtype=dtype)
else:
if axis is None:
axis = tuple(k for k in range(1, self.ndim + 1))
elif isinstance(axis, int):
axis = axis + 1
elif isinstance(axis, tuple):
axis = tuple(k + 1 for k in axis)
coefficients = self.array_coefficients.sum(axis=axis, dtype=dtype)
return ArrayPolynomial(
array_coefficients=coefficients,
monomial_labels=copy(self._monomial_labels),
constant_term=constant_term,
)
@property
def real(self) -> "ArrayPolynomial":
"""Return the real part of self."""
constant_term = None
array_coefficients = None
if self.constant_term is not None:
constant_term = self.constant_term.real
if self.array_coefficients is not None:
array_coefficients = self.array_coefficients.real
return ArrayPolynomial(
array_coefficients=array_coefficients,
monomial_labels=copy(self._monomial_labels),
constant_term=constant_term,
)
[docs]
def add(
self,
other: Union["ArrayPolynomial", ArrayLike],
monomial_filter: Optional[Callable] = None,
) -> "ArrayPolynomial":
"""Add two polynomials with bounds on which terms to keep.
Optionally, a function ``monomial_filter`` can be provided to limit which monomials appear
in the output. It must accept as input a ``Multiset`` and return a ``bool``, and a term with
label given by ``multiset`` will be included only if ``monomial_filter(multiset) == True``,
and will not be computed if ``monomial_filter(multiset) == False``.
Args:
other: Other to add to self.
monomial_filter: Function determining which terms to compute and keep.
Returns:
ArrayPolynomial achieved by adding both self and other.
Raises:
QiskitError: if other cannot be cast as an ArrayPolynomial.
"""
if _isArrayLike(other):
other = ArrayPolynomial(constant_term=other)
if isinstance(other, ArrayPolynomial):
return _array_polynomial_addition(self, other, monomial_filter=monomial_filter)
raise QiskitError(
"Only types castable as an ArrayPolynomial can be added to an ArrayPolynomial."
)
[docs]
def matmul(
self,
other: Union["ArrayPolynomial", ArrayLike],
monomial_filter: Optional[Callable] = None,
) -> "ArrayPolynomial":
"""Matmul self @ other with bounds on which terms to keep.
Optionally, a function ``monomial_filter`` can be provided to limit which monomials appear
in the output. It must accept as input a ``Multiset`` and return a ``bool``, and a term with
label given by ``multiset`` will be included only if ``monomial_filter(multiset) == True``,
and will not be computed if ``monomial_filter(multiset) == False``.
Args:
other: Other to add to self.
monomial_filter: Function determining which terms to compute and keep.
Returns:
ArrayPolynomial achieved by matmul of self and other.
Raises:
QiskitError: if other cannot be cast as an ArrayPolynomial.
"""
if _isArrayLike(other):
other = ArrayPolynomial(constant_term=other)
if isinstance(other, ArrayPolynomial):
return _array_polynomial_distributive_binary_op(
self, other, lambda A, B: A @ B, monomial_filter=monomial_filter
)
raise QiskitError(f"Type {type(other)} not supported by ArrayPolynomial.matmul.")
[docs]
def mul(
self,
other: Union["ArrayPolynomial", ArrayLike],
monomial_filter: Optional[Callable] = None,
) -> "ArrayPolynomial":
"""Entrywise multiplication of two ArrayPolynomials with bounds on which terms to keep.
Optionally, a function ``monomial_filter`` can be provided to limit which monomials appear
in the output. It must accept as input a ``Multiset`` and return a ``bool``, and a term with
label given by ``multiset`` will be included only if ``monomial_filter(multiset) == True``,
and will not be computed if ``monomial_filter(multiset) == False``.
Args:
other: Other to add to self.
monomial_filter: Function determining which terms to compute and keep.
Returns:
ArrayPolynomial achieved by matmul of self and other.
Raises:
QiskitError: if other cannot be cast as an ArrayPolynomial.
"""
if _isArrayLike(other):
other = ArrayPolynomial(constant_term=other)
if isinstance(other, ArrayPolynomial):
return _array_polynomial_distributive_binary_op(
self, other, lambda A, B: A * B, monomial_filter=monomial_filter
)
raise QiskitError(f"Type {type(other)} not supported by ArrayPolynomial.mul.")
def __add__(self, other: Union["ArrayPolynomial", ArrayLike]) -> "ArrayPolynomial":
"""Dunder method for addition of two ArrayPolynomials."""
return self.add(other)
def __radd__(self, other: Union["ArrayPolynomial", ArrayLike]) -> "ArrayPolynomial":
"""Dunder method for right-addition of two ArrayPolynomials."""
return self.add(other)
def __neg__(self) -> "ArrayPolynomial":
constant_term = None
if self.constant_term is not None:
# pylint: disable=invalid-unary-operand-type
constant_term = -self.constant_term
array_coefficients = None
if self.array_coefficients is not None:
# pylint: disable=invalid-unary-operand-type
array_coefficients = -self.array_coefficients
return ArrayPolynomial(
constant_term=constant_term,
monomial_labels=self.monomial_labels,
array_coefficients=array_coefficients,
)
def __sub__(self, other: Union["ArrayPolynomial", ArrayLike]) -> "ArrayPolynomial":
return self + (-other)
def __rsub__(self, other: Union["ArrayPolynomial", ArrayLike]) -> "ArrayPolynomial":
return other + (-self)
def __mul__(self, other: Union["ArrayPolynomial", ArrayLike]) -> "ArrayPolynomial":
"""Dunder method for entry-wise multiplication."""
return self.mul(other)
def __rmul__(self, other: Union["ArrayPolynomial", ArrayLike]) -> "ArrayPolynomial":
"""Dunder method for right-multiplication."""
return self.mul(other)
def __matmul__(self, other: Union["ArrayPolynomial", ArrayLike]) -> "ArrayPolynomial":
"""Dunder method for matmul."""
return self.matmul(other)
def __rmatmul__(self, other: Union["ArrayPolynomial", ArrayLike]) -> "ArrayPolynomial":
"""Dunder method for rmatmul."""
if _isArrayLike(other):
other = ArrayPolynomial(constant_term=other)
if isinstance(other, ArrayPolynomial):
return other.matmul(self)
raise QiskitError(f"Type {type(other)} not supported by ArrayPolynomial.__rmatmul__.")
def __getitem__(self, idx) -> "ArrayPolynomial":
"""Index the ArrayPolynomial similarly to an array."""
constant_term = None
array_coefficients = None
if self.constant_term is not None:
constant_term = self.constant_term[idx]
if self.array_coefficients is not None:
array_coefficients = self.array_coefficients[(slice(None),) + idx]
return ArrayPolynomial(
array_coefficients=array_coefficients,
monomial_labels=copy(self._monomial_labels),
constant_term=constant_term,
)
def __len__(self) -> int:
"""Number of terms in the polynomial."""
num_terms = 0
if self.array_coefficients is not None:
num_terms = num_terms + len(self.array_coefficients)
if self.constant_term is not None:
num_terms = num_terms + 1
return num_terms
def __call__(self, c: Optional[ArrayLike] = None) -> ArrayLike:
"""Evaluate the polynomial.
Args:
c: Array of variables.
Returns:
Value of the polynomial at c.
"""
if self._array_coefficients is not None and self._constant_term is not None:
monomials = self.compute_monomials(c)
return self._constant_term + _numpy_multi_dispatch(
self._array_coefficients, monomials, path="tensordot", axes=(0, 0)
)
elif self._array_coefficients is not None:
monomials = self.compute_monomials(c)
return _numpy_multi_dispatch(
self._array_coefficients, monomials, path="tensordot", axes=(0, 0)
)
else:
return self._constant_term
def _get_monomial_compute_function(multisets: List[Multiset]) -> Callable:
"""Construct a vectorized function for computing multivariable monomial terms indicated by
multisets.
The returned function takes in the individual variables as an array, and returns an array of
computed monomial terms in the order indicated by multisets.
The returned function is vectorized in the sense that the supplied first order terms can be
arrays.
The algorithm computes monomial terms of increasing order, recursively utilizing lower order
terms.
Args:
multisets: list of multisets.
Returns:
Callable: Vectorized function for computing monomials.
"""
if multisets is None or len(multisets) == 0:
return lambda c: None
complete_multiset_list = _get_all_submultisets(multisets)
(
first_order_terms,
first_order_range,
left_indices,
right_indices,
update_ranges,
) = _get_recursive_monomial_rule(complete_multiset_list)
multiset_len = len(complete_multiset_list)
location_list = np.array(
[complete_multiset_list.index(multiset) for multiset in multisets], dtype=int
)
def monomial_function(c):
shape = [multiset_len] + list(c.shape[1:])
mono_vec = np.empty(shape=shape, dtype=complex)
mono_vec[first_order_range[0] : first_order_range[1]] = c[first_order_terms]
for left_index, right_index, update_range in zip(
left_indices, right_indices, update_ranges
):
mono_vec[update_range[0] : update_range[1]] = (
mono_vec[left_index] * mono_vec[right_index]
)
return mono_vec[location_list]
return monomial_function
# pylint: disable=invalid-name
def _get_monomial_compute_function_jax(multisets: List) -> Callable:
"""JAX version of _get_monomial_compute_function."""
if multisets is None or len(multisets) == 0:
return lambda c: None
complete_multiset_list = _get_all_submultisets(multisets)
first_order_terms, _, left_indices, right_indices, _ = _get_recursive_monomial_rule(
complete_multiset_list
)
location_list = np.array(
[complete_multiset_list.index(multiset) for multiset in multisets], dtype=int
)
# initial function sets up required first order terms
def monomial_function_init(c):
return c[first_order_terms]
monomial_function = monomial_function_init
# function for generating next update and compositions to avoid looping reference issues
def get_next_update_func(left_index, right_index):
def update_next(mono_vec):
return jnp.append(mono_vec, mono_vec[left_index] * mono_vec[right_index], axis=0)
return update_next
def compose_functions(f, g):
def new_func(x):
return g(f(x))
return new_func
# recursively compose updates with monomial_function
for left_index, right_index in zip(left_indices, right_indices):
next_update = get_next_update_func(left_index, right_index)
monomial_function = compose_functions(monomial_function, next_update)
# return only the requested terms
def trimmed_output_function(c):
return monomial_function(c)[location_list]
return trimmed_output_function
def _get_recursive_monomial_rule(complete_multisets: List) -> Tuple:
"""Helper function for _get_monomial_compute_function and _get_monomial_compute_function_jax;
computes a representation of the algorithm for computing monomials that is used by both
functions.
complete_multisets is assumed to be closed under taking submultisets and in canonical order.
Args:
complete_multisets: Description of monomial terms.
Returns:
Tuple: Collection of lists organizing computation for both
_get_monomial_compute_function and _get_monomial_compute_function_jax.
"""
# first, construct representation of recursive rule explicitly in terms of multisets
first_order_terms = []
left_terms = []
right_terms = []
current_left = -1
current_right_list = []
current_len = 2
# convert multisets to list representation
complete_multisets = [_multiset_to_sorted_list(multiset) for multiset in complete_multisets]
for multiset in complete_multisets:
if len(multiset) == 1:
first_order_terms.append(multiset[0])
else:
if multiset[0] != current_left or len(multiset) != current_len:
current_len = len(multiset)
if current_left != -1:
left_terms.append(current_left)
right_terms.append(current_right_list)
current_left = multiset[0]
current_right_list = [multiset[1:]]
else:
current_right_list.append(multiset[1:])
# if current_left is still -1, then only first order terms exist
if current_left == -1:
return np.array(first_order_terms), [0, len(first_order_terms)], [], [], []
# proceed assuming at least one term above first order exists
# append the last one
left_terms.append(current_left)
right_terms.append(current_right_list)
# convert representation of rule in terms of multisets into one in terms of
# array indices
# set up arrays
first_order_terms = np.array(first_order_terms, dtype=int)
left_indices = []
right_indices = []
for left_term, right_term in zip(left_terms, right_terms):
left_indices.append(complete_multisets.index([left_term]))
right_index_list = []
for term in right_term:
right_index_list.append(complete_multisets.index(term))
right_indices.append(np.array(right_index_list, dtype=int))
# set up index updating ranges
first_order_range = [0, len(first_order_terms)]
update_ranges = []
current_idx = first_order_range[1]
for right_index in right_indices:
next_idx = current_idx + len(right_index)
update_ranges.append([current_idx, next_idx])
current_idx = next_idx
return (
np.array(first_order_terms),
first_order_range,
np.array(left_indices),
right_indices,
update_ranges,
)
def _array_polynomial_distributive_binary_op(
ap1: ArrayPolynomial,
ap2: ArrayPolynomial,
binary_op: Callable,
monomial_filter: Optional[Callable] = None,
) -> ArrayPolynomial:
"""Apply a distributive binary op on two array polynomials."""
# determine list of Multisets required for monomial labels, including filtering
all_multisets = []
# if no filter is provided, set to always return True
if monomial_filter is None:
def monomial_filter(_):
return True
if ap1.constant_term is not None:
for multiset in ap2.monomial_labels:
if monomial_filter(multiset) and multiset not in all_multisets:
all_multisets.append(multiset)
if ap2.constant_term is not None:
for multiset in ap1.monomial_labels:
if monomial_filter(multiset) and multiset not in all_multisets:
all_multisets.append(multiset)
for I, J in product(ap1.monomial_labels, ap2.monomial_labels):
IuJ = I + J
if monomial_filter(IuJ) and IuJ not in all_multisets:
all_multisets.append(IuJ)
all_multisets = _sorted_multisets(all_multisets)
# setup constant term
new_constant_term = None
if (
ap1.constant_term is not None
and ap2.constant_term is not None
and monomial_filter(Multiset({}))
):
new_constant_term = binary_op(ap1.constant_term, ap2.constant_term)
# return constant case
if not all_multisets:
return ArrayPolynomial(constant_term=new_constant_term)
# iteratively construct custom multiplication rule,
# temporarily treating the constant terms as index -1
operation_rule = []
for multiset in all_multisets:
rule_indices = []
if multiset in ap1.monomial_labels:
idx = ap1.monomial_labels.index(multiset)
rule_indices.append([idx, -1])
if multiset in ap2.monomial_labels:
idx = ap2.monomial_labels.index(multiset)
rule_indices.append([-1, idx])
if len(multiset) > 1:
for I, J in zip(*_submultisets_and_complements(multiset)):
if I in ap1.monomial_labels and J in ap2.monomial_labels:
rule_indices.append(
[ap1.monomial_labels.index(I), ap2.monomial_labels.index(J)]
)
# if non-empty,
if rule_indices:
operation_rule.append((np.ones(len(rule_indices)), np.array(rule_indices)))
lmats = None
if ap1.constant_term is not None:
lmats = unp.expand_dims(ap1.constant_term, 0)
else:
lmats = unp.expand_dims(unp.zeros_like(ap1.array_coefficients[0]), 0)
if ap1.array_coefficients is not None:
lmats = _numpy_multi_dispatch(lmats, ap1.array_coefficients, path="append", axis=0)
rmats = None
if ap2.constant_term is not None:
rmats = unp.expand_dims(ap2.constant_term, 0)
else:
rmats = unp.expand_dims(unp.zeros_like(ap2.array_coefficients[0]), 0)
if ap2.array_coefficients is not None:
rmats = _numpy_multi_dispatch(rmats, ap2.array_coefficients, path="append", axis=0)
custom_binary_op = _CustomBinaryOp(
operation_rule=operation_rule,
binary_op=binary_op,
index_offset=1,
)
new_array_coefficients = custom_binary_op(lmats, rmats)
return ArrayPolynomial(
array_coefficients=new_array_coefficients,
monomial_labels=all_multisets,
constant_term=new_constant_term,
)
def _array_polynomial_addition(
ap1: ArrayPolynomial,
ap2: ArrayPolynomial,
monomial_filter: Optional[Callable] = None,
) -> ArrayPolynomial:
"""Add two ArrayPolynomials."""
for a, b in zip(ap1.shape[::-1], ap2.shape[::-1]):
if not (a == 1 or b == 1 or a == b):
raise QiskitError(
"ArrayPolynomial addition requires shapes be broadcastable to eachother."
)
if monomial_filter is None:
def monomial_filter(_):
return True
# construct constant term
new_constant_term = None
# if constant term is to be included, determine what it is
if monomial_filter(Multiset({})):
if ap1.constant_term is not None and ap2.constant_term is not None:
new_constant_term = ap1.constant_term + ap2.constant_term
elif ap1.constant_term is not None:
new_constant_term = ap1.constant_term
elif ap2.constant_term is not None:
new_constant_term = ap2.constant_term
# exit early if both polynomials are constant
if ap1.array_coefficients is None and ap2.array_coefficients is None:
return ArrayPolynomial(constant_term=new_constant_term)
# construct list of admissable multisets and sort into canonical order
new_multisets = []
for multiset in ap1.monomial_labels + ap2.monomial_labels:
if monomial_filter(multiset) and multiset not in new_multisets:
new_multisets.append(multiset)
new_multisets = _sorted_multisets(new_multisets)
# construct index order mapping for each polynomial
idx1 = []
idx2 = []
for multiset in new_multisets:
if multiset in ap1.monomial_labels:
idx1.append(ap1.monomial_labels.index(multiset))
else:
idx1.append(-1)
if multiset in ap2.monomial_labels:
idx2.append(ap2.monomial_labels.index(multiset))
else:
idx2.append(-1)
# if either is empty, pad with single -1, then convert to array
idx1 = idx1 or [-1]
idx2 = idx2 or [-1]
idx1 = np.array(idx1)
idx2 = np.array(idx2)
# append zero to the coefficient arrays
array_coefficients1 = np.zeros((1,) + ap1.shape, dtype=complex)
array_coefficients2 = np.zeros((1,) + ap1.shape, dtype=complex)
if ap1.array_coefficients is not None:
array_coefficients1 = unp.append(ap1.array_coefficients, array_coefficients1, axis=0)
if ap2.array_coefficients is not None:
array_coefficients2 = unp.append(ap2.array_coefficients, array_coefficients2, axis=0)
new_coefficients = array_coefficients1[idx1] + array_coefficients2[idx2]
return ArrayPolynomial(
array_coefficients=new_coefficients,
monomial_labels=new_multisets,
constant_term=new_constant_term,
)