diff --git a/b_asic/signal_generator.py b/b_asic/signal_generator.py index ec3ceb6b17b9810a0ba8fd023b818088b1aafb72..e9bc99b4cc5faeb4d7260bdafb727cc4d124faf8 100644 --- a/b_asic/signal_generator.py +++ b/b_asic/signal_generator.py @@ -13,7 +13,7 @@ if you want more information. from math import pi, sin from numbers import Number -from typing import Optional, Sequence +from typing import Optional, Sequence, Union import numpy as np @@ -408,3 +408,79 @@ class _DivGenerator(SignalGenerator): else f"{self._b}" ) return f"{a} / {b}" + + +class Upsample(SignalGenerator): + """ + Signal generator that upsamples the value of a signal generator or a sequence. + + *factor* - 1 zeros are inserted between every value. If a sequence, it will + automatically be zero-padded. + + Parameters + ---------- + data : SignalGenerator or 1-D array + The signal generator or array to upsample. + factor : int + Upsampling factor. + phase : int, default 0 + The phase of the upsampling. + """ + + def __init__( + self, + data: Union[SignalGenerator, Sequence[complex]], + factor: int, + phase: int = 0, + ) -> None: + if not isinstance(data, SignalGenerator): + data = ZeroPad(data) + self._generator = data + self._factor = factor + self._phase = phase + + def __call__(self, time: int) -> complex: + index = (time - self._phase) / self._factor + if int(index) == index: + return self._generator(int(index)) + else: + return 0.0 + + def __repr__(self) -> str: + return f"Upsample({self._generator!r}, {self._factor}, {self._phase})" + + +class Downsample(SignalGenerator): + """ + Signal generator that downsamples the value of a signal generator or a sequence. + + Return every *factor*:th value. If a sequence, it will automatically be zero-padded. + + Parameters + ---------- + data : SignalGenerator or 1-D array + The signal generator or array to downsample. + factor : int + Downsampling factor. + phase : int, default 0 + The phase of the downsampling. + """ + + def __init__( + self, + data: Union[SignalGenerator, Sequence[complex]], + factor: int, + phase: int = 0, + ) -> None: + if not isinstance(data, SignalGenerator): + data = ZeroPad(data) + self._generator = data + self._factor = factor + self._phase = phase + + def __call__(self, time: int) -> complex: + index = time * self._factor + self._phase + return self._generator(index) + + def __repr__(self) -> str: + return f"Downsample({self._generator!r}, {self._factor}, {self._phase})" diff --git a/test/test_signal_generator.py b/test/test_signal_generator.py index 787bf0d384d0c8d12be336d755c48d35c3743440..65266a25ae4ab9d7a989eeba2207dd877e5fcbf5 100644 --- a/test/test_signal_generator.py +++ b/test/test_signal_generator.py @@ -5,12 +5,14 @@ import pytest from b_asic.signal_generator import ( Constant, Delay, + Downsample, FromFile, Gaussian, Impulse, Sinusoid, Step, Uniform, + Upsample, ZeroPad, _AddGenerator, _DivGenerator, @@ -284,3 +286,61 @@ def test_fromfile(datadir): with pytest.raises(ValueError, match="could not convert string"): g = FromFile(datadir.join('bad.csv')) + + +def test_upsample(): + g = Upsample([0.4, 0.6], 2) + assert g(-1) == 0 + assert g(0) == 0.4 + assert g(1) == 0 + assert g(2) == 0.6 + assert g(3) == 0.0 + + assert str(g) == "Upsample(ZeroPad([0.4, 0.6]), 2, 0)" + + g = Upsample([0.4, 0.6], 2, 1) + assert g(-1) == 0 + assert g(0) == 0.0 + assert g(1) == 0.4 + assert g(2) == 0.0 + assert g(3) == 0.6 + assert g(4) == 0.0 + + assert str(g) == "Upsample(ZeroPad([0.4, 0.6]), 2, 1)" + + g = Upsample([0.4, 0.6], 3) + assert g(-1) == 0.0 + assert g(0) == 0.4 + assert g(1) == 0.0 + assert g(2) == 0.0 + assert g(3) == 0.6 + assert g(4) == 0.0 + + assert str(g) == "Upsample(ZeroPad([0.4, 0.6]), 3, 0)" + + g = Upsample(Sinusoid(0.5, 0.25), 3) + assert g(0) == pytest.approx(sqrt(2) / 2) + assert g(1) == 0.0 + assert g(2) == 0.0 + assert g(3) == pytest.approx(sqrt(2) / 2) + assert g(4) == 0.0 + assert g(5) == 0.0 + assert g(6) == pytest.approx(-sqrt(2) / 2) + + assert str(g) == "Upsample(Sinusoid(0.5, 0.25), 3, 0)" + + +def test_downsample(): + g = Downsample([0.4, 0.6], 2) + assert g(-1) == 0 + assert g(0) == 0.4 + assert g(1) == 0.0 + assert g(2) == 0.0 + assert str(g) == "Downsample(ZeroPad([0.4, 0.6]), 2, 0)" + + g = Downsample([0.4, 0.6], 2, 1) + assert g(-1) == 0 + assert g(0) == 0.6 + assert g(1) == 0.0 + assert g(2) == 0.0 + assert str(g) == "Downsample(ZeroPad([0.4, 0.6]), 2, 1)"