simulate helical filament/tube segments with random number of segments in each helical filament/tube, rotation around helical axis, inplane direction and polarity, out-of-plane tilt, and center shift

simulateHelicalSegments.py

import os, sys, random, math
import pandas as pd
import EMAN2

try:
    mapFile, lstFile, nHelices, twist, maxTilt, maxShift = sys.argv[1:]
    projFile = os.path.splitext(lstFile)[0]+".mrcs"
    nHelices = int(nHelices)
    twist = float(twist)
    maxTilt = float(maxTilt)
    maxShift = float(maxShift)
except:
    print("Usage: %s <input map file> <output lst file> <number of helices> <inter-segment twist in degrees> <max tilt in degree> <max shift in pixel>" % (sys.argv[0]))
    sys.exit(-1)

m = EMAN2.EMData(mapFile)
nx, ny = m.get_xsize(), m.get_ysize()

pids = []
eulers = []
eulersTrue = []
centers = []
centersTrue = []
helicaltubes = []

alt = 90
pid = 0
for hi in range(nHelices):
    az0 = random.random()*360   # random rotation around helical axis
    alt_tilt = alt + (2*random.random()-1)*maxTilt   # random out-of-plane tilt
    phi = random.random()*360   # random inplane direction
    phi_polarity = math.fmod(phi + random.choice((0, 180)), 360)    # simulate unknown polarity
    nSegments = random.randint(1, 10)   # random number of segments in each helix
    for si in range(nSegments):
       print("Helix %d/%d - segment %d/%d" % (hi, nHelices, si, nSegments))
       az = az0 + si*twist
       t = EMAN2.Transform({"type":"eman", "az":az, "alt":alt_tilt, "phi":phi_polarity})
       dx = (2*random.random()-1)*maxShift  # random center shift
       dy = (2*random.random()-1)*maxShift
       t.set_trans(dx, dy)
       proj = m.project("standard", t)
       proj.write_image(projFile, pid)
       pids.append(pid)
       pid+=1
       eulers.append(str(alt)+","+str(0)+","+str(phi))
       eulersTrue.append(str(alt_tilt)+","+str(az)+","+str(phi_polarity))
       centers.append(str(nx/2)+","+str(ny/2))
       centersTrue.append(str(nx/2+dx)+","+str(ny/2+dy))
       helicaltubes.append(hi)
       
data = pd.DataFrame()
data["pid"] = pids
data["filename"] = projFile
data["euler"] = eulers
data["eulerTrue"] = eulersTrue
data["center"] = centers
data["centerTrue"] = centersTrue
data["helicaltube"] = helicaltubes

def dataframe2lst(data, lstFile):
    import numpy as np
    keys = list(data)
    keys.remove("pid")
    keys.remove("filename")

    lines = data['pid'].astype(str) + '\t' + data['filename']
    for k in keys:
        mask = data[k].notnull()
        lines[mask] += '\t' + k + '=' + data[k][mask].astype(str)

    maxlen = max(lines.str.len())
    lines = lines.str.pad(maxlen, side='right')

    with open(lstFile, "w") as lstfp:
        lstfp.write("#LSX\n#If you edit this file, you MUST rerun lstfast.py on it before using it!\n# %d\n" % (maxlen+1))
        lstfp.write('\n'.join(lines))
        lstfp.write('\n')

dataframe2lst(data, lstFile)