Add test for input dimension (#299)

Co-authored-by: Daniele Nerini <daniele.nerini@gmail.com>

Author: katelbach

pysteps/tests/test_extrapolation_semilagrangian.py

@@ -24,6 +24,26 @@ def test_semilagrangian():
     assert_array_almost_equal(result, expected)
 
 
+def test_wrong_input_dimensions():
+    p_1d = np.ones(8)
+    p_2d = np.ones((8, 8))
+    p_3d = np.ones((8, 8, 2))
+    v_2d = np.ones((8, 8))
+    v_3d = np.stack([v_2d, v_2d])
+
+    num_timesteps = 1
+
+    invalid_inputs = [
+        (p_1d, v_3d),
+        (p_2d, v_2d),
+        (p_3d, v_2d),
+        (p_3d, v_3d),
+    ]
+    for precip, velocity in invalid_inputs:
+        with pytest.raises(ValueError):
+            extrapolate(precip, velocity, num_timesteps)
+
+
 def test_ascending_time_step():
     precip = np.ones((8, 8))
     v = np.ones((8, 8))