Computes the voxelwise btable from initial bvecs/bvals and calc_grad_perc_dev output (from fullwarp output of gradunwrap.py)

compute_spatial_btable.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-

# first draft of script to compute voxelwise bvec and bval after GNL
# work on the output of calc_grad_perc_dev
# calc_grad_perc_dev works on the fullwarp output of gradunwrap.py

import numpy as np
import nibabel as nib
import argparse


DESCRIPTION = """
Compute bvecs and bvals at each voxel according to gradient percent deviation map.
"""

def buildArgsParser():
    p = argparse.ArgumentParser(description=DESCRIPTION)
    p.add_argument('bvecs', action='store', type=str,
                            help='Path of the bvecs file')

    p.add_argument('bvals', action='store', type=str,
                            help='Path of the bvals file')

    p.add_argument('devX', action='store', type=str,
                            help='Path of X gradient deviation file')

    p.add_argument('devY', action='store', type=str,
                            help='Path of Y gradient deviation file')

    p.add_argument('devZ', action='store', type=str,
                            help='Path of Z gradient deviation file')

    p.add_argument('outfile', action='store', type=str,
                            help='Path of the output btable file')


    p.add_argument('--mask', dest='mask', action='store', type=str,
                            help='Path of the mask file. If none given, computes on the full volume.')

    return p


def main():

    # parse inpout
    parser = buildArgsParser()
    args = parser.parse_args()

    bvecsfile = args.bvecs
    bvalsfile = args.bvals
    devXfile = args.devX
    devYfile = args.devY
    devZfile = args.devZ
    outfile = args.outfile
    maskfile = args.mask


    # load data
    bvecs = np.genfromtxt(bvecsfile)
    if bvecs.shape[1] != 3:
        bvecs = bvecs.T

    bvals = np.genfromtxt(bvalsfile)

    devX_img = nib.load(devXfile)
    devY_img = nib.load(devYfile)
    devZ_img = nib.load(devZfile)

    devX = devX_img.get_data()
    devY = devY_img.get_data()
    devZ = devZ_img.get_data()

    if maskfile is None:
        mask = np.ones(devX.shape[:3])
        print('No mask used, beware of inaccurate volume boundary.')
    else:
        mask = nib.load(maskfile).get_data()


    # convert percentage to fraction
    devX *= 0.01
    devY *= 0.01
    devZ *= 0.01


    # make the grad non lin tensor
    dev = np.concatenate((devX[...,None], devY[...,None], devZ[...,None]), axis=4)

    # "q" gradient
    bscaled_grad = (bvecs*np.sqrt(bvals)[:,None])


    new_bscaled_grad = np.zeros(mask.shape+bvecs.shape)
    for idx in np.ndindex(mask.shape):
        if mask[idx]:
            # distort bvecs
            new_bscaled_grad[idx] = bscaled_grad.dot(dev[idx])

    # renormalize gradient direction
    new_b = np.linalg.norm(new_bscaled_grad, axis=4)
    new_grad = new_bscaled_grad / new_b[...,None]
    # nan removal (from the b division at b0)
    new_grad[...,bvals<10,:] = 0

    # new b is the norm squared of the distorded gradient
    new_b = new_b**2

    # make the (X,Y,Z,N,4) btable
    btable = np.concatenate((new_grad, new_b[...,None]), axis=4)

    # save
    output_img = nib.Nifti1Image(btable, devX_img.affine, devX_img.header)
    nib.save(output_img, outfile)


if __name__ == '__main__':
    main()