plotting fixes

This commit is contained in:
2024-06-20 11:13:15 +02:00
parent cdd666e298
commit a2ec986d0c
2 changed files with 49 additions and 11 deletions

View File

@@ -7,6 +7,10 @@ from typing import Any, Callable, Literal, Optional, Sequence, Union
import matplotlib.colors
import matplotlib.gridspec as gs
import matplotlib.pyplot as plt
from matplotlib.ticker import (
LinearLocator,
mtransforms,
)
import numpy as np
from matplotlib.axes import Axes
from matplotlib.colors import ColorConverter, ListedColormap
@@ -27,6 +31,22 @@ RangeType = tuple[float, float, Union[str, Callable]]
NO_LIM = object()
class InverseLocator(LinearLocator):
def tick_values(self, vmin, vmax):
vmin, vmax = mtransforms.nonsingular(vmin, vmax, expander=0.05)
if (vmin, vmax) in self.presets:
return self.presets[(vmin, vmax)]
if self.numticks == 0:
return []
ticklocs = 1 / np.linspace(1 / vmax, 1 / vmin, self.numticks)
ticklocs = 20 * np.unique(np.round(ticklocs / 20))
return self.raise_if_exceeds(ticklocs)
@dataclass
class ImageData:
x: np.ndarray
@@ -1210,16 +1230,23 @@ def auto_wl_range(wl: np.ndarray, wl_int: np.ndarray, threshold: float) -> PlotR
def summary_plot(
specs: Spectrum,
z: Sequence[float] | None = None,
wl_range: PlotRange | None = None,
spec_range: PlotRange | None = None,
time_range: PlotRange | None = None,
db_min: float = -50.0,
lin_min: float = 1e-3,
axes: tuple[Axes, Axes] | None = None,
wl_db="1D",
spec_db="1D",
time_db=False,
cmap: str | matplotlib.colors.LinearSegmentedColormap = None,
) -> tuple[ImageData, ImageData]:
wl_int = specs.wl_int
if spec_range is not None and spec_range.unit.type != "WL":
spec_int = specs.afreq_int
x_axis = specs.w_disp
freq_ax = True
else:
spec_int = specs.wl_int
x_axis = specs.wl_disp
freq_ax = False
time_int = specs.time_int
if z is None:
@@ -1229,11 +1256,11 @@ def summary_plot(
z = np.asarray(z)
calc_limit, wl_disp_limit = (
(10 ** (0.1 * db_min - 1), db_min) if wl_db and wl_db != "linear 1D" else (lin_min, 0)
(10 ** (0.1 * db_min - 1), db_min) if spec_db and spec_db != "linear 1D" else (lin_min, 0)
)
if wl_range is None:
imin, imax = math.span_above(wl_int, wl_int.max() * calc_limit)
wl_range = PlotRange(specs.wl_disp[imin] * 1e9, specs.wl_disp[imax] * 1e9, "nm")
if spec_range is None:
imin, imax = math.span_above(spec_int, spec_int.max() * calc_limit)
spec_range = PlotRange(specs.wl_disp[imin] * 1e9, specs.wl_disp[imax] * 1e9, "nm")
calc_limit, time_disp_limit = (
(10 ** (0.1 * db_min - 1), db_min) if time_db and time_db != "linear 1D" else (lin_min, 0)
@@ -1247,7 +1274,7 @@ def summary_plot(
else:
left, right = axes
image_spec = transform_2D_data(wl_int, wl_range, specs.wl_disp, z, log=wl_db)
image_spec = transform_2D_data(spec_int, spec_range, x_axis, z, log=spec_db)
left.imshow(
image_spec.data,
extent=image_spec.extent,
@@ -1266,4 +1293,13 @@ def summary_plot(
vmin=time_disp_limit,
cmap=cmap,
)
if freq_ax:
secax = left.secondary_xaxis(
"top",
functions=(
lambda x: units.nm_rads(spec_range.unit(x)),
lambda x: spec_range.unit.inv(units.nm_rads(x)),
),
)
secax.xaxis.set_major_locator(InverseLocator(numticks=5))
return image_spec, image_time

View File

@@ -267,6 +267,7 @@ class Spectrum(np.ndarray):
gate_width: float = 100e-15,
wavelength: bool = True,
autocrop: bool | float = 1e-5,
step_size: int = 4,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
dt = self.t[1] - self.t[0]
sigma = gate_width / (2 * np.sqrt(2 * np.log(2))) / dt
@@ -276,7 +277,7 @@ class Spectrum(np.ndarray):
1 / dt,
window=("gaussian", sigma),
nperseg=nperseg,
noverlap=nperseg - 4,
noverlap=nperseg - step_size,
detrend=False,
scaling="psd",
boundary=None,
@@ -291,8 +292,9 @@ class Spectrum(np.ndarray):
if wavelength:
f = units.m_hz(f)
s = units.to_WL(s.T, f).T
f = f[::-1]
s = s[::-1]
o = f.argsort()
f = f[o]
s = s[o]
if autocrop:
thr = s.max()