# Source code for oqupy.helpers

# Copyright 2022 The TEMPO Collaboration
#
# you may not use this file  in compliance with the License.
# You may obtain a copy of the License at
#
#
# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
"""
Handy helper functions.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.axes import Axes

from oqupy.correlations import BaseCorrelations
from oqupy.tempo import TempoParameters

[docs]def plot_correlations_with_parameters(
correlations: BaseCorrelations,
parameters: TempoParameters,
ax: Axes = None) -> Axes:
"""Plot the correlation function on a grid that corresponds to some
tempo parameters. For comparison, it also draws a solid line that is 10%
longer and has two more sampling points per interval.

Parameters
----------
correlations: BaseCorrelations
The correlation object we are interested in.
parameters: TempoParameters
The tempo parameters that determine the grid.
"""
infinity = False
infinity = True
else:
infinity = False

dt = parameters.dt
dkmax = parameters.dkmax

times_infl = dt/3.0 * np.arange((dkmax+1)*3 - 2)

corr = np.vectorize(correlations.correlation)

corr_infl = corr(times_infl)
sample = [3*i for i in range(dkmax+1)]
corr_extra = corr(times_extra)

show = False
if ax is None:
fig, ax = plt.subplots()
show = True
ax.set_xlabel(r"$\tau$")
ax.set_ylabel(r"$C(\tau)$")
ax.plot(
times_infl, np.real(corr_infl),
color="C0", linestyle="-", label="real")
ax.scatter(
times_infl[sample], np.real(corr_infl[sample]),
marker="d", color="C0")
ax.plot(
times_infl, np.imag(corr_infl),
color="C1", linestyle="-", label="imag")
ax.scatter(
times_infl[sample], np.imag(corr_infl[sample]),
marker="o", color="C1")
ax.plot(times_extra, np.real(corr_extra), color="C0", linestyle="-")
ax.plot(times_extra, np.imag(corr_extra), color="C1", linestyle="-")

if infinity:
ax.fill_between(
times_extra, np.real(corr_extra),
0.0, color="C0", alpha=0.30)
ax.fill_between(
times_extra, np.imag(corr_extra),
0.0, color="C1", alpha=0.30)
ax.fill_between(
0.0, color="C0", alpha=0.30)
ax.fill_between(