mirror of
https://github.com/vale981/HOPSFlow-Paper
synced 2025-03-04 09:11:40 -05:00
75 lines
2.3 KiB
Python
75 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)
|