3
d-                 @   s   d dl mZmZ d dljZd dlZd dlZ	ddl
mZ ddlmZmZmZmZmZmZ ddlmZ ejd	ZG d
d deZG dd deZG dd deZdd ZdddgdfddZdS )    )Pool	cpu_countN   )logging   )traitsTraitedSpecBaseInterfaceInputSpecFileInputMultiPath	isdefined   )DipyBaseInterfaceznipype.interfacec               @   s<  e Zd ZeedddddZeedddddZeedddddZedddZe	j
d	d
dge	jdddZe	jd(e	je	je	jdddZe	jddddZeddddZedddZedddZedddZe	jddddZe	j
e	jddgdddZeddddZed dd!dZed"dd#dZed$dd%dZe	jddd&dZd'S ))SimulateMultiTensorInputSpecT)existsz%list of fibers (principal directions))	mandatorydesczvolume fraction of each fiberz+volume fractions of isotropic compartimentszmask to simulate data)r   r   g~jth?giUMuO?gq@HF?z%Diffusivity of isotropic compartments)
usedefaultr   _Q[?-C6*?zSingle fiber tensorr   znumber of processeszbaseline T2 signal)r   r   r   zgradients filezinput bvecs filezinput bvals file    zEnumber of gradient directions (when table is automatically generated)i  i  z8list of b-values (when table is automatically generated))valuer   r   zsim_dwi.nii.gzz*output file with fractions to be simluatedzsim_msk.nii.gzzfile with the mask simulatedzbvec.simzsimulated b vectorszbval.simzsimulated b valueszsignal-to-noise ratio (dB)N)r   r   r   )__name__
__module____qualname__r   r
   in_dirsin_fracin_vfmsin_maskr   ListZFloatdiff_isoTuplediff_sfZIntn_procbaseline	gradientsin_bvecin_bvalnum_dirsbvaluesout_fileout_maskout_bvecout_bvalsnr r/   r/   A/tmp/pip-build-7vycvbft/nipype/nipype/interfaces/dipy/simulate.pyr      sX   r   c               @   s<   e Zd ZedddZedddZedddZedddZdS )SimulateMultiTensorOutputSpecTzsimulated DWIs)r   r   z	mask filezsimulated b vectorszsimulated b valuesN)r   r   r   r
   r*   r+   r,   r-   r/   r/   r/   r0   r1   R   s   r1   c               @   s(   e Zd ZdZeZeZdd Zdd Z	dS )SimulateMultiTensora  
    Interface to MultiTensor model simulator in dipy
    http://nipy.org/dipy/examples_built/simulate_multi_tensor.html

    Example
    -------

    >>> import nipype.interfaces.dipy as dipy
    >>> sim = dipy.SimulateMultiTensor()
    >>> sim.inputs.in_dirs = ['fdir00.nii', 'fdir01.nii']
    >>> sim.inputs.in_frac = ['ffra00.nii', 'ffra01.nii']
    >>> sim.inputs.in_vfms = ['tpm_00.nii.gz', 'tpm_01.nii.gz',
    ...                       'tpm_02.nii.gz']
    >>> sim.inputs.baseline = 'b0.nii'
    >>> sim.inputs.in_bvec = 'bvecs'
    >>> sim.inputs.in_bval = 'bvals'
    >>> sim.run()                                   # doctest: +SKIP
    c       %         s   ddl m} t| jjrNt| jjrNtj| jj}tj| jjj}|||}nt	| jj
| jj}t|j}tjtj| jj|jj tjtj| jj|j tj| jj}|j}|j}	|j}
t| jj}t| jj|krtdt| jj}tjtj| jjj }|dkr|dtj!f }tj"|dd}t| jj#rftj$tj| jj#j }d||d	k< d	||dk < ntj%|	}d||d	k< tj&|d	d}t||dk }tj| jj}tj'tj|j }tj&|d	d}|dkr|dtj!f }x&t(|D ]|df  |9  < qW tj"|dd}xtt(d|D ]ftj)|dkrxtj*|}||dk d ||dk< |df  |8  < d	||d	k < tj"|dd}q"W x,t(|jd D ]|df  |8  < qW tj&|d	d}tj+||fdd}tj,||
d j-d
 tj,tj"|dd|
d j-d |j. }|j/tj0 |j1dd tj,||
|j-tj| jj2 ||dk }d xt(|D ]| jj }tj'tj|j }tj3j4|dddtj!f }d||tj5t6j7k < || }d kr||dk j. ntj8||dk fqdW xt(|D ]|}tj9j:|d}tj3j4|dd}tj;dd	d	g||tj5t6j7k df< d||tj5t6j7k < ||dtj!f  }tj8|fqW t<| jj=}t<| jj> |g|  fddt(|D  }|j? |dk }g }xVt(|D ]J|j@|df jA fddt(|| D ||| | jjBd qW | jjC} | dkr:tD } ytE| dd}!W n  tFk
rj   tE| d}!Y nX tGjHdt|| tj;|!jItJ|}"tj|"d |krtdtj%|	d |	d |	d |f}#|"|#|dk< |j. }$|$j/tjK |$j1dd tj,|#jLtjK|
|$j-tj| jjM |S )Nr   )gradient_tablez7Number of sticks and their volume fractions must match.r   .r   )Zaxisg      ?g        zfractions.nii.gzztotal_vf.nii.gzmmsecc                s   g | ]} | gd  qS )r   r/   ).0d)ba_evalsr/   r0   
<listcomp>   s    z6SimulateMultiTensor._run_interface.<locals>.<listcomp>c                s$   g | ]}t  ||d  f qS )r   )tuple)r6   j)dirsir/   r0   r9      s    )	fractionssticksr%   mevalsS0r.   2   )	processesZmaxtasksperchild)rC   z:Starting simulation of %d voxels, %d diffusion directions.z3Computed directions do not match numberof b-values.r   )Ndipy.core.gradientsr3   r   inputsr'   r&   npZloadtxtT_generate_gradientsr(   r)   lenbvalsZsavetxtopabspathr,   bvecsr-   nbloadr$   headershapeZaffiner   r   RuntimeErrorr   ZsqueezeZconcat_imagesZdataobjZnewaxissumr   Z
asanyarrayzerosZclipZ
nan_to_numrangeany
zeros_likeZconcatenateZNifti1Imageto_filenamecopyZset_data_dtypeZuint8Zset_xyzt_unitsr+   ZlinalgZnormZfinfofloatepshstackrandomZrandnarraylistr"   r    Z	get_fdataappendtolistr.   r#   r   r   	TypeErrorIFLOGGERinfomap_compute_voxelfloat32Zastyper*   )%selfZruntimer3   rK   rN   gtabndirsZb0_imhdrrR   ZaffZnsticksZnballsZvfsZtotal_vfZmskZnvoxZffsimffsZtotal_fferrorsr>   ZmhdrZfracsffdwr7   Zsf_evalsr@   Zb0argsr#   poolresultsignalZsimhdrr/   )r8   r<   r=   r0   _run_interfacep   s    





$ 
z"SimulateMultiTensor._run_interfacec             C   sX   | j  j }tj| jj|d< tj| jj|d< tj| jj|d< tj| jj|d< |S )Nr*   r+   r,   r-   )	Z_outputsgetrL   rM   rF   r*   r+   r,   r-   )ri   outputsr/   r/   r0   _list_outputs  s    z!SimulateMultiTensor._list_outputsN)
r   r   r   __doc__r   Z
input_specr1   Zoutput_specrv   ry   r/   r/   r/   r0   r2   Y   s    $r2   c             C   s   ddl m} | d }| d }tj|jtjd}tj|}|dkrtj|| d }| d dkrf| d nd	}y(||| d
 | d | d ||d\}}W n tk
r   Y nX |j	 S )a=  
    Simulate DW signal for one voxel. Uses the multi-tensor model and
    three isotropic compartments.

    Apparent diffusivity tensors are taken from [Alexander2002]_
    and [Pierpaoli1996]_.

    .. [Alexander2002] Alexander et al., Detection and modeling of non-Gaussian
      apparent diffusion coefficient profiles in human brain data, MRM
      48(2):331-340, 2002, doi: `10.1002/mrm.10209
      <https://doi.org/10.1002/mrm.10209>`_.
    .. [Pierpaoli1996] Pierpaoli et al., Diffusion tensor MR imaging
      of the human brain, Radiology 201:637-648. 1996.
    r   )multi_tensorr>   r%   )Zdtypeg        d   r.   Nr@   rA   r?   )rA   Zanglesr>   r.   )
Zdipy.sims.voxelr{   rG   rX   rK   rh   rT   r_   	Exceptionrb   )rr   r{   rm   rj   ru   Zsf_vfr.   _r/   r/   r0   rg     s&    
rg   @   i  i  c             C   s,  ddl }ddlm}m}m} ddlm} |j|jj	|  }d|j |jj	|  }	|||	d}
||
d\}}|j
|j }|j}|j }|j|jd |d  }x@|dd D ]0}|j||f}|j|||j|jd  f}qW x@td|D ]2}|j }|jdd |j }|jd|jd	 qW |||S )
z
    Automatically generate a `gradient table
    <http://nipy.org/dipy/examples_built/gradients_spheres.html#example-gradients-spheres>`_

    r   N)disperse_chargesSphere
HemiSphere)r3   r   )thetaphii  r   r   )numpyZdipy.core.spherer   r   r   rE   r3   pir^   ZrandZ
atleast_1drb   verticesrZ   ZonesrR   Zvstackr]   rV   insertrU   )rk   valuesZnb0srG   r   r   r   r3   r   r   Zhsph_initialZhsph_updatedZ	potentialr   rN   rK   vr=   r/   r/   r0   rI   G  s(    "rI   )multiprocessingr   r   os.pathpathrL   r   rG   ZnibabelrO    r   baser   r   r	   r
   r   r   r   	getLoggerrd   r   r1   r2   rg   rI   r/   r/   r/   r0   <module>   s   
 
< E*