Figures: Simple Correlation figures

This commit is contained in:
Eric Teunis de Boone 2023-06-23 19:16:37 +02:00
parent 4fa10f59c1
commit 821a2a340b
4 changed files with 266 additions and 0 deletions

11
figures/methods/Makefile Normal file
View file

@ -0,0 +1,11 @@
SUBDIRS := $(subst Makefile,,$(wildcard */Makefile))
.PHONY: all dist dist-clean $(SUBDIRS)
all: dist $(SUBDIRS)
dist: \
#
$(SUBDIRS):
@$(MAKE) -C $@

View file

@ -0,0 +1,12 @@
.PHONY: all
all: pdf png
pdf png : src/correlation_figure.py
$< waveforms.$@ correlation.$@
waveforms.pdf correlation.pdf : src/correlation_figure.py
pdf
waveforms.png correlation.png : src/correlation_figure.py
png

View file

@ -0,0 +1,231 @@
#!/usr/bin/env python3
# vim: fdm=marker fmr=<<<,>>>
__doc__ = \
"""
Create one/two figures exemplifiying the correlation between waveforms.
"""
import numpy as np
import matplotlib.pyplot as plt
## Correlating function
def my_correlation(in1, template, lags=None, normalise=True):
""" From 11_pulsed_timing """
template_length = len(template)
in1_length = len(in1)
if lags is None:
lags = np.arange(-template_length+1, in1_length + 1)
# do the correlation jig
corrs = np.zeros_like(lags, dtype=float)
for i, l in enumerate(lags):
if l <= 0: # shorten template at the front
in1_start = 0
template_end = template_length
template_start = -template_length - l
in1_end = max(0, min(in1_length, -template_start)) # 0 =< l + template_length =< in1_lengt
elif l > in1_length - template_length:
# shorten template from the back
in1_end = in1_length
template_start = 0
in1_start = min(l, in1_length)
template_end = max(0, in1_length - l)
else:
in1_start = min(l, in1_length)
in1_end = min(in1_start + template_length, in1_length)
# full template
template_start = 0
template_end = template_length
# Slice in1 and template
in1_slice = in1[in1_start:in1_end]
template_slice = template[template_start:template_end]
corrs[i] = np.dot(in1_slice, template_slice)
if normalise:
corrs /= np.amax(corrs)
return corrs, (in1, template, lags)
## Some Functions
def boxcar(x, x_low, x_high, x2=1):
return np.heaviside(x-x_low, x2) * (1-np.heaviside(x-x_high, x2))
def f1(x, x0=0, x1=1, x2=2, x3=3, h=1):
""" triangle with boxcar inbetween """
if x0 == x1:
slant_up = 0
else:
norm = (x1-x0)
slant_up = (x-x0)/norm*boxcar(x, x0, x1)
if x2 == x3:
slant_down = 0
else:
norm = (x3-x2)
slant_down = (x3-x)/norm * boxcar(x, x2, x3)
y = h*(slant_up + slant_down + boxcar(x, x1, x2))
return y
def triangle(x, x0=0, w=1, h=1):
x2 = x1 = x0
x0 = x1 - w
x3 = x1 + 2
return f1(x, x0, x1, x2, x3, h)
def slanted_boxcar(x, x0=0, slant_width = 1, box_width=1, h=1):
x1 = x0 - box_width/2
x2 = x0 + box_width/2
x0 = x1 - slant_width/2
x3 = x2 + slant_width/2
return f1(x, x0, x1, x2, x3, h)
def sine(x, x0=0, x1=1, n_periods=0.5, phase=0):
L = x1-x0
return np.sin(2*np.pi* x * n_periods/L + phase)*boxcar(x, x0, x1)
def f2(x, x0=0, x1=2, h=1):
""" logarithm """
return h*np.log(1+x-x0)*boxcar(x, x0, x1)
## Figures
def plot_waveform(y, x=None, ax=None, fig_kwargs={}, **line_kwargs):
if ax is None:
fig, ax = plt.subplots(**fig_kwargs)
if x is not None:
ax.set_xlabel("Time")
else:
ax.set_xlabel("Sample no.")
ax.set_ylabel("Amplitude")
if x is not None:
y, x = x, y
else:
x = np.arange(0, len(y))
ax.plot(x, y, **line_kwargs)
return ax
def correlation_figure(y1, y2, lags=None, lag_dt=1, ax=None, fig_kwargs={}, **line_kwargs):
""" figure """
if ax is None:
fig, ax = plt.subplots(**fig_kwargs)
ax.set_xlabel("Delay [samples]")
ax.set_ylabel("Normalised Correlation")
corrs, (_, __, lags) = my_correlation(y1, y2, lags=lags, normalise=True)
ax.plot(lag_dt * lags, corrs, **line_kwargs)
return ax, (corrs, lags)
def main(N_x = 1e4, x_delay=5, separate_figures=True, fig_kwargs={}):
x = np.linspace(-2, 8, N_x)
y1 = slanted_boxcar(x, x0=0, slant_width=0, box_width=3) + 1/4*sine(x, x0=-3/2, x1=3/2, phase=np.pi/2)
y2 = slanted_boxcar(x, x0=0+x_delay, slant_width=1, box_width=0)
if not separate_figures:
fig, axs = plt.subplots(2,1, **fig_kwargs)
else:
axs = (None, None)
figs = []
if True:
line_kwargs = dict(alpha=0.7)
ax = plot_waveform(x, y1, ax=axs[0], label="Waveform 1", fig_kwargs=fig_kwargs)
plot_waveform(x, y2, ax=ax, label="Waveform 2")
ax.set_xlabel("Time")
ax.set_ylabel("Amplitude")
ax.grid()
ax.get_figure().tight_layout()
figs.append(ax.get_figure())
if True:
lag_dt = x[1] - x[0]
ax2, (corrs, lags) = correlation_figure(y2, y1, ax=axs[1], lag_dt=lag_dt, fig_kwargs=fig_kwargs)
ax2.grid()
ax2.set_xlabel("Delay")
ax2.set_ylabel("Normalised Correlation")
best_lag_idx = np.argmax(corrs)
best_lag_x = lags[best_lag_idx] * lag_dt
if True:
wx = 4
ax2.set_xlim(x_delay - wx, x_delay + wx)
ax2.axvline(best_lag_x, ls='dashed', alpha=0.7, color='r')
ax2.get_figure().tight_layout()
figs.append(ax2.get_figure())
if not separate_figures:
return fig
else:
return figs
if __name__ == "__main__":
from argparse import ArgumentParser
import os.path as path
import os
import sys
# Append parent directory to import path so pyfiglib can be found
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
import pyfiglib as pfl
parser = ArgumentParser(description=__doc__)
parser.add_argument("fname", metavar="path/to/figure[/]", nargs="*", help="Location for generated figure, will append __file__ if a directory. If not supplied, figure is shown.", default=os.getcwd())
args = parser.parse_args()
default_extensions = ['.pdf', '.png']
default_names = ['waveforms', 'correlation']
if args.fname == 'none':
args.fname = None
###
figs = main(separate_figures=True, fig_kwargs=dict(figsize=(6,3)))
### Save or show figures
if not args.fname:
# empty list, False, None
plt.show()
else:
for i, f in enumerate(figs):
f.tight_layout()
if len(args.fname) == 1 and len(figs) != 1:
for ext in default_extensions:
f.savefig(path.join(args.fname[0], default_names[i]) + "." + ext, transparent=True)
else:
f.savefig(args.fname[i], transparent=True)

12
figures/pyfiglib.py Normal file
View file

@ -0,0 +1,12 @@
"""
Various functions usable by multiple scripts
"""
from matplotlib import rcParams
import matplotlib.pyplot as plt
import os.path as path
import numpy as np
rcParams["text.usetex"] = True
rcParams["font.family"] = "serif"
rcParams["font.size"] = "12"
rcParams["grid.linestyle"] = 'dotted'