Final touches

TorchDSP · Sep 5, 2023 · 4964f48 · 4964f48
1 parent 2fbc788
commit 4964f48
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 54 deletions.
diff --git a/tests/test_transforms_benchmark.py b/tests/test_transforms_benchmark.py
@@ -25,60 +25,60 @@ def generate_data():
 
 
 transforms_list = [
-    # (
-    #     "random_resample_up",
-    #     RandomResample(1.5, num_iq_samples=128, keep_samples=False),
-    #     RandomResample(1.5, num_iq_samples=4096, keep_samples=False),
-    # ),
-    # (
-    #     "random_resample_down",
-    #     RandomResample(0.75, num_iq_samples=128, keep_samples=False),
-    #     RandomResample(0.75, num_iq_samples=4096, keep_samples=False),
-    # ),
+    (
+        "random_resample_up",
+        RandomResample(1.5, num_iq_samples=128, keep_samples=False),
+        RandomResample(1.5, num_iq_samples=4096, keep_samples=False),
+    ),
+    (
+        "random_resample_down",
+        RandomResample(0.75, num_iq_samples=128, keep_samples=False),
+        RandomResample(0.75, num_iq_samples=4096, keep_samples=False),
+    ),
     ("add_noise", AddNoise(-10), AddNoise(-10)),
-    # ("time_varying_noise", TimeVaryingNoise(-30, -10), TimeVaryingNoise(-30, -10)),
-    # (
-    #     "rayleigh_fading",
-    #     RayleighFadingChannel(0.05, (1.0, 0.5, 0.1)),
-    #     RayleighFadingChannel(0.05, (1.0, 0.5, 0.1)),
-    # ),
-    # ("phase_shift", RandomPhaseShift(0.5), RandomPhaseShift(0.5)),
-    # ("time_shift", RandomTimeShift(-100.5), RandomTimeShift(-2.5)),
-    # (
-    #     "time_crop",
-    #     TimeCrop("random", length=128),
-    #     TimeCrop("random", length=4096),
-    # ),
-    # ("time_reversal", TimeReversal(False), TimeReversal(False)),
-    # ("frequency_shift", RandomFrequencyShift(-0.25), RandomFrequencyShift(-0.25)),
-    # (
-    #     "delayed_frequency_shift",
-    #     RandomDelayedFrequencyShift(0.2, 0.25),
-    #     RandomDelayedFrequencyShift(0.2, 0.25),
-    # ),
-    # (
-    #     "oscillator_drift",
-    #     LocalOscillatorDrift(0.01, 0.001),
-    #     LocalOscillatorDrift(0.01, 0.001),
-    # ),
-    # ("gain_drift", GainDrift(0.01, 0.001, 0.1), GainDrift(0.01, 0.001, 0.1)),
-    # (
-    #     "iq_imbalance",
-    #     IQImbalance(3, np.pi / 180, 0.05),
-    #     IQImbalance(3, np.pi / 180, 0.05),
-    # ),
-    # ("roll_off", RollOff(0.05, 0.98), RollOff(0.05, 0.98)),
-    # ("add_slope", AddSlope(), AddSlope()),
-    # ("spectral_inversion", SpectralInversion(), SpectralInversion()),
-    # ("channel_swap", ChannelSwap(), ChannelSwap()),
-    # ("magnitude_rescale", RandomMagRescale(0.5, 3), RandomMagRescale(0.5, 3)),
-    # (
-    #     "drop_samples",
-    #     RandomDropSamples(0.01, 50, ["zero"]),
-    #     RandomDropSamples(0.01, 50, ["zero"]),
-    # ),
-    # ("quantize", Quantize(32, ["floor"]), Quantize(32, ["floor"])),
-    # ("clip", Clip(0.85), Clip(0.85)),
+    ("time_varying_noise", TimeVaryingNoise(-30, -10), TimeVaryingNoise(-30, -10)),
+    (
+        "rayleigh_fading",
+        RayleighFadingChannel(0.05, (1.0, 0.5, 0.1)),
+        RayleighFadingChannel(0.05, (1.0, 0.5, 0.1)),
+    ),
+    ("phase_shift", RandomPhaseShift(0.5), RandomPhaseShift(0.5)),
+    ("time_shift", RandomTimeShift(-100.5), RandomTimeShift(-2.5)),
+    (
+        "time_crop",
+        TimeCrop("random", length=128),
+        TimeCrop("random", length=4096),
+    ),
+    ("time_reversal", TimeReversal(False), TimeReversal(False)),
+    ("frequency_shift", RandomFrequencyShift(-0.25), RandomFrequencyShift(-0.25)),
+    (
+        "delayed_frequency_shift",
+        RandomDelayedFrequencyShift(0.2, 0.25),
+        RandomDelayedFrequencyShift(0.2, 0.25),
+    ),
+    (
+        "oscillator_drift",
+        LocalOscillatorDrift(0.01, 0.001),
+        LocalOscillatorDrift(0.01, 0.001),
+    ),
+    ("gain_drift", GainDrift(0.01, 0.001, 0.1), GainDrift(0.01, 0.001, 0.1)),
+    (
+        "iq_imbalance",
+        IQImbalance(3, np.pi / 180, 0.05),
+        IQImbalance(3, np.pi / 180, 0.05),
+    ),
+    ("roll_off", RollOff(0.05, 0.98), RollOff(0.05, 0.98)),
+    ("add_slope", AddSlope(), AddSlope()),
+    ("spectral_inversion", SpectralInversion(), SpectralInversion()),
+    ("channel_swap", ChannelSwap(), ChannelSwap()),
+    ("magnitude_rescale", RandomMagRescale(0.5, 3), RandomMagRescale(0.5, 3)),
+    (
+        "drop_samples",
+        RandomDropSamples(0.01, 50, ["zero"]),
+        RandomDropSamples(0.01, 50, ["zero"]),
+    ),
+    ("quantize", Quantize(32, ["floor"]), Quantize(32, ["floor"])),
+    ("clip", Clip(0.85), Clip(0.85)),
 ]
 
 

diff --git a/torchsig/transforms/transforms.py b/torchsig/transforms/transforms.py
@@ -2969,7 +2969,7 @@ class Quantize(SignalTransform):
     def __init__(
         self,
         num_levels: IntParameter = UniformDiscreteRD(
-            np.asarray([16, 24, 32, 40, 48, 56, 64])
+            np.asarray([16, 24, 32, 40, 48, 56, 64], dtype=int)
         ),
         round_type: List[str] = (["floor", "middle", "ceiling"]),
     ) -> None:

diff --git a/torchsig/utils/types.py b/torchsig/utils/types.py
@@ -16,6 +16,9 @@ def to_distribution(dist):
         if isinstance(dist, tuple):
             return UniformContinuousRD(dist[0], dist[1])
 
+        if isinstance(dist, list):
+            return UniformDiscreteRD(dist)
+
     def __call__(self, num: int = 1):
         raise NotImplementedError
 
@@ -26,6 +29,9 @@ def __init__(self, constant: float) -> None:
         self.constant = constant
 
     def __call__(self, num: int = 1):
+        if num == 1:
+            return self.constant
+
         return np.repeat(self.constant, repeats=num)