mirror of
https://github.com/vale981/fibre_walk_project_code
synced 2025-03-05 09:51:39 -05:00
136 lines
3.5 KiB
Python
136 lines
3.5 KiB
Python
from . import data
|
|
import matplotlib.pyplot as plt
|
|
from scipy.ndimage import uniform_filter1d
|
|
from scipy.interpolate import make_interp_spline
|
|
import numpy as np
|
|
|
|
|
|
def wrap_plot(f):
|
|
def wrapped(*args, ax=None, setup_function=plt.subplots, **kwargs):
|
|
fig = None
|
|
if not ax:
|
|
fig, ax = setup_function()
|
|
|
|
ret_val = f(*args, ax=ax, **kwargs)
|
|
return (fig, ax, ret_val) if ret_val else (fig, ax)
|
|
|
|
return wrapped
|
|
|
|
|
|
def lighten_color(color, amount=0.5):
|
|
"""
|
|
Lightens the given color by multiplying (1-luminosity) by the given amount.
|
|
Input can be matplotlib color string, hex string, or RGB tuple.
|
|
|
|
Examples:
|
|
>> lighten_color('g', 0.3)
|
|
>> lighten_color('#F034A3', 0.6)
|
|
>> lighten_color((.3,.55,.1), 0.5)
|
|
"""
|
|
import matplotlib.colors as mc
|
|
import colorsys
|
|
|
|
try:
|
|
c = mc.cnames[color]
|
|
except:
|
|
c = color
|
|
c = colorsys.rgb_to_hls(*mc.to_rgb(c))
|
|
return colorsys.hls_to_rgb(c[0], 1 - amount * (1 - c[1]), c[2])
|
|
|
|
|
|
def fancy_error(x, y, err, ax, **kwargs):
|
|
line = ax.plot(
|
|
x,
|
|
y,
|
|
**kwargs,
|
|
)
|
|
|
|
err = ax.fill_between(
|
|
x,
|
|
y + err,
|
|
y - err,
|
|
color=lighten_color(line[0].get_color(), 0.5),
|
|
alpha=0.5,
|
|
)
|
|
|
|
return line, err
|
|
|
|
|
|
@wrap_plot
|
|
def plot_scan(
|
|
data: data.ScanData,
|
|
laser=False,
|
|
output=True,
|
|
steps: bool | int = False,
|
|
normalize=False,
|
|
smoothe_output: bool | int = False,
|
|
ax=None,
|
|
**kwargs,
|
|
):
|
|
if not (laser or output):
|
|
raise ValueError("At least one of 'laser' or 'output' must be True.")
|
|
|
|
time, output_data, laser_data = (
|
|
(data.time, data.output, data.laser)
|
|
if isinstance(steps, bool)
|
|
else data.for_step(steps)
|
|
)
|
|
|
|
if laser:
|
|
if normalize:
|
|
laser_data = (laser_data - laser_data.min()) / (
|
|
laser_data.max() - laser_data.min()
|
|
)
|
|
|
|
lines = ax.plot(time, laser_data, **kwargs)
|
|
|
|
if output:
|
|
if normalize:
|
|
output_data = (output_data - output_data.min()) / (
|
|
output_data.max() - output_data.min()
|
|
)
|
|
|
|
if smoothe_output:
|
|
if not isinstance(smoothe_output, int):
|
|
smoothe_output = 60
|
|
|
|
window = len(output_data) // smoothe_output
|
|
output_data = uniform_filter1d(output_data, window)
|
|
|
|
lines = ax.plot(time, output_data, **kwargs)
|
|
|
|
if isinstance(steps, bool) and steps:
|
|
peaks = data.laser_steps()
|
|
peaks = [0, *peaks, len(data.time) - 1]
|
|
|
|
vertical = output_data.min()
|
|
for peak in peaks[1:-1]:
|
|
ax.axvline(
|
|
data.time[peak],
|
|
color=lighten_color(lines[0].get_color()),
|
|
linestyle="--",
|
|
zorder=-10,
|
|
)
|
|
|
|
for i, (begin, end) in enumerate(zip(peaks[:-1], peaks[1:])):
|
|
ax.text(
|
|
(data.time[begin] + data.time[end]) / 2,
|
|
vertical,
|
|
f"{i}",
|
|
ha="center",
|
|
va="bottom",
|
|
)
|
|
|
|
|
|
@wrap_plot
|
|
def plot_transmission(data: data.ScanData, timepoints=1000, ax=None, **kwargs):
|
|
amplitude = data.output_end_averages()
|
|
times = data.laser_step_times()
|
|
|
|
smoothed = make_interp_spline(times, amplitude, k=3)
|
|
smooth_times = np.linspace(times.min(), times.max(), timepoints)
|
|
|
|
lines = ax.plot(smooth_times, smoothed(smooth_times), **kwargs)
|
|
plt.plot(
|
|
times, amplitude, "o", color=lighten_color(lines[0].get_color()), alpha=0.5
|
|
)
|