mirror of
https://github.com/vale981/HOPSFlow-Paper
synced 2025-03-04 17:21:42 -05:00
76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
|
"""Random and unversioned utility functions."""
|
|||
|
|
import numpy as np
|
|||
|
|
from numpy.typing import NDArray
|
|||
|
|
from statsmodels.nonparametric.smoothers_lowess import lowess
|
|||
|
|
|
|||
|
|
|
|||
|
|
def strobe_times(time: NDArray[np.float64], frequency: float, tolerance: float = 1e-4):
|
|||
|
|
r"""
|
|||
|
|
Given a time array ``time`` and an angular frequency ``frequency`` (ω) the
|
|||
|
|
time points (and their indices) coinciding with :math:`2π / ω \cdot n` within the
|
|||
|
|
``tolerance`` are being returned.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
stroboscope_interval = 2 * np.pi / frequency
|
|||
|
|
sorted_times = np.sort(time)
|
|||
|
|
tolerance = min(tolerance, (sorted_times[1:] - sorted_times[:-1]).min() / 2)
|
|||
|
|
strobe_indices = np.where((time % stroboscope_interval) <= tolerance)[0]
|
|||
|
|
|
|||
|
|
if len(strobe_indices) == 0:
|
|||
|
|
raise ValueError("Can't match the strobe interval to the times.")
|
|||
|
|
|
|||
|
|
strobe_times = time[strobe_indices]
|
|||
|
|
|
|||
|
|
return strobe_times, strobe_indices
|
|||
|
|
|
|||
|
|
|
|||
|
|
def linspace_with_strobe(
|
|||
|
|
begin: float, end: float, N: int, strobe_frequency: float
|
|||
|
|
) -> NDArray[np.float64]:
|
|||
|
|
"""
|
|||
|
|
Like ``linspace`` but so that the time points defined by the
|
|||
|
|
stroboscope angular frequency ``strobe_frequency`` are included.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
return np.unique(
|
|||
|
|
np.sort(
|
|||
|
|
np.concatenate(
|
|||
|
|
[
|
|||
|
|
np.linspace(begin, end, N),
|
|||
|
|
np.arange(begin, end, 2 * np.pi / strobe_frequency),
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def ergotropy(
|
|||
|
|
ρs: NDArray[np.complex128], H: NDArray[np.complex128]
|
|||
|
|
) -> NDArray[np.float64]:
|
|||
|
|
"""
|
|||
|
|
Calculates the ergotropy of the states ``ρs`` with respect to the
|
|||
|
|
hamiltonian ``H``.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
energies = np.linalg.eigh(H)[0]
|
|||
|
|
|
|||
|
|
pops = np.array([np.linalg.eigh(ρ)[0][::-1] for ρ in ρs])
|
|||
|
|
return (
|
|||
|
|
energies * np.diagonal(ρs, axis1=1, axis2=2).real[:, ::-1] - pops * energies
|
|||
|
|
).real.sum(axis=1)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def smoothen(
|
|||
|
|
t: NDArray[np.float64], y: NDArray[np.float64], **kwargs
|
|||
|
|
) -> NDArray[np.float64]:
|
|||
|
|
"""
|
|||
|
|
Uses :any:`statsmodels.nonparametric.smoothers_lowess.lowess` to
|
|||
|
|
smoothen ``y`` indexed by ``t``.
|
|||
|
|
|
|||
|
|
The ``kwargs`` are passed straight into
|
|||
|
|
:any:`statsmodels.nonparametric.smoothers_lowess.lowess`.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
kwargs = dict(frac=0.1, it=0) | kwargs
|
|||
|
|
return lowess(y, t, is_sorted=True, return_sorted=False, **kwargs)
|