""" Class for solve function results"""
# Required for Sphinx to follow autodoc_type_aliases
from __future__ import annotations
from typing import TypedDict, Any, Callable
from ..core.numpy_backend import np
from numpy.typing import ArrayLike
from ..core import Qobj, QobjEvo, expect
__all__ = ["Result"]
class _QobjExpectEop:
"""
Pickable e_ops callable that calculates the expectation value for a given
operator.
Parameters
----------
op : :obj:`.Qobj`
The expectation value operator.
"""
def __init__(self, op):
self.op = op
def __call__(self, t, state):
return expect(self.op, state)
class ExpectOp:
"""
A result e_op (expectation operation).
Parameters
----------
op : object
The original object used to define the e_op operation, e.g. a
:~obj:`Qobj` or a function ``f(t, state)``.
f : function
A callable ``f(t, state)`` that will return the value of the e_op
for the specified state and time.
append : function
A callable ``append(value)``, e.g. ``expect[k].append``, that will
store the result of the e_ops function ``f(t, state)``.
Attributes
----------
op : object
The original object used to define the e_op operation.
"""
def __init__(self, op, f, append):
self.op = op
self._f = f
self._append = append
def __call__(self, t, state):
"""
Return the expectation value for the given time, ``t`` and
state, ``state``.
"""
return self._f(t, state)
def _store(self, t, state):
"""
Store the result of the e_op function. Should only be called by
:class:`~Result`.
"""
self._append(self._f(t, state))
class _BaseResult:
"""
Common method for all ``Result``.
"""
def __init__(self, options, *, solver=None, stats=None):
self.solver = solver
if stats is None:
stats = {}
self.stats = stats
self._state_processors = []
self._state_processors_require_copy = False
# make sure not to store a reference to the solver
options_copy = options.copy()
if hasattr(options_copy, "_feedback"):
options_copy._feedback = None
self.options = options_copy
# Almost all integrators already return a copy that is safe to use.
self._integrator_return_copy = options.get("method", None) in [
"adams", "lsoda", "bdf", "dop853", "diag",
"euler", "platen", "explicit1.5",
"milstein", "pred_corr", "taylor1.5",
"milstein_imp", "taylor1.5_imp", "rouchon",
]
def _e_ops_to_dict(self, e_ops):
"""Convert the supplied e_ops to a dictionary of Eop instances."""
if e_ops is None:
e_ops = {}
elif isinstance(e_ops, (list, tuple)):
e_ops = {k: e_op for k, e_op in enumerate(e_ops)}
elif isinstance(e_ops, dict):
pass
else:
e_ops = {0: e_ops}
return e_ops
def add_processor(self, f, requires_copy=False):
"""
Append a processor ``f`` to the list of state processors.
Parameters
----------
f : function, ``f(t, state)``
A function to be called each time a state is added to this
result object. The state is the state passed to ``.add``, after
applying the pre-processors, if any.
requires_copy : bool, default False
Whether this processor requires a copy of the state rather than
a reference. A processor must never modify the supplied state, but
if a processor stores the state it should set ``require_copy`` to
true.
"""
self._state_processors.append(f)
self._state_processors_require_copy |= requires_copy
class ResultOptions(TypedDict):
store_states: bool | None
store_final_state: bool
[docs]
class Result(_BaseResult):
"""
Base class for storing solver results.
Parameters
----------
e_ops : :obj:`.Qobj`, :obj:`.QobjEvo`, function or list or dict of these
The ``e_ops`` parameter defines the set of values to record at
each time step ``t``. If an element is a :obj:`.Qobj` or
:obj:`.QobjEvo` the value recorded is the expectation value of that
operator given the state at ``t``. If the element is a function, ``f``,
the value recorded is ``f(t, state)``.
The values are recorded in the ``e_data`` and ``expect`` attributes of
this result object. ``e_data`` is a dictionary and ``expect`` is a
list, where each item contains the values of the corresponding
``e_op``.
options : dict
The options for this result class.
solver : str or None
The name of the solver generating these results.
stats : dict or None
The stats generated by the solver while producing these results. Note
that the solver may update the stats directly while producing results.
kw : dict
Additional parameters specific to a result sub-class.
Attributes
----------
times : list
A list of the times at which the expectation values and states were
recorded.
states : list of :obj:`.Qobj`
The state at each time ``t`` (if the recording of the state was
requested).
final_state : :obj:`.Qobj`:
The final state (if the recording of the final state was requested).
expect : list of arrays of expectation values
A list containing the values of each ``e_op``. The list is in
the same order in which the ``e_ops`` were supplied and empty if
no ``e_ops`` were given.
Each element is itself a list and contains the values of the
corresponding ``e_op``, with one value for each time in ``.times``.
The same lists of values may be accessed via the ``.e_data`` dictionary
and the original ``e_ops`` are available via the ``.e_ops`` attribute.
e_data : dict
A dictionary containing the values of each ``e_op``. If the ``e_ops``
were supplied as a dictionary, the keys are the same as in
that dictionary. Otherwise the keys are the index of the ``e_op``
in the ``.expect`` list.
The lists of expectation values returned are the *same* lists as
those returned by ``.expect``.
e_ops : dict
A dictionary containing the supplied e_ops as ``ExpectOp`` instances.
The keys of the dictionary are the same as for ``.e_data``.
Each value is object where ``.e_ops[k](t, state)`` calculates the
value of ``e_op`` ``k`` at time ``t`` and the given ``state``, and
``.e_ops[k].op`` is the original object supplied to create the
``e_op``.
solver : str or None
The name of the solver generating these results.
stats : dict or None
The stats generated by the solver while producing these results.
options : dict
The options for this result class.
"""
times: list[float]
states: list[Qobj]
options: ResultOptions
e_data: dict[Any, list[Any]]
def __init__(
self,
e_ops: dict[Any, Qobj | QobjEvo | Callable[[float, Qobj], Any]],
options: ResultOptions,
*,
solver: str = None,
stats: dict[str, Any] = None,
**kw,
):
super().__init__(options, solver=solver, stats=stats)
raw_ops = self._e_ops_to_dict(e_ops)
self.e_data = {k: [] for k in raw_ops}
self.e_ops = {}
for k, op in raw_ops.items():
f = self._e_op_func(op)
self.e_ops[k] = ExpectOp(op, f, self.e_data[k].append)
self.add_processor(self.e_ops[k]._store)
self.times = []
self.states = []
self._final_state = None
self._post_init(**kw)
def _e_op_func(self, e_op):
"""
Convert an e_op entry into a function, ``f(t, state)`` that returns
the appropriate value (usually an expectation value).
Sub-classes may override this function to calculate expectation values
in different ways.
"""
if isinstance(e_op, Qobj):
return _QobjExpectEop(e_op)
elif isinstance(e_op, QobjEvo):
return e_op.expect
elif callable(e_op):
return e_op
raise TypeError(f"{e_op!r} has unsupported type {type(e_op)!r}.")
def _post_init(self):
"""
Perform post __init__ initialisation. In particular, add state
processors or pre-processors.
Sub-class may override this. If the sub-class wishes to register the
default processors for storing states, it should call this parent
``.post_init()`` method.
Sub-class ``.post_init()`` implementation may take additional keyword
arguments if required.
"""
store_states = self.options["store_states"]
store_states = store_states or (
len(self.e_ops) == 0 and store_states is None
)
if store_states:
self.add_processor(self._store_state, requires_copy=True)
store_final_state = self.options["store_final_state"]
if store_final_state and not store_states:
self.add_processor(self._store_final_state, requires_copy=True)
def _store_state(self, t, state):
"""Processor that stores a state in ``.states``."""
self.states.append(state)
def _store_final_state(self, t, state):
"""Processor that writes the state to ``._final_state``."""
self._final_state = state
def _pre_copy(self, state):
"""Return a copy of the state. Sub-classes may override this to
copy a state in different manner or to skip making a copy
altogether if a copy is not necessary.
"""
return state.copy()
def add(self, t, state):
"""
Add a state to the results for the time ``t`` of the evolution.
Adding a state calculates the expectation value of the state for
each of the supplied ``e_ops`` and stores the result in ``.expect``.
The state is recorded in ``.states`` and ``.final_state`` if specified
by the supplied result options.
Parameters
----------
t : float
The time of the added state.
state : typically a :obj:`.Qobj`
The state a time ``t``. Usually this is a :obj:`.Qobj` with
suitable dimensions, but it sub-classes of result might support
other forms of the state.
Notes
-----
The expectation values, i.e. ``e_ops``, and states are recorded by
the state processors (see ``.add_processor``).
Additional processors may be added by sub-classes.
"""
self.times.append(t)
if (
self._state_processors_require_copy
and not self._integrator_return_copy
):
state = self._pre_copy(state)
for op in self._state_processors:
op(t, state)
def __repr__(self):
lines = [
f"<{self.__class__.__name__}",
f" Solver: {self.solver}",
]
if self.stats:
lines.append(" Solver stats:")
lines.extend(f" {k}: {v!r}" for k, v in self.stats.items())
if self.times:
lines.append(
f" Time interval: [{self.times[0]}, {self.times[-1]}]"
f" ({len(self.times)} steps)"
)
lines.append(f" Number of e_ops: {len(self.e_ops)}")
if self.states:
lines.append(" States saved.")
elif self.final_state is not None:
lines.append(" Final state saved.")
else:
lines.append(" State not saved.")
lines.append(">")
return "\n".join(lines)
[docs]
def plot_expect(
self,
*,
fig=None,
axes=None,
labels=None,
title=None,
xlabel="Time",
ylabel="Expectation value",
show_legend=True,
separate_axes=False,
**plot_kwargs,
):
"""
Plot the expectation values from the solver result.
Parameters
----------
fig : matplotlib.figure.Figure, optional
User-provided figure. If ``None`` and *axes* is also ``None``,
a new figure is created.
axes : matplotlib.axes.Axes, optional
User-provided axes. If ``None`` and *fig* is also ``None``,
new axes are created.
labels : list of str, optional
Labels for each expectation-value curve. When ``None``, labels
are taken from the *e_ops* keys (the original dictionary keys
when *e_ops* was passed as a ``dict``, otherwise
``"e_ops[0]"``, ``"e_ops[1]"``, …).
title : str, optional
Title for the plot. When ``None``, the solver name stored in
the result is used.
xlabel : str, optional
Label for the *x*-axis. Default ``"Time"``.
ylabel : str, optional
Label for the *y*-axis. Default ``"Expectation value"``.
show_legend : bool, optional
Whether to display the legend. Default ``True``.
separate_axes : bool, optional
If ``True``, each expectation value is plotted in its own
subplot. Default ``False``.
**plot_kwargs
Additional keyword arguments forwarded to
``matplotlib.axes.Axes.plot``.
Returns
-------
fig : matplotlib.figure.Figure
The figure containing the plot.
axes : matplotlib.axes.Axes or array of Axes
The axes containing the plot.
Raises
------
ValueError
If no expectation-value data is available.
"""
try:
import matplotlib.pyplot as plt
except ImportError as err:
raise ImportError(
"matplotlib is required for plotting. "
"Install it with: pip install matplotlib"
) from err
if not self.e_data:
raise ValueError(
"No expectation-value data to plot. "
"Ensure that e_ops were supplied to the solver."
)
if labels is None:
labels = [
key if isinstance(key, str) else f"e_ops[{key}]"
for key in self.e_data.keys()
]
n_e_ops = len(self.expect)
if title is None:
title = self.solver
if separate_axes:
custom_axes = True
if fig is None and axes is None:
fig, axes = plt.subplots(
n_e_ops, 1, sharex=True,
figsize=(6, 3 * n_e_ops),
squeeze=False,
)
axes = axes.flatten()
custom_axes = False
elif axes is None:
axes = np.array([
fig.add_subplot(n_e_ops, 1, i + 1)
for i in range(n_e_ops)
])
custom_axes = False
elif fig is None:
if not hasattr(axes, '__len__'):
axes = np.array([axes])
fig = axes[0].get_figure()
for ax, expectation, label in zip(axes, self.expect, labels):
ax.plot(
self.times, expectation, label=label, **plot_kwargs
)
ax.set_ylabel(ylabel)
if show_legend:
ax.legend()
if custom_axes:
ax.set_xlabel(xlabel)
ax.set_title(title)
if not custom_axes:
axes[0].set_title(title)
axes[-1].set_xlabel(xlabel)
else:
if fig is None and axes is None:
fig, axes = plt.subplots()
elif axes is None:
axes = fig.add_subplot(111)
elif fig is None:
fig = axes.get_figure()
for label, expectation in zip(labels, self.expect):
axes.plot(
self.times, expectation, label=label, **plot_kwargs
)
axes.set_xlabel(xlabel)
axes.set_ylabel(ylabel)
axes.set_title(title)
if show_legend:
axes.legend()
return fig, axes
@property
def expect(self) -> list[ArrayLike]:
return [np.array(e_op) for e_op in self.e_data.values()]
@property
def final_state(self) -> Qobj:
if self._final_state is not None:
return self._final_state
if self.states:
return self.states[-1]
return None