# -*- coding: utf-8 -*-
# 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.
# pylint: disable=invalid-name
r"""
Compute perturbation theory terms for an LMDE.
"""
from typing import List, Optional, Callable
from scipy.integrate._ivp.ivp import OdeResult
from multiset import Multiset
from qiskit import QiskitError
from qiskit_dynamics import ArrayLike
from qiskit_dynamics import DYNAMICS_NUMPY as unp
from qiskit_dynamics.solvers.solver_functions import _is_diffrax_method, _is_jax_method
from qiskit_dynamics.perturbation.multiset_utils import _clean_multisets
from qiskit_dynamics.perturbation.perturbation_utils import (
_merge_multiset_expansion_order_labels,
_merge_list_expansion_order_labels,
)
from qiskit_dynamics.perturbation.dyson_magnus import (
_solve_lmde_dyson,
_solve_lmde_magnus,
_solve_lmde_dyson_jax,
_solve_lmde_magnus_jax,
)
[docs]
def solve_lmde_perturbation(
perturbations: List[Callable],
t_span: ArrayLike,
expansion_method: str,
expansion_order: Optional[int] = None,
expansion_labels: Optional[List[Multiset]] = None,
perturbation_labels: Optional[List[Multiset]] = None,
generator: Optional[Callable] = None,
y0: Optional[ArrayLike] = None,
dyson_in_frame: Optional[bool] = True,
integration_method: Optional[str] = "DOP853",
t_eval: Optional[ArrayLike] = None,
**kwargs,
) -> OdeResult:
r"""Compute time-dependent perturbation theory terms for an LMDE.
This function computes multi-variable Dyson or Magnus expansion terms via the algorithm in
:footcite:`puzzuoli_algorithms_2023`, or Dyson-like terms via the algorithm in
:footcite:`haas_engineering_2019`. See the
:ref:`review on time-dependent perturbation theory <perturbation review>`
to understand the details and notation used in this documentation.
Which expansion is used is specified by the ``expansion_method`` argument, which impacts the
interpretation of several of the function arguments (described below). Regardless of
``expansion_method``, the main computation is performed by solving a differential equation,
utilizing :func:`.solve_ode`, and as such several of the function arguments are direct inputs
into this function:
- ``integration_method`` is the ODE method used (passed as ``method`` to
:func:`.solve_ode`), ``t_span`` is the integration interval, and ``t_eval`` is an optional
set of points to evaluate the perturbation terms at.
- ``kwargs`` are passed directly to :func:`.solve_ode`, enabling passing through of
tolerance or step size arguments.
Other arguments which are treated the same regardless off ``expansion_method`` are:
- ``generator`` is the unperturbed generator, and the computation is performed in the
toggling frame of this generator.
If ``expansion_method in ['dyson', 'magnus']``, this function computes either multivariable
Dyson series or Magnus expansion terms. That is, given a (finitely truncated) power series for
the generator:
.. math::
G(t, c_0, \dots, c_{r-1}) = G_\emptyset(t)
+ \sum_{k=1}^\infty \sum_{I \in \mathcal{I}_k(r)} c_I G_I(t),
this function computes, in the toggling frame of :math:`G_\emptyset(t)` given by ``generator``,
either a collection of multivariable Dyson terms :math:`\mathcal{D}_I(t)` or multivariable
Magnus terms :math:`\mathcal{O}_I(t)`, whose definitions are given in the
:ref:`perturbation theory review <perturbation review>`. In this case, the arguments to the
function are interpreted as follows:
- ``perturbations`` and ``perturbation_labels`` specify the truncated generator power
series. ``perturbations`` provides a list of python callable functions for the non-zero
:math:`G_I(t)`, and the ``perturbation_labels`` is a list of the corresponding multiset
labels :math:`I`, in the form of ``Multiset`` instances. If not specified, the labels are
assumed to be ``[Multiset({0: 1}), ..., Multiset({len(perturbations) - 1: 1})]``.
- ``expansion_order`` and ``expansion_labels`` specify which terms in the chosen expansion
are to be computed. ``expansion_order`` specifies that all expansion terms up to a given
order are to be computed, and ``expansion_labels`` specifies individual terms to be
computed, specified as ``Multiset`` instances. At least one of ``expansion_order`` and
``expansion_labels`` must be specified. If both are specified, then all terms up to
``expansion_order`` will be computed, along with any additional specific terms given by
``expansion_labels``.
Note that in the above, this function requires that the Multisets consist of non-negative
integers. Arguments requiring lists of ``Multiset`` instances also accept lists of any valid
format acceptable to the ``Multiset`` constructor (modulo the non-negative integer constraint).
If ``expansion_method == 'dyson_like'``, the setup is different. In this case, for a list of
matrix-valued functions :math:`G_0(t), \dots, G_{r-1}(t)`, this function computes integrals of
the form
.. math::
\int_{t_0}^{t_F} dt_1 \int_{t_0}^{t_1} dt_2 \dots \int_{t_0}^{t_{k-1}}dt_k
\tilde{G}_{i_1}(t_1) \dots \tilde{G}_{i_k}(t_k),
for lists of integers :math:`[i_1, \dots, i_k]`, and similar to the other cases,
:math:`\tilde{G}_j(t) = V(t)^\dagger G_j(t)V(t)`, i.e. the computation is performed in the
toggling frame specified by ``generator``.
- ``perturbations`` gives the list of matrix functions as callables
:math:`G_0(t), \dots, G_{r-1}(t)`.
- ``perturbation_labels`` is not used in this mode.
- ``expansion_order`` specifies that all possible integrals of the above form should be
computed up to a given order (i.e. integrals up to a given order with all possible
orderings of the :math:`G_0(t), \dots, G_{r-1}(t)`).
- ``expansion_labels`` allows for specification of specific terms to be computed. In this
case, a term is specified by a list of ``int``\s, where the length of the list is the
order of the integral, and the :math:`G_0(t), \dots, G_{r-1}(t)` appear in the integral in
the order given by the list.
Finally, additional optional arguments which can be used in the ``'dyson'`` and ``'dyson_like'``
cases are:
- ``dyson_in_frame`` controls which frame the results are returned in. The default is
``True``, in which case the results are returned as described above. If ``False``, the
returned results include a pre-factor of :math:`V(t)`, the solution of the unperturbed
generator, e.g. :math:`V(t)\mathcal{D}_I(t)`. If ``expansion_method=='magnus'``, this
argument has no effect on the computation.
- ``y0`` is the initial state for the LMDE given by the unperturbed generator. The effect of
this argument on the output is to multiply all outputs by ``y0`` on the right. If ``y0``
is supplied, the argument ``dyson_in_frame`` must be ``False``, as the operator
:math:`V(t)` is never explicitly computed, and therefore it cannot be removed from the
results. If ``y0`` is 1d, it will first be transformed into a 2d column vector to ensure
consistent usage of matrix multiplication.
Regardless of the value of ``expansion_method``, results are returned in an ``OdeResult``
instance in the same manner as :func:`.solve_ode`. The result object stores the solution of the
LMDE for ``generator`` and ``y0`` as if these were passed directly to :func:`.solve_ode` before,
as well as the computed perturbation theory terms in the additional attribute
``perturbation_data``. If ``expansion_method in ['dyson', 'magnus']``, the ``perturbation_data``
attribute stores a :class:`.PowerSeriesData` instance, and if
``expansion_method == 'dyson_like'``, it stores a :class:`.DysonLikeData` instance. In either
case, these are data container classes with the following attributes:
- ``metadata``: Containing expansion information.
- ``labels``: Index labels for all computed perturbation terms. In the case of the Dyson or
Magnus expansion, the labels are ``Multiset`` instances, and in the `'dyson_like'` case,
they are lists of ``int``\s.
- ``data``: A 4d array storing all computed terms. The first axis indexes the expansion
terms in the same ordering as ``labels``, and the second axis indexes the perturbation
terms evaluated at the times in ``results.t`` in the same manner as ``results.y``.
Additionally, both the :meth:`.PowerSeriesData.get_item` and :meth:`.DysonLikeData.get_item`
methods can be used to retrieve the results for a given perturbation term according to its
label. E.g. the results for the term with label ``[0, 1]`` is retrievable via
``results.perturbation_data.get_term([0, 1])``.
Args:
perturbations: List of matrix-valued callables.
t_span: Integration bounds.
expansion_method: Either ``'dyson'``, ``'magnus'``, or ``'dyson_like'``.
expansion_order: Order of perturbation terms to compute up to. Specifying this argument
results in computation of all terms up to the given order. Can be used in conjunction
with ``expansion_labels``.
expansion_labels: Specific perturbation terms to compute. If both ``expansion_order``
and ``expansion_labels`` are specified, then all terms up to ``expansion_order`` are
computed, along with the additional terms specified in ``expansion_labels``.
perturbation_labels: Optional description of power series terms specified by
``perturbations``. To only be used with ``'dyson'`` and ``'magnus'`` methods.
generator: Optional frame generator. Defaults to ``0``.
y0: Optional initial state for frame generator LMDE. Defaults to the identity matrix.
dyson_in_frame: For ``expansion_method`` ``'dyson'`` or ``'dyson_like'``, whether or not
to remove the frame transformation pre-factor from the Dyson terms.
integration_method: Integration method to use. Must be supported by :func:`.solve_ode`.
t_eval: Points at which to evaluate the system.
**kwargs: Additional arguments to pass to ode integration method used to compute terms.
Returns:
OdeResult: Results object containing standard ODE results for LMDE given by ``generator``
and ``y0``, with additional ``perturbation_data`` attribute containing the requested
perturbation theory terms.
Raises:
QiskitError: If problem with inputs, either ``expansion_method`` is unsupported, or both of
``expansion_order`` and ``expansion_labels`` unspecified.
.. footbibliography::
"""
# validation checks
if y0 is not None:
if "magnus" in expansion_method:
raise QiskitError("""Argument y0 cannot be used for expansion_method=='magnus'.""")
if dyson_in_frame:
raise QiskitError(
"""If expansion_method in ['dyson', 'dyson_like']
and y0 passed, dyson_in_frame must be False."""
)
# if 1d in a dyson case, turn into a column vector
if y0.ndim == 1:
y0 = unp.asarray([y0]).transpose()
if perturbation_labels is not None and expansion_method == "dyson_like":
raise QiskitError(
"""perturbation_labels argument not usable with
expansion_method='dyson_like'."""
)
# clean and validate perturbation_labels, and setup expansion terms to compute
if expansion_method in ["dyson", "magnus"]:
if perturbation_labels is None:
perturbation_labels = [Multiset({idx: 1}) for idx in range(len(perturbations))]
else:
# validate perturbation_labels
perturbations_len = len(perturbation_labels)
perturbation_labels = [Multiset(x) for x in perturbation_labels]
cleaned_perturbation_labels = _clean_multisets(perturbation_labels)
if len(cleaned_perturbation_labels) != perturbations_len:
raise QiskitError("perturbation_labels argument contains duplicates as multisets.")
expansion_labels = _merge_multiset_expansion_order_labels(
perturbation_labels=perturbation_labels,
expansion_order=expansion_order,
expansion_labels=expansion_labels,
)
elif expansion_method in ["dyson_like"]:
expansion_labels = _merge_list_expansion_order_labels(
perturbation_num=len(perturbations),
expansion_order=expansion_order,
expansion_labels=expansion_labels,
)
use_jax = _is_jax_method(integration_method) or _is_diffrax_method(integration_method)
if expansion_method in ["dyson", "dyson_like"]:
dyson_like = expansion_method == "dyson_like"
if not use_jax:
return _solve_lmde_dyson(
perturbations=perturbations,
t_span=t_span,
dyson_terms=expansion_labels,
perturbation_labels=perturbation_labels,
generator=generator,
y0=y0,
dyson_in_frame=dyson_in_frame,
dyson_like=dyson_like,
integration_method=integration_method,
t_eval=t_eval,
**kwargs,
)
else:
return _solve_lmde_dyson_jax(
perturbations=perturbations,
t_span=t_span,
dyson_terms=expansion_labels,
perturbation_labels=perturbation_labels,
generator=generator,
y0=y0,
dyson_in_frame=dyson_in_frame,
dyson_like=dyson_like,
integration_method=integration_method,
t_eval=t_eval,
**kwargs,
)
elif expansion_method == "magnus":
if not use_jax:
return _solve_lmde_magnus(
perturbations=perturbations,
t_span=t_span,
magnus_terms=expansion_labels,
perturbation_labels=perturbation_labels,
generator=generator,
y0=y0,
integration_method=integration_method,
t_eval=t_eval,
**kwargs,
)
else:
return _solve_lmde_magnus_jax(
perturbations=perturbations,
t_span=t_span,
magnus_terms=expansion_labels,
perturbation_labels=perturbation_labels,
generator=generator,
y0=y0,
integration_method=integration_method,
t_eval=t_eval,
**kwargs,
)
# raise error if none apply
raise QiskitError(f"expansion_method {expansion_method} not supported.")