Make a pool of particles in a pseudo-tomogram for use in I3

createPtclCube.py

import argparse
import subprocess
import math
import glob

# Take a number of particles, designed for a single day of data collection
# and using i3paste place them in a large tomogram volume to avoid the limit
# of 1000 tomograms in I3. 

# The size of the tomogram is given by:
# 1. Take the cube root of the number of particles.
# 2. Cube the ceiling of that result this provides ptcl coords in the form
#    (I, J, K) for I = 1,...,NX; J = 1,...,NY; K = 1,...,NZ.
# 3. Multiply that by the particles boxsize plus padding and add one more 
#    value of padding to account for both edges.

# The particles are inserted into the tomogram starting at the bottom left front 
# of the tomogram and are then inserted first along the X axis followed by the
# Y axis and finally slowest along the Z axis.

# Padding should also used to handle the transformation of the particles
# and this is considered when determining the paste coordinates as follows:
# for particle(I, J, K):
# tx(I,J,K) = p + (I - 1) * (nx + p)
# ty(I,J,K) = p + (J - 1) * (ny + p)
# tz(I,J,K) = p + (K - 1) * (nz + p)

# A trf file must also be generated with the new particle centers which are
# given by the particles paste coordinates (tx_i, ty_i, tz_i) for particle i
# and box size nx, ny, nz:
# oj_i = tj_i + ceiling(nj / 2) for j = x, y, z

def getTomogramSize(numPtcls):
  tomoFactor = math.pow(numPtcls, 1.0 / 3.0)
  tomoFactor = math.ceil(tomoFactor)
  return tomoFactor

def makeTomogram(filename, tomoFactor, boxDims, padding):
  tomoSizeX = int((tomoFactor * (boxDims[0] + padding)) + padding)
  tomoSizeY = int((tomoFactor * (boxDims[1] + padding)) + padding)
  tomoSizeZ = int((tomoFactor * (boxDims[2] + padding)) + padding)
  retcode = subprocess.call(["i3new", "-z", "-val", "0", "-s", 
                             str(tomoSizeX), str(tomoSizeY), str(tomoSizeZ),
                             filename])
  return retcode

def getPasteCoords(tomoIndex, boxDims, padding):
  tx = ((tomoIndex[0]) * (boxDims[0] + padding)) + padding
  ty = ((tomoIndex[1]) * (boxDims[1] + padding)) + padding
  tz = ((tomoIndex[2]) * (boxDims[2] + padding)) + padding
  return (int(tx), int(ty), int(tz))

def getCenterCoords(pasteCoords, boxDims):
  ox = pasteCoords[0] + math.ceil(boxDims[0] / 2.)
  oy = pasteCoords[1] + math.ceil(boxDims[1] / 2.)
  oz = pasteCoords[2] + math.ceil(boxDims[2] / 2.)
  return (int(ox), int(oy), int(oz))

def insertParticle(ptclFilename, tomoFilename, tomoIndex, boxDims, padding):
  pasteCoords = getPasteCoords(tomoIndex, boxDims, padding)
  retcode = subprocess.call([
      "i3paste", "-p",  str(pasteCoords[0]), str(pasteCoords[1]),
      str(pasteCoords[2]), ptclFilename, tomoFilename])
  return retcode

def getParticleTransform(posFilename, subset, tomoIndex, boxDims, padding):
  trf = ""
  with open(posFilename) as pos:
    oldI3Trf = pos.readline().rstrip('\n').split()
    pasteCoords = getPasteCoords(tomoIndex, boxDims, padding)
    centerCoords = getCenterCoords(pasteCoords, boxDims)
    oldShifts = [math.modf(float(oldI3Trf[4])), math.modf(float(oldI3Trf[5])),
                 math.modf(float(oldI3Trf[6]))]

    trf += "{} ".format(subset)
    trf += "{:d} {:d} {:d} ".format(centerCoords[0] + int(oldShifts[0][1]),
                                    centerCoords[1] + int(oldShifts[1][1]),
                                    centerCoords[2] + int(oldShifts[2][1]))

    trf += "{:14.10f} {:14.10f} {:14.10f} ".format(oldShifts[0][0],
                                                   oldShifts[1][0],
                                                   oldShifts[2][0])

    rotm = subprocess.check_output(["i3euler", oldI3Trf[7], oldI3Trf[8],
                                    oldI3Trf[9]])

    rotm = rotm.split()
    invRotm = subprocess.check_output(["i3mat", "transp", rotm[0], rotm[1],
                                       rotm[2], rotm[3], rotm[4], rotm[5],
                                       rotm[6], rotm[7], rotm[8]])

    invRotm = [ float(x) for x in invRotm.split() ]
    trf += ("{:14.10f} " * 9).format(*invRotm)
    trf += "\n"
  return trf  

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument('basename', help='Basename for particle filenames')
  parser.add_argument('subsetName', help='Subset name for use in trf file.')
  parser.add_argument('--boxSizeX', '-x', type=int, help='Particle size X',
                      required=True)
  parser.add_argument('--boxSizeY', '-y', type=int, help='Particle size Y',
                      required=True)
  parser.add_argument('--boxSizeZ', '-z', type=int, help='Particle size Z',
                      required=True)
  parser.add_argument('--padding', '-p', type=int, help='Padding to use',
                      required=True)
  args = parser.parse_args()

  ptclList = glob.glob(args.basename + '*.mrc')

  tomoFactor = getTomogramSize(len(ptclList))
  boxDims = (args.boxSizeX, args.boxSizeY, args.boxSizeZ)
  tomoName = args.basename + '_pool.mrc'
  makeTomogram(tomoName, tomoFactor, boxDims, args.padding)
  trfData = ""

  for k in range(tomoFactor):
    for j in range(tomoFactor):
      for i in range(tomoFactor):
        tomoIndex = (i,j,k)
        ptclIndex = i + (j * tomoFactor) + (k * tomoFactor * tomoFactor)
        if ptclIndex < len(ptclList):
          ptclName = ptclList[ptclIndex]
          posFile = ptclName[:-4] + '.pos'
          insertParticle(ptclName, tomoName, tomoIndex, boxDims, args.padding)
          trf = getParticleTransform(posFile, args.subsetName, tomoIndex,
                                     boxDims, args.padding)
          trfData += trf

  with open(args.subsetName + '.trf', 'w') as trfFile:
    trfFile.write(trfData)