Test fft-generators (#213)

* Add basic tests for fft noise gen

* Fix example

* Aesthetics

examples/plot_noise_generators.py

@@ -3,11 +3,11 @@
 Generation of stochastic noise
 ==============================
 
-This example script shows how to run the stochastic noise field generators 
+This example script shows how to run the stochastic noise field generators
 included in pysteps.
 
 These noise fields are used as perturbation terms during an extrapolation
-nowcast in order to represent the uncertainty in the evolution of the rainfall 
+nowcast in order to represent the uncertainty in the evolution of the rainfall
 field.
 """
 
@@ -67,7 +67,7 @@
 
 # Compute the observed and fitted 1D PSD
 L = np.max(Fp["input_shape"])
-if L % 2 == 0:
+if L % 2 == 1:
     wn = np.arange(0, int(L / 2) + 1)
 else:
     wn = np.arange(0, int(L / 2))

pysteps/tests/test_noise_fftgenerators.py

@@ -0,0 +1,53 @@
+import numpy as np
+
+from pysteps.noise import fftgenerators
+from pysteps.tests.helpers import get_precipitation_fields
+
+
+PRECIP = get_precipitation_fields(
+    num_prev_files=0,
+    num_next_files=0,
+    return_raw=False,
+    metadata=False,
+    upscale=2000,
+)
+PRECIP = PRECIP.filled()
+
+
+def test_noise_param_2d_fft_filter():
+
+    fft_filter = fftgenerators.initialize_param_2d_fft_filter(PRECIP)
+
+    assert isinstance(fft_filter, dict)
+    assert all([key in fft_filter for key in ["field", "input_shape", "model", "pars"]])
+
+    out = fftgenerators.generate_noise_2d_fft_filter(fft_filter)
+
+    assert isinstance(out, np.ndarray)
+    assert out.shape == PRECIP.shape
+
+
+def test_noise_nonparam_2d_fft_filter():
+
+    fft_filter = fftgenerators.initialize_nonparam_2d_fft_filter(PRECIP)
+
+    assert isinstance(fft_filter, dict)
+    assert all([key in fft_filter for key in ["field", "input_shape"]])
+
+    out = fftgenerators.generate_noise_2d_fft_filter(fft_filter)
+
+    assert isinstance(out, np.ndarray)
+    assert out.shape == PRECIP.shape
+
+
+def test_noise_nonparam_2d_ssft_filter():
+
+    fft_filter = fftgenerators.initialize_nonparam_2d_ssft_filter(PRECIP)
+
+    assert isinstance(fft_filter, dict)
+    assert all([key in fft_filter for key in ["field", "input_shape"]])
+
+    out = fftgenerators.generate_noise_2d_ssft_filter(fft_filter)
+
+    assert isinstance(out, np.ndarray)
+    assert out.shape == PRECIP.shape