add noise, add test

Author: oesteban

nipype/interfaces/dipy/simulate.py

@@ -29,7 +29,7 @@
     have_dipy = False
 else:
     import numpy as np
-    from dipy.sims.voxel import (multi_tensor,
+    from dipy.sims.voxel import (multi_tensor, add_noise,
                                  all_tensor_evecs)
     from dipy.core.gradients import gradient_table
 
@@ -60,6 +60,7 @@ class SimulateMultiTensorInputSpec(BaseInterfaceInputSpec):
                     desc='file with the mask simulated')
     out_bvec = File('bvec.sim', usedefault=True, desc='simulated b vectors')
     out_bval = File('bval.sim', usedefault=True, desc='simulated b values')
+    snr = traits.Int(30, usedefault=True, desc='signal-to-noise ratio (dB)')
 
 
 class SimulateMultiTensorOutputSpec(TraitedSpec):
@@ -143,7 +144,8 @@ def _run_interface(self, runtime):
         np.savetxt(op.abspath(self.inputs.out_bvec), gtab.bvecs.T)
         np.savetxt(op.abspath(self.inputs.out_bval), gtab.bvals.T)
 
-        args = [tuple(np.hstack((r, gtab))) for r in args]
+        snr = self.inputs.snr
+        args = [tuple(np.hstack((r, gtab, snr))) for r in args]
 
         n_proc = self.inputs.n_proc
         if n_proc == 0:
@@ -154,7 +156,8 @@ def _run_interface(self, runtime):
         except TypeError:
             pool = Pool(processes=n_proc)
 
-        iflogger.info('Starting simulation of %d voxels' % len(args))
+        iflogger.info(('Starting simulation of %d voxels, %d diffusion'
+                       ' directions.') % (len(args), len(gtab.bvals)))
         result = pool.map(_compute_voxel, args)
         ndirs = np.shape(result)[1]
 
@@ -229,6 +232,9 @@ def _compute_voxel(args):
     except IndexError:
         pass
 
+    if snr is not None and snr >= 0:
+        signal[1:] = add_noise(signal[1:], snr, 1)
+
     return signal * S0
 
 

nipype/interfaces/dipy/tests/test_auto_SimulateMultiTensor.py

@@ -0,0 +1,53 @@
+# AUTO-GENERATED by tools/checkspecs.py - DO NOT EDIT
+from nipype.testing import assert_equal
+from nipype.interfaces.dipy.simulate import SimulateMultiTensor
+
+def test_SimulateMultiTensor_inputs():
+    input_map = dict(baseline=dict(mandatory=True,
+    ),
+    bval=dict(),
+    bvalues=dict(usedefault=True,
+    ),
+    bvec=dict(),
+    gradients=dict(),
+    ignore_exception=dict(nohash=True,
+    usedefault=True,
+    ),
+    in_dirs=dict(mandatory=True,
+    ),
+    in_frac=dict(mandatory=True,
+    ),
+    in_mask=dict(),
+    in_vfms=dict(mandatory=True,
+    ),
+    n_proc=dict(usedefault=True,
+    ),
+    num_dirs=dict(usedefault=True,
+    ),
+    out_bval=dict(usedefault=True,
+    ),
+    out_bvec=dict(usedefault=True,
+    ),
+    out_file=dict(usedefault=True,
+    ),
+    out_mask=dict(usedefault=True,
+    ),
+    )
+    inputs = SimulateMultiTensor.input_spec()
+
+    for key, metadata in input_map.items():
+        for metakey, value in metadata.items():
+            yield assert_equal, getattr(inputs.traits()[key], metakey), value
+
+def test_SimulateMultiTensor_outputs():
+    output_map = dict(out_bval=dict(),
+    out_bvec=dict(),
+    out_file=dict(),
+    out_mask=dict(),
+    )
+    outputs = SimulateMultiTensor.output_spec()
+
+    for key, metadata in output_map.items():
+        for metakey, value in metadata.items():
+            yield assert_equal, getattr(outputs.traits()[key], metakey), value
+