Script for generating multi-shell diffusion gradient tables for Siemens platforms (Note: executes against MRtrix3 >= 3.0.0 version)

dvsgen

#!/usr/bin/env python
import math, os, sys



def usage(cmdline): #pylint: disable=unused-variable
  cmdline.set_author('Robert E. Smith (robert.smith@florey.edu.au)')
  cmdline.set_synopsis('Generate diffusion gradient tables in the DVS format for Siemens scanners')
  cmdline.add_argument('sets', type=int, help='The number of sets into which to break the scheme')
  cmdline.add_argument('bzeros', type=int, help='Total number of b=0 images')
  cmdline.add_argument('shells', nargs='+', type=int, metavar='<bvalue volumes>', help='Additional b-value shells, provided as bvalue-volumecount pairs')
  cmdline.add_argument('output', help='The output DVS file')
  cmdline.add_argument('-export_grad_mrtrix', metavar='file', help='Export the gradient table in MRtrix3 format')



def execute(): #pylint: disable=unused-variable
  from mrtrix3 import MRtrixError #pylint: disable=no-name-in-module, import-outside-toplevel
  from mrtrix3 import app, matrix, path, run #pylint: disable=no-name-in-module, import-outside-toplevel

  app.check_output_path(app.ARGS.output)

  if app.ARGS.sets < 1:
    raise MRtrixError('Number of sets must be at least 1')
  if len(app.ARGS.shells) % 2:
    raise MRtrixError('Must provide shell information as bvalue - volumecount pairs')
  if app.ARGS.bzeros / app.ARGS.sets < 2:
    raise MRtrixError('Require at least two b=0 volumes per set')

  app.make_scratch_dir()
  app.goto_scratch_dir()

  class Shell(object):
    def __init__(self, sets, bvalue, volumes):
      self.bvalue = bvalue
      self.volumes = volumes
      self.dirgen_path = 'dirs_' + str(bvalue) + '.txt'
      self.dirsplit_paths = [ os.path.splitext(self.dirgen_path)[0] + '_' + str(index) + '.txt' for index in range(0, sets) ] if sets > 1 else [ self.dirgen_path ]
      self.dirflip_paths = [ os.path.splitext(filepath)[0] + '_flip.txt' for filepath in self.dirsplit_paths ]
      self.dirorder_paths = [ os.path.splitext(filepath)[0] + '_order.txt' for filepath in self.dirflip_paths ]

  shells = [ ]

  bmax = max(app.ARGS.shells[0::2])

  for bvalue, volumes in zip(app.ARGS.shells[0::2], app.ARGS.shells[1::2]):
    shell = Shell(app.ARGS.sets, int(bvalue), int(volumes))
    # Generate the directions
    run.command('dirgen ' + str(shell.volumes) + ' ' + shell.dirgen_path + ' -niter 1M -restarts 1K -cartesian')
    # Split the directions into discrete subsets
    if app.ARGS.sets > 1:
      run.command('dirsplit ' + shell.dirgen_path + ' ' + ' '.join(shell.dirsplit_paths) + ' -cartesian')
    for splitfile, flipfile, orderfile in zip(shell.dirsplit_paths, shell.dirflip_paths, shell.dirorder_paths):
      # Flip the directions to balance eddy currents
      run.command('dirflip ' + splitfile + ' ' + flipfile + ' -cartesian')
    shells.append(shell)

  # Manual combination of the different sets
  grads = [ ]
  for set_index in range(0, app.ARGS.sets):

    # Make use of dirmerge instead of dirorder:
    #   this allows for some control of eddy current accumulation via unipolar model
    # Don't include b=0 volumes
    dirmerge_path = 'dirmerge_' + str(set_index) + '.txt'
    run.command('dirmerge -unipolar 0.5 1 '
                + ' '.join(str(item.bvalue) + ' ' + item.dirflip_paths[set_index] for item in shells)
                + ' ' + dirmerge_path)
    grad = [row[0:4] for row in matrix.load_matrix(dirmerge_path)]

    # How many b=0 images in this set?
    bzero_count = round(app.ARGS.bzeros * (set_index+1) / float(app.ARGS.sets)) - round(app.ARGS.bzeros * set_index / float(app.ARGS.sets))

    # Add the b=0 volumes
    # If at least 3 b=0 volumes, have a pair at the start of the acquisition
    if bzero_count >= 3:
      bzero_volume_increment = len(grad) / (bzero_count-2)
      new_grad = [[0, 0, 0, 0], [0, 0, 0, 0]]
    else:
      bzero_volume_increment = len(grad) / (bzero_count-1)
      new_grad = [[0, 0, 0, 0]]
    bzero_accumulator = 0.0
    bzeros_to_insert = bzero_count - len(new_grad)
    for bzero_index in range(0, bzeros_to_insert):
      bzero_accumulator += bzero_volume_increment
      volumes_from_grad = round(bzero_accumulator)
      new_grad.extend(grad[0:volumes_from_grad])
      grad = grad[volumes_from_grad:]
      new_grad.append([0, 0, 0, 0])
      bzero_accumulator -= volumes_from_grad
    grad = new_grad

    # We export the file at this level, since there's one file for each set
    if app.ARGS.export_grad_mrtrix:
      filename = path.from_user(app.ARGS.export_grad_mrtrix, False)
      if app.ARGS.sets > 1:
        filename = os.path.splitext(filename)[0] + '_' + str(set_index+1) + (os.path.splitext(filename)[1] if len(os.path.splitext(filename))>1 else '.b')
      # TODO Custom export
      matrix.save_matrix(filename, grad)

    volume_index_digits = math.floor(math.log10(len(grad)-1))

    dvs = '[directions=' + str(len(grad)-1) + ']\r\n'
    dvs += 'CoordinateSystem = xyz\r\n'
    dvs += 'Normalisation = none\r\n'
    # DVS file omits the first b=0 volume
    for index, line_in in enumerate(grad[1:]):
      v = line_in[0:3]
      b = line_in[3]
      # Scale vector based on its ratio to the max b-value
      v = [value * math.sqrt(b) / math.sqrt(bmax) for value in v]
      dvs += 'Vector[' + str(index) + '] ' + (' ' * (volume_index_digits - math.floor(math.log10(max([1, index]))))) + '= ( ' \
             + ', '.join("{:18.15f}".format(value) for value in v) \
             + ' )\r\n'

    filename = path.from_user(app.ARGS.output, False)
    if app.ARGS.sets > 1:
      filename = os.path.splitext(filename)[0] + '_' + str(set_index+1) + '.dvs'
    with open(filename, 'w') as f:
      f.write(dvs)


# Execute the script
import mrtrix3
mrtrix3.execute() #pylint: disable=no-member