-
Notifications
You must be signed in to change notification settings - Fork 0
/
convert_dmt_aa2relion4.py
165 lines (130 loc) · 6.13 KB
/
convert_dmt_aa2relion4.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
#/usr/bin/python3
# Convert the entire cilia from AA to relion4
# Must run adjustOrigin from AA before
# Add HelicalTubeID now
# HB 08/2022
import numpy as np
import pandas as pd
import argparse, os, re
from eulerangles import euler2euler
from eulerangles import convert_eulers
def write_star_4(dfin, outfile, tableType):
out = open(outfile, 'w')
out.write("# version 30001 from aa\n\n")
out.write("data_" + tableType +"\n\n")
out.write("loop_\n")
for i in range(len(dfin.columns)):
out.write('_rln{:s} #{:d}\n'.format(dfin.columns[i], i+1))
out.write(dfin.to_string(index=False, header=False))
out.write("\n")
out.close()
def preprocess_spider_doc(spiderdoc):
cmd = 'sed -i \'/^ ;/d\' ' + spiderdoc
os.system(cmd)
def preprocess_bstar(starFile):
cmd = 'grep \'^\\s*[0-9]\' ' + starFile + ' > ' + starFile.replace('.star', '.txt')
os.system(cmd)
"""Convert aa doc & star to dynamo table"""
def aa_to_relion(starFile, docFile, tomoName, tomoNo, binFactor, pixelSize, doubletId):
# Read the doc file
header_list=["no", "norec", "phi", "theta", "psi", "OriginX", "OriginY", "OriginZ", "cc"]
df = pd.read_csv(docFile, delim_whitespace=True, names=header_list)
fulldata = df.to_numpy()
# Extract phi, theta, psi (AA format) and reverse sign of phi & psi
eulers_zyz = fulldata[:, 2:5]*-1
eulers_zyz[:,1] = eulers_zyz[:,1]*-1
eulers_dynamo = euler2euler(eulers_zyz, source_axes='zyz', source_intrinsic=True, source_right_handed_rotation=True,
target_axes='zxz', target_intrinsic=True,target_right_handed_rotation=True,invert_matrix=False)
# Read the star file (ignore header for now)
star_header = ["no", "c2", "c3", "c4", "CoordinateX", "CoordinateY", "CoordinateZ", "c8", "c9", "c10", "c11", "c12", "c13", "c14", "c15", "c16"]
df2 = pd.read_csv(starFile, delim_whitespace=True, names=star_header)
fullstar = df2.to_numpy()
# Extract origin
origin = fullstar[:, 4:7]
nrows, ncols = origin.shape
# Hard Code Here
header_list = ["TomoName", "TomoParticleId", "HelicalTubeID", "CoordinateX", "CoordinateY", "CoordinateZ", "OriginXAngst", "OriginYAngst", "OriginZAngst", "AngleRot", "AngleTilt", "AnglePsi", "ClassNumber", "RandomSubset"]
df_relion = pd.DataFrame(columns = header_list)
df_relion['TomoParticleId'] = np.arange(len(df2), dtype=np.int16) + 1
df_relion['HelicalTubeID'] = np.ones(len(df2['CoordinateX']), dtype=np.int16)*doubletId
df_relion['CoordinateX'] = df2['CoordinateX']*binFactor;
df_relion['CoordinateY'] = df2['CoordinateY']*binFactor;
df_relion['CoordinateZ'] = df2['CoordinateZ']*binFactor;
# To adjust originXYZ
df_relion['OriginXAngst'] = np.zeros(len(df_relion['CoordinateX']))
df_relion['OriginYAngst'] = np.zeros(len(df_relion['CoordinateX']))
df_relion['OriginZAngst'] = np.zeros(len(df_relion['CoordinateX']))
# Reset angle for debug
eulers_relion = convert_eulers(eulers_dynamo, source_meta='dynamo', target_meta='warp')
df_relion['AngleRot'] = eulers_relion[:,0]
df_relion['AngleTilt'] = eulers_relion[:,1]
df_relion['AnglePsi'] = eulers_relion[:,2]
df_relion['ClassNumber'] = np.ones(len(df_relion['CoordinateX']), dtype=np.int8)
for i in range(len(df2['CoordinateX'])):
df_relion.loc[i, ('TomoName')] = tomoName
a = np.empty((len(df_relion['CoordinateX']),), dtype=np.int8)
a[::2] = 1
a[1::2] = 2
df_relion['RandomSubset'] = a
return df_relion
if __name__=='__main__':
# get name of input starfile, output starfile, output stack file
print('Script to convert from AxonemeAlign to Relion. HB 2021')
parser = argparse.ArgumentParser(description='Convert doc & star file to Relion 4.0 input file')
parser.add_argument('--i', help='Input list file',required=True)
parser.add_argument('--ostar', help='Output star file',required=True)
parser.add_argument('--angpix', help='Input pixel size',required=True)
parser.add_argument('--bin', help='Bin of current tomo',required=True)
parser.add_argument('--frac_dose', help='Tomo fractional dose',required=True, default=2)
args = parser.parse_args()
listDoublet = open(args.i, 'r')
pixelSize = float(args.angpix)
binFactor = float(args.bin)
tomoList = {}
tomoNo = 0;
df_all = None
# Template for tomo_description
orderList = 'input/order_list.csv'
tomo_header_list = ["TomoName", "TomoTiltSeriesName", "TomoImportCtfFindFile", "TomoImportImodDir", "TomoImportFractionalDose", "TomoImportOrderList", "TomoImportCulledFile"]
df_tomo = pd.DataFrame(columns = tomo_header_list)
for line in listDoublet:
if line.startswith('#'):
continue
record = line.split()
# Check tomo
# This is not so robust for tomoa & tomob name yet
tomoSubName = record[0].replace('_ida_v1', '')
tomoSubName = tomoSubName[:-4]
# Replace a, b, c in case. Not exact more than 3 tomo
tomoName = re.sub('[a-z]$', '', tomoSubName)
doubletId = int(record[1][-1])
if tomoList.get(tomoName) == None:
print(tomoName)
tomoNo += 1
tomoList[tomoName] = tomoNo
df_tomo.loc[tomoNo-1, 'TomoName'] = tomoName
df_tomo.loc[tomoNo-1, 'TomoTiltSeriesName'] = 'tomograms/' + tomoName + '/' + tomoName + '.mrc'
df_tomo.loc[tomoNo-1, 'TomoImportCtfFindFile'] = 'tomograms/' + tomoName + '/' + tomoName + '_output.txt'
df_tomo.loc[tomoNo-1, 'TomoImportImodDir'] = 'tomograms/' + tomoName
df_tomo.loc[tomoNo-1, 'TomoImportFractionalDose'] = args.frac_dose
df_tomo.loc[tomoNo-1, 'TomoImportOrderList'] = orderList
df_tomo.loc[tomoNo-1, 'TomoImportCulledFile'] = 'tomograms/' + tomoName + '/' + tomoName + '_culled.mrc'
print(' -->' + str(doubletId))
# This part need to be fixed
starFile = 'star_corr/' + record[1] + '.star'
docFile = 'doc_corr/doc_total_' + record[0] + '.spi'
# Remove the comment in spider file
preprocess_bstar(starFile)
preprocess_spider_doc(docFile)
# Convert
df_relion = aa_to_relion(starFile.replace('.star', '.txt'), docFile, tomoName, tomoNo, binFactor, pixelSize, doubletId)
if df_all is None:
df_all = df_relion.copy()
else:
df_all = df_all.append(df_relion)
# Renumber
df_all['TomoParticleId'] = np.arange(len(df_all), dtype=np.int16) + 1
print("Writing " + args.ostar)
write_star_4(df_all, args.ostar, 'particles')
print("Writing tomograms_" + args.ostar)
write_star_4(df_tomo, 'tomograms_' + args.ostar, '')