101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from scgenerator.math import abs2, wspace
|
|
from scgenerator.physics.units import m_rads
|
|
from scgenerator.physics.pulse import find_lobe_limits
|
|
from scgenerator.spectra import Spectrum
|
|
|
|
|
|
def test_export():
|
|
t = np.linspace(-1e-12, 1e-12, 1024)
|
|
w = wspace(t) + m_rads(800e-9)
|
|
|
|
spec = np.fft.fft(np.exp(-((t / 1e-13) ** 2)))
|
|
spec = Spectrum(spec, w, t)
|
|
|
|
new_spec = Spectrum.from_bytes(bytes(spec))
|
|
assert np.all(new_spec.w == spec.w)
|
|
assert np.all(new_spec.t == spec.t)
|
|
assert np.all(new_spec.l == spec.l)
|
|
assert np.all(new_spec == spec)
|
|
assert new_spec.ifft is spec.ifft
|
|
|
|
t = np.linspace(-1e-12, 1e-12, 512)
|
|
w = np.fft.rfftfreq(len(t), d=t[1] - t[0]) * 2 * np.pi + m_rads(800e-9)
|
|
spec = np.exp(2j * np.pi * np.random.rand(4, 2, 5, 257))
|
|
spec = Spectrum(spec, w, t, np.fft.irfft)
|
|
|
|
new_spec = Spectrum.from_bytes(bytes(spec))
|
|
assert np.all(new_spec.w == spec.w)
|
|
assert np.all(new_spec.t == spec.t)
|
|
assert np.all(new_spec.l == spec.l)
|
|
assert np.all(new_spec == spec)
|
|
assert new_spec.ifft is spec.ifft
|
|
|
|
|
|
def test_center_gravity():
|
|
t = np.linspace(-1e-12, 1e-12, 1024)
|
|
w = wspace(t) + m_rads(800e-9)
|
|
|
|
spec = np.fft.fft(np.exp(-((t / 1e-13) ** 2)))
|
|
spec = Spectrum(spec, w, t)
|
|
assert spec.center_of_gravity == pytest.approx(0)
|
|
|
|
spec = np.fft.fft(
|
|
np.array([np.exp(-((t / 1e-13) ** 2)), np.exp(-(((t - 1e-15) / 1e-13) ** 2))])
|
|
)
|
|
spec = Spectrum(spec, w, t)
|
|
assert spec.center_of_gravity * 1e15 == pytest.approx((0, 1))
|
|
|
|
spec = np.fft.fft(
|
|
np.array(
|
|
[
|
|
[
|
|
np.exp(-((t / 1e-13) ** 2)),
|
|
np.exp(-(((t - 1e-15) / 1e-13) ** 2)),
|
|
np.exp(-(((t + 2e-15) / 1e-13) ** 2)),
|
|
],
|
|
[
|
|
np.exp(-((t / 1e-13) ** 2)),
|
|
np.exp(-(((t + 1e-15) / 1e-13) ** 2)),
|
|
np.exp(-(((t - 2e-15) / 1e-13) ** 2)),
|
|
],
|
|
]
|
|
)
|
|
)
|
|
spec = Spectrum(spec, w, t)
|
|
assert spec.center_of_gravity[0] * 1e15 == pytest.approx((0, 1, -2))
|
|
assert spec.center_of_gravity[1] * 1e15 == pytest.approx((0, -1, 2))
|
|
|
|
|
|
def test_energy():
|
|
t = np.linspace(-1e-12, 1e-12, 1024)
|
|
w = wspace(t) + m_rads(800e-9)
|
|
|
|
spec = np.random.rand(1024)
|
|
spec = Spectrum(spec, w, t)
|
|
assert spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t)
|
|
|
|
spec = np.random.rand(2, 1024)
|
|
spec = Spectrum(spec, w, t)
|
|
assert np.all(spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t, axis=-1))
|
|
assert spec.energy.shape == (2,)
|
|
|
|
spec = np.random.rand(3, 2, 1024)
|
|
spec = Spectrum(spec, w, t)
|
|
assert np.all(spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t, axis=-1))
|
|
assert spec.energy.shape == (3, 2)
|
|
|
|
|
|
def test_mask():
|
|
wl = 800e-9
|
|
bw = 50e-9
|
|
n = 8192
|
|
t = np.linspace(-1e-12, 1e-12, n)
|
|
w = wspace(t) + m_rads(wl)
|
|
spec = Spectrum(np.ones(n) + 0j, w, t)
|
|
mased = spec.mask_wl(wl, bw)
|
|
l, r = find_lobe_limits(mased.wl_disp, mased.wl_int)[1]
|
|
assert r - l == pytest.approx(bw, rel=1e-2, abs=1e-10)
|