diff --git a/tests/test_004_wakes_and_impedances.py b/tests/test_004_wakes_and_impedances.py new file mode 100644 index 0000000..306b4d1 --- /dev/null +++ b/tests/test_004_wakes_and_impedances.py @@ -0,0 +1,65 @@ +import os, sys +import numpy as np +import matplotlib.pyplot as plt + +sys.path.append('../wakis') +from wakis import WakeSolver as wk + +class TestImpedancesAndWakes: + def test_sin_wake(self): + fr = 0.5e9 + A = 100 + t = np.linspace(0, 100*1e-9, 3000) + wake = A*np.sin(2*np.pi*fr*t) + + f, Z = wk.calc_impedance_from_wake([t, wake]) + tt, wwake = wk.calc_wake_from_impedance([f, Z], samples=3000) + ff, Zz = wk.calc_impedance_from_wake([tt, wwake],) + ttt, wwwake = wk.calc_wake_from_impedance([ff,Zz], samples=3000) + + assert np.allclose(wake, wwake, atol=1), '1st Transformed wake failed' + assert np.allclose(wake, wwwake, atol=1), '2nd Transformed wake failed' + + def test_sin_impedance(self): + fr = 0.5e9 + A = 100 + t = np.linspace(0, 100*1e-9, 3000) + wake = A*np.sin(2*np.pi*fr*t) + + f, Z = wk.calc_impedance_from_wake([t, wake]) + tt, wwake = wk.calc_wake_from_impedance([f, Z], samples=3000) + ff, Zz = wk.calc_impedance_from_wake([tt, wwake],) + + assert np.allclose(np.max(np.abs(Z)), A, atol=1), '1st Transformed impedance Max. failed' + assert np.allclose(np.max(np.abs(Zz)), A, atol=1), '2nd Transformed impedance Max. failed' + + assert np.allclose(f[np.argmax(Z)], fr, atol=1e6), '1st Transformed impedance fr failed' + assert np.allclose(ff[np.argmax(Zz)], fr, atol=1e6), '2nd Transformed impedance fr. failed' + + def plot_sin(self): + fr = 0.5e9 + A = 100 + t = np.linspace(0, 100*1e-9, 3000) + wake = A*np.sin(2*np.pi*fr*t) + + f, Z = wk.calc_impedance_from_wake([t, wake]) + tt, wwake = wk.calc_wake_from_impedance([f, Z], samples=3000) + ff, Zz = wk.calc_impedance_from_wake([tt, wwake],) + ttt, wwwake = wk.calc_wake_from_impedance([ff,Zz], samples=3000) + + fig, (ax1,ax2) = plt.subplots(2,1) + ax1.plot(t, wake, '-g', alpha=0.8, label='analytic') + ax1.plot(tt, wwake, '--r', alpha=0.5, label='calc') + ax1.plot(ttt, wwwake, '--b', alpha=0.5, label='calc, iter2') + ax1.set_xlabel('time [s]') + + ax2.plot([fr,fr], [0., A], '-g', alpha=0.8, label='analytic') + ax2.plot(f, np.abs(Z), '--r', alpha=0.5, label='calc') + ax2.plot(ff, np.abs(Zz), '--b', alpha=0.5, label='calc, iter2') + ax2.set_xlabel('frequency [Hz]') + ax2.legend() + + fig.tight_layout() + fig.savefig('test_004_sin.png') + plt.show() +