Reformatting docs for perturbation module (#336)

qiskit_dynamics/perturbation/__init__.py

@@ -20,36 +20,33 @@
 .. currentmodule:: qiskit_dynamics.perturbation
 
 This module contains tools for numerically computing and utilizing perturbation theory terms.
-Perturbation theory is an advanced topic; a brief review of the concepts and notation required
-to understand the contents in this module are given in the
-:ref:`Time-dependent perturbation theory and multi-variable
-series expansions review <perturbation review>` discussion.
+Perturbation theory is an advanced topic; a brief review of the concepts and notation required to
+understand the contents in this module are given in the :ref:`Time-dependent perturbation theory and
+multi-variable series expansions review <perturbation review>` discussion.
 
 .. _td perturbation theory:
 
 Time-dependent perturbation theory
 ==================================
 
-The function :func:`.solve_lmde_perturbation` computes
-Dyson series :footcite:`dyson_radiation_1949` and
-Magnus expansion :footcite:`magnus_exponential_1954,blanes_magnus_2009` terms
-in a multi-variable setting via algorithms in :footcite:`puzzuoli_sensitivity_2022`.
-It can also be used to compute Dyson-like integrals using the algorithm in
-:footcite:`haas_engineering_2019`. Results are returned in either a :class:`PowerSeriesData`
-or :class:`DysonLikeData` class, which are data classes with functionality for indexing
-and accessing specific perturbation terms. See the function documentation for further details.
+The function :func:`.solve_lmde_perturbation` computes Dyson series :footcite:`dyson_radiation_1949`
+and Magnus expansion :footcite:`magnus_exponential_1954,blanes_magnus_2009` terms in a
+multi-variable setting via algorithms in :footcite:`puzzuoli_sensitivity_2022`. It can also be used
+to compute Dyson-like integrals using the algorithm in :footcite:`haas_engineering_2019`. Results
+are returned in either a :class:`PowerSeriesData` or :class:`DysonLikeData` class, which are data
+classes with functionality for indexing and accessing specific perturbation terms. See the function
+documentation for further details.
 
 Truncated power-series representation and multisets
 ===================================================
 
-The class :class:`.ArrayPolynomial` represents an array-valued multivariable polynomial
-(i.e. a truncated power series), and provides functionality for
-both evaluating and transforming array-valued polynomials.
+The class :class:`.ArrayPolynomial` represents an array-valued multivariable polynomial (i.e. a
+truncated power series), and provides functionality for both evaluating and transforming
+array-valued polynomials.
 
-This module makes use of the `multiset package <https://pypi.org/project/multiset/>`_ for
-indexing multi-variable power series. See the
-:ref:`multiset and power series notation section <multiset power series>`
-of the perturbation review for an explanation of this convention.
+This module makes use of the `multiset package <https://pypi.org/project/multiset/>`_ for indexing
+multi-variable power series. See the :ref:`multiset and power series notation section <multiset
+power series>` of the perturbation review for an explanation of this convention.
 
 
 Perturbation module functions

qiskit_dynamics/perturbation/array_polynomial.py

@@ -59,23 +59,21 @@ class ArrayPolynomial:
 
     where in the above:
 
-        - :math:`S` is a finite set of multisets
-          indicating non-zero monomial terms,
+        - :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.
+    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``.
+        - ``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
 
@@ -84,8 +82,8 @@ class ArrayPolynomial:
         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
+    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
 
@@ -103,10 +101,10 @@ class ArrayPolynomial:
         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.
+    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
 
@@ -120,15 +118,15 @@ class ArrayPolynomial:
         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
+    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.
+    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
@@ -357,11 +355,10 @@ def add(
     ) -> "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``.
+        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.
@@ -389,11 +386,10 @@ def matmul(
     ) -> "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``.
+        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.
@@ -420,11 +416,10 @@ def mul(
     ) -> "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``.
+        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.
@@ -553,14 +548,14 @@ 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 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.
+    The algorithm computes monomial terms of increasing order, recursively utilizing lower order
+    terms.
 
     Args:
         multisets: list of multisets.

qiskit_dynamics/perturbation/custom_binary_op.py

@@ -29,39 +29,35 @@
 
 
 class _CustomBinaryOp:
-    r"""A binary operation between arrays of dimension >1d built from taking linear combinations
-    of a base binary operation acting on sub-arrays.
+    r"""A binary operation between arrays of dimension >1d built from taking linear combinations of
+    a base binary operation acting on sub-arrays.
 
-    This class constructs customized binary operations between
-    lists of arrays :math:`A = (A_i)` and :math:`B = (B_i)` of the form:
+    This class constructs customized binary operations between lists of arrays :math:`A = (A_i)` and
+    :math:`B = (B_i)` of the form:
 
     .. math::
         (A \times B)_i = \sum_{jk} a_{ijk} f(A_j, B_k),
 
-    where :math:`a_{ijk}` is an array of complex scalars, and :math:`f` is a binary operation
-    (a common example being matrix multiplication).
+    where :math:`a_{ijk}` is an array of complex scalars, and :math:`f` is a binary operation (a
+    common example being matrix multiplication).
 
-    At instantiation the binary operation :math:`f`, as well
-    as the array :math:`a_{ijk}`, are specified. The array :math:`a_{ijk}` is given in a
-    sparse format, as a list where each
-    entry specifies the 2d sub-array :math:`a_{i}` as a 2-tuple
-    with entries:
+    At instantiation the binary operation :math:`f`, as well as the array :math:`a_{ijk}`, are
+    specified. The array :math:`a_{ijk}` is given in a sparse format, as a list where each entry
+    specifies the 2d sub-array :math:`a_{i}` as a 2-tuple with entries:
 
         - The non-zero entries of :math:`a_i` given as a list.
         - A 2d-array with each entry being the pair of indices ``[j,k]``.
 
-    Internally, this specification is translated into one more suited for efficient
-    evaluation:
-        - A specification of the unique pairs of required evaluations
-          :math:`f(A_j, B_k)` in terms of two arrays giving the left
-          indices and the right indices.
-        - A 2-tuple of 2d arrays specifying the linear combination of
-          unique evaluations of :math:`f` to compute the :math:`i^{th}` entry of the
-          binary operations. The :math:`i^{th}` entry of each array is:
+    Internally, this specification is translated into one more suited for efficient evaluation:
+        - A specification of the unique pairs of required evaluations :math:`f(A_j, B_k)` in terms
+          of two arrays giving the left indices and the right indices.
+        - A 2-tuple of 2d arrays specifying the linear combination of unique evaluations of
+          :math:`f` to compute the :math:`i^{th}` entry of the binary operations. The :math:`i^{th}`
+          entry of each array is:
             - The list of coefficients in the linear combination.
             - The index of the unique product for each coefficient.
-          These arrays are padded with the value ``-1`` to make each
-          linear combo the same length (relevant for JAX evaluation).
+          These arrays are padded with the value ``-1`` to make each linear combo the same length
+          (relevant for JAX evaluation).
     """
 
     def __init__(
@@ -149,27 +145,25 @@ def _compile_custom_operation_rule(
     unique_evaluation_len: Optional[int] = None,
     linear_combo_len: Optional[int] = None,
 ) -> Tuple[np.array, np.array]:
-    """Compile the list of unique evaluations and linear combinations required
-    to implement a given operation_rule.
+    """Compile the list of unique evaluations and linear combinations required to implement a given
+    operation_rule.
 
     See _CustomBinaryOp doc string for formatting details.
 
     Args:
-        operation_rule: Custom operation rule in the sparse format in the
-                        _CustomBinaryOp doc string.
-        index_offset: Integer specifying a shift to apply in the 2nd and 3rd indices in the
-                      sparse representation of :math:`a_{ijk}`.
-        unique_evaluation_len: Integer specifying a minimum length to represent the
-                               unique multiplications list. The unique multiplication list
-                               is padded with entries ``[-1, -1]`` to meet the minimum length.
-        linear_combo_len: Minimum length for linear combo specification. Coefficients are
-                          padded with zeros, and the unique multiplication indices are
-                          padded with ``-1``.
+        operation_rule: Custom operation rule in the sparse format in the _CustomBinaryOp doc
+            string.
+        index_offset: Integer specifying a shift to apply in the 2nd and 3rd indices in the sparse
+            representation of :math:`a_{ijk}`.
+        unique_evaluation_len: Integer specifying a minimum length to represent the unique
+            multiplications list. The unique multiplication list is padded with entries ``[-1, -1]``
+            to meet the minimum length.
+        linear_combo_len: Minimum length for linear combo specification. Coefficients are padded
+            with zeros, and the unique multiplication indices are padded with ``-1``.
 
     Returns:
-        Tuple[np.array, np.array]: Multiplication rule compiled into a list of
-                                   unique products and a list of linear combinations of
-                                   unique products for implementing the custom dot rule.
+        Tuple[np.array, np.array]: Multiplication rule compiled into a list of unique products and a
+            list of linear combinations of unique products for implementing the custom dot rule.
     """
 
     # force numpy usage for algebra specification
@@ -251,8 +245,8 @@ def _compute_unique_evaluations(
 def _compute_linear_combos(
     unique_evaluations: np.array, linear_combo_rule: Tuple[np.array, np.array]
 ) -> np.array:
-    r"""Compute linear combinations of the entries in the array ``unique_mults``
-    according to ``linear_combo_rule``. The :math:`j^{th}` entry of the output is given by
+    r"""Compute linear combinations of the entries in the array ``unique_mults`` according to
+    ``linear_combo_rule``. The :math:`j^{th}` entry of the output is given by
     :math:`sum_k c[j, k] unique_mults[idx[j, k]]`, where ``linear_combo_rule`` is ``(c, idx)``.
     """
     M0 = np.zeros_like(unique_evaluations[0])