added tests and stuff

dipy · Nov 20, 2015 · da4d661 · da4d661
1 parent ba606b2
commit da4d661
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/dipy/denoise/tests/test_noise_estimate.py b/dipy/denoise/tests/test_noise_estimate.py
@@ -41,12 +41,30 @@ def test_piesno():
     # Test using the median as the initial estimation
     initial_estimation = (np.median(sigma) /
                           np.sqrt(2 * _inv_nchi_cdf(1, 1, 0.5)))
+
     sigma, mask = _piesno_3D(rician_noise, N=1, alpha=0.01, l=1, eps=1e-10,
                              return_mask=True,
                              initial_estimation=initial_estimation)
 
     assert_(np.abs(sigma - 50) / sigma < 0.03)
 
+    sigma = _piesno_3D(rician_noise, N=1, alpha=0.01, l=1, eps=1e-10,
+                         return_mask=False,
+                         initial_estimation=initial_estimation)
+    assert_(np.abs(sigma - 50) / sigma < 0.03)
+
+    sigma = _piesno_3D(np.zeros_like(rician_noise), N=1, alpha=0.01, l=1, eps=1e-10,
+                     return_mask=False,
+                     initial_estimation=initial_estimation)
+
+    assert_(np.all(sigma == 0))
+
+    sigma, mask = _piesno_3D(np.zeros_like(rician_noise), N=1, alpha=0.01, l=1, eps=1e-10,
+                 return_mask=True,
+                 initial_estimation=initial_estimation)
+
+    assert_(np.all(sigma == 0))
+    assert_(np.all(mask == 0))
 
 def test_estimate_sigma():