def ssnr_sequential___individual_data_collect(self, r, op):
if (self is None):
get_mrc_func = IV.get_mrc
else:
get_mrc_func = self.cache.get_mrc
v = get_mrc_func(r['subtomogram'])
if 'angle' in r:
v = GR.rotate_pad_mean(v, angle=N.array(r['angle'], dtype=N.float), loc_r=N.array(r['loc'], dtype=N.float))
if (op is not None) and ('segmentation_tg' in op) and ('template' in r) and ('segmentation' in r['template']):
phi = IV.read_mrc_vol(r['template']['segmentation'])
phi_m = (phi > 0.5)
del phi
(ang_inv, loc_inv) = AAL.reverse_transform_ang_loc(r['angle'], r['loc'])
phi_mr = GR.rotate(phi_m, angle=ang_inv, loc_r=loc_inv, default_val=0)
del phi_m
del ang_inv, loc_inv
import aitom.tomominer.pursuit.multi.util as PMU
v_s = PMU.template_guided_segmentation(v=v, m=phi_mr, op=op['segmentation_tg'])
del phi_mr
if v_s is not None:
v_f = N.isfinite(v_s)
if v_f.sum() > 0:
v_s[N.logical_not(v_f)] = v_s[v_f].mean()
v = v_s
del v_s
v = NF.fftshift(NF.fftn(v))
m = get_mrc_func(r['mask'])
if 'angle' in r:
m = GR.rotate_mask(m, angle=N.array(r['angle'], dtype=N.float))
v[(m < op['mask_cutoff'])] = 0.0
return {'v': v, 'm': m, }
def random_rotate(v):
angle = GAL.random_rotation_angle_zyz()
vr = GR.rotate(v, angle=angle, default_val=0.0) # loc_r is none
# print('angle:', angle)
# print('loc_max:', loc_max)
# print('loc_r:',type(loc_r),loc_r)
return vr, angle
def normalize(record, op):
if os.path.isfile(record['pose']['subtomogram']):
return {
'record': record,
}
ls = level_set(record=record, op=op['segmentation'])
if ls is None:
return
phi = N.zeros(ls['phi'].shape)
phi[(ls['phi'] > 0)] = ls['phi'][(ls['phi'] > 0)]
c = PNU.center_mass(phi)
mid_co = (N.array(phi.shape) / 2)
if N.sqrt(N.square(
(c - mid_co)).sum()) > (N.min(phi.shape) *
op['center_mass_max_displacement_proportion']):
return
rm = PNU.pca(v=phi, c=c)['v']
record['pose']['c'] = c.tolist()
record['pose']['rm'] = rm.tolist()
phi_pn = GR.rotate(phi, rm=rm, c1=c, default_val=0)
v_org_pn = GR.rotate_pad_mean(ls['v_org'], rm=rm, c1=c)
return {
'ls': ls,
'phi': phi,
'phi_pn': phi_pn,
'v_org_pn': v_org_pn,
'record': record,
}
def forward(x0, sc, c, phi, theta, psi, orientation, sigma_c, sigma_l, sigma_d, t1, bm):
"""
The forward tracing algorithm (recursive)
@param:
x0: index of the starting voxel
sc: the forward (along y) search cone as a numpy array
c: a numpy array storing the correlation coefficients from template matching
phi: a numpy array storing the first ZYZ rotation angle of the template from template matching
theta: a numpy array storing the second ZYZ rotation angle of the template from template matching
psi: a numpy array storing the thrid ZYZ rotation angle of the template from template matching
orientaion: a numpy array storing the template orientations in vector form from template matching
sigma_c: user-defined parameter for smoothness (smoother lines with smaller sigma_c)
sigma_l: user-defined parameter for linearity
sigma_d: user-defined parameter for distance
t1: threshold for similarity
bm: a numpy array storing the binary mask that is to be updated
@returns:
bm when search_result.max() <= t1
"""
# get ZYZ template rotation angles
angles = (phi[x0], theta[x0], psi[x0])
# rotate search cone according to template rotation angles
sc_curr = GR.rotate(sc, angle=angles, default_val=0.0)
# get x0 location on the search cone for index calculation later
x0_loc = convert(np.where(sc_curr == sc_curr.min()))
# convert the format of sc: >0-->True
sc_mask = sc_curr > 0
# create a np array the size of the sc to store search results
# highest similarity in search_result means this voxel is part of the filament
search_result = np.full(sc_mask.shape, -1.0)
# for every voxel in the search cone, calculate similarity and update search_result
for idx, val in np.ndenumerate(sc_mask):
if val is True:
# x = idx - x0_loc + x0 #convert to the index of x
x = (x0[0] - x0_loc[0] + idx[0], x0[1] - x0_loc[1] + idx[1], x0[2] - x0_loc[2] + idx[2])
# z, y, x boundaries
if (x[0] <= c.shape[0] - 1) and (x[1] <= c.shape[1] - 1) and (x[2] <= c.shape[2] - 1):
# calculate similarity between x and x0
s = similarity(x0, x, c[x], sigma_c, sigma_l, sigma_d, orientation)
if s > t1:
# store voxel location when similarity is above threshold t1
search_result[idx] = s
# continue forward()
if search_result.max() > t1:
# highest similarity in the search cone
x = convert(np.where(search_result == search_result.max()))
# convert coordinate
x = (x0[0] - x0_loc[0] + x[0], x0[1] - x0_loc[1] + x[1], x0[2] - x0_loc[2] + x[2])
# update binary mask
bm[x] = True
# recursively move
return forward(x, sc, c, phi, theta, psi, orientation, sigma_c, sigma_l, sigma_d, t1, bm)
# forward
# stop searching when all similarity values within sc are below threshold t1
else:
return bm
def forward(x0, sc):
'''
The forward tracing algorithm
@param x0: index of the starting voxel
@param sc: the forward (along y) search cone as a numpy array
@returns when search_result.max() <= t1, where t1 is a similarity threshold specified by the user
updates global binary mask variable bm
'''
angles = (phi[x0], theta[x0], psi[x0]) #get ZYZ template rotation angles
sc_curr = GR.rotate(
sc, angle=angles, default_val=0.0
) #rotate search cone according to template rotation angles
x0_loc = convert(np.where(sc_curr == sc_curr.min(
))) #get x0 location on the search cone for index calculation later
sc_mask = sc_curr > 0 #convert the format of sc: >0-->True
#create a np array the size of the sc to store search results
#highest similarity in search_result means this voxel is part of the filament
search_result = np.full(sc_mask.shape, -1.0)
#for every voxel in the search cone, calculate similarity and update search_result
for idx, val in np.ndenumerate(sc_mask):
if val == True:
x = (x0[0] - x0_loc[0] + idx[0], x0[1] - x0_loc[1] + idx[1],
x0[2] - x0_loc[2] + idx[2]
) #x = idx - x0_loc + x0 #convert to the index of x
if (x[0] <= c.shape[0] - 1) and (x[1] <= c.shape[1] - 1) and (
x[2] <= c.shape[2] - 1): #z, y, x boundaries
s = similarity(x0, x, c[x], sigma_c, sigma_l,
sigma_d) #calculate similarity between x and x0
if s > t1:
search_result[
idx] = s #store voxel location when similarity is above threshold t1
if search_result.max() > t1: #continue forward()
x = convert(np.where(search_result == search_result.max())
) #highest similarity in the search cone
x = (x0[0] - x0_loc[0] + x[0], x0[1] - x0_loc[1] + x[1],
x0[2] - x0_loc[2] + x[2]) #convert coordinate
bm[x] = True #update binary mask
return forward(x, sc) #recursively move forward
else: #stop searching when all similarity values within sc are below threshold t1
return
def angle_rotate(v, angle):
vr = GR.rotate(v, angle=angle, default_val=0.0)
return vr
'voltage': 300,
'Cs': 2.0,
'sigma': 0.4
}
}
# generate a density map v that contains a toy structure
v = MU.generate_toy_model(dim_siz=64) # generate a pseudo density map
print(v.shape)
# randomly rotate and translate v
loc_proportion = 0.1
loc_max = N.array(v.shape, dtype=float) * loc_proportion
angle = GAL.random_rotation_angle_zyz()
loc_r = (N.random.random(3) - 0.5) * loc_max
vr = GR.rotate(v, angle=angle, loc_r=loc_r, default_val=0.0)
# generate simulated subtomogram vb from v
vb = TSRSC.do_reconstruction(vr, op, verbose=True)
print('vb', 'mean', vb.mean(), 'std', vb.std(), 'var', vb.var())
# save v and vb as 3D grey scale images
TIF.put_mrc(vb, '/tmp/vb.mrc', overwrite=True)
TIF.put_mrc(v, '/tmp/v.mrc', overwrite=True)
# save images of the slices of the corresponding 3D iamges for visual inspection
import aitom.image.io as IIO
import aitom.tomominer.image.vol.util as TIVU
IIO.save_png(TIVU.cub_img(vb)['im'], "/tmp/vb.png")
IIO.save_png(TIVU.cub_img(v)['im'], "/tmp/v.png")
def trace(t1, t2, sigma_c, sigma_l, sigma_d, sc_path, out_path, c_path, phi_path, theta_path, psi_path):
"""
The main function for filament tracing
@param:
t1: threshold for similarity, defualt=0.000001
t2: threshold for correlation coefficient, default= 0.0007
sigma_c: measure for smootheness, default=1.0
sigma_l: measure for linearity, default=1.0
sigma_d: measure for distance, default=1.0
sc_path: file path of search cone, see search cone linked at the top of the script
out_path: output file path for the tracing result (a binary mask)
c_path: file path of the correlation coefficients (e.g. c.npy from 007_template_matching.py)
phi_path: file path of the first ZYZ rotation angle(e.g. phi.npy from 007_template_matching.py)
theta_path: file path of the second ZYZ rotation angle(e.g. theta.npy from 007_template_matching.py)
psi_path: file path of the thrid ZYZ rotation angle(e.g. psi.npy from 007_template_matching.py)
@return:
bm, a binary mask storing the search result
Note that if out_path is not None, trace() would store the search result at the given location
"""
# forward() and backward() might reach recursion limit
sys.setrecursionlimit(10000)
assert (sc_path is not None) and (c_path is not None), "File path error"
assert (phi_path is not None) and (theta_path is not None) and (psi_path is not None), "File path error"
if t1 is None:
t1 = 0.000001
if t2 is None:
t2 = 0.0007
if sigma_c is None:
sigma_c = 1.0
if sigma_l is None:
sigma_l = 1.0
if sigma_d is None:
sigma_d = 1.0
# load template matching results (see 007_template_matching_tutorial.py)
print("Loading template matching results...")
c = np.load(c_path) # (10, 400, 398)
phi = np.load(phi_path)
theta = np.load(theta_path)
psi = np.load(psi_path)
print("Map shape: ", c.shape)
# preprocess cross correlation matrix: negative values => 0.0
print("Preproces correlation coefficients...")
c[c < 0.0] = 0.0
print("Maximum corr = ", c.max())
print("Preprocessing orientations...")
# orientation = np.load('./orietnation.npy', allow_pickle=True)#can store orientation.npy for faster testing
# empty np array to store template orientations
orientation = np.empty(c.shape, dtype=object)
# assume template is parallel to the y axis
v = np.array([0, 1, 0])
# for all voxels
for idx, x in np.ndenumerate(c):
# get ZYZ angle for template orientation
angle = (phi[idx], theta[idx], psi[idx])
# convert to rotation matrix
rm = AA.rotation_matrix_zyz(angle)
# template orientation in [x, y, z]
orientation[idx] = rm.dot(v)
# binary mask to store tracing results
bm = np.full(c.shape, False, dtype=bool)
# read in search cone: SC is along the y axis, i.e. [0,1,0]
print("Preparing search cone...")
mrc = mrcfile.open(sc_path, mode='r+', permissive=True)
# rotated already, along the y axis for now #forward search cone
sc = mrc.data
print('search cone size', sc.shape) # (10, 11, 11)
# mark where x0 is with a negative value
sc[(5, 0, 5)] = -100.0
# 180 degree rotation of the search cone
rm = np.array([[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]])
# about the XYZ axes
# rotate to get the backward search cone
sc_back = GR.rotate(sc, rm=rm, default_val=0.0)
# tracing: go through corr for all c.max()>=t2
print("Start tracing actin filaments...")
while c.max() >= t2:
# tracing starts with the highest corr
corr = c.max()
# get the index of c.max #e.g.(array([5]), array([177]), array([182]))
x0 = np.where(c == corr)
assert (x0[0].size == 1), "multiple maximum"
# convert index to a tuple
x0 = (x0[0][0], x0[1][0], x0[2][0])
print("Now at ", x0, " with corr=", corr)
bm[x0] = True
# recursively search forward
bm = forward(x0, sc, c, phi, theta, psi, orientation, sigma_c, sigma_l, sigma_d, t1, bm)
bm = backward(x0, sc_back, c, phi, theta, psi, orientation, sigma_c, sigma_l, sigma_d, t1, bm)
# set c.max() to -1.0 and go to the next c.max()
c[x0] = -1.0
# output bm
if out_path is not None:
print("Generating output to ", out_path)
bm = bm.astype(np.int16)
mrc = mrcfile.new(out_path, overwrite=True)
mrc.set_data(bm)
mrc.close()
print("TRACING DONE")
return bm
def scan(op):
'''
This function contains:
template and map preprocessing
template rotations using ZYZ Euler angles
cross-correlation calculation
output: c.npy, phi.npy, theta.npy, psi.npy
'''
if not os.path.isdir(op['out_dir']): os.makedirs(op['out_dir'])
re = {} #create output file paths
re['c'] = os.path.join(op['out_dir'], '%s-c.npy'%(op['id'], ))
re['phi'] = os.path.join(op['out_dir'], '%s-phi.npy'%(op['id'], ))
re['theta'] = os.path.join(op['out_dir'], '%s-theta.npy'%(op['id'], ))
re['psi'] = os.path.join(op['out_dir'], '%s-psi.npy'%(op['id'], ))
#if the output files already exits, template matching would not run
if os.path.isfile(re['c']) and os.path.isfile(re['phi']) and os.path.isfile(re['theta']) and os.path.isfile(re['psi']): return re
#load template as t and preprocess
mrc = mrcfile.open(op['template'], mode='r+',permissive=True)
t = mrc.data
print('template size',t.shape)
tm = N.isfinite(t) # real space mask
t_mean = t[tm].mean()
t[N.logical_not(tm)] = t_mean
tm = tm.astype(N.float)
#load map as v
mrc = mrcfile.open(op['map'], mode='r+',permissive=True)
v = mrc.data
v = v.astype(N.float)
print ('map size', v.shape) ; sys.stdout.flush()
diff_time_v = []
cur_time = time.time()
c_max = None
phi_max = None
theta_max = None
psi_max = None
#template rotations
for i, (phi, theta, psi) in enumerate(op['angles']):
#print(i, (phi, theta, psi))
#template rotation with affine transformation
tr = GR.rotate(t, angle=(phi, theta, psi), default_val=t_mean)
tr = tr.astype(N.float)
if op['mode'] == 'convolve': #convolution
#c = FC.convolve(v=v, t=tr)
c = SNF.convolve(input=v, weights=tr, mode='reflect')
elif op['mode'] == 'normalized-cor': #correlation
tmr = GR.rotate(tm, angle=(phi, theta, psi), default_val=0.0)
tr[tmr < 0.5] = float('NaN')
c = FC.pearson_correlation_simple(v=v, t=tr) #calculate Pearson cross correlation
else:
raise Exception('mode')
#record the maximum cross-correlation coefficient and the orientation
#c_max - a 3d numpy array of maximum cross-correlation coefficients
#phi_max, theta_max, psi_max - 3d numpy arrays that store the maximum phi, theta, psi values (template orientation in ZYZ convention)
if c_max is None:
c_max = c
phi_max = N.zeros(c.shape) + phi
theta_max = N.zeros(c.shape) + theta
psi_max = N.zeros(c.shape) + psi
else:
ind = (c > c_max)
c_max[ind] = c[ind]
phi_max[ind] = phi
theta_max[ind] = theta
psi_max[ind] = psi
#run time
diff_time = time.time()-cur_time
diff_time_v.append(diff_time)
remain_time = N.array(diff_time_v).mean() * (len(op['angles']) - i - 1)
print ('angle', i+1, 'of', len(op['angles']), '; time used', diff_time, '; time remain', remain_time) ; sys.stdout.flush()
cur_time = time.time()
#output
if not os.path.isfile(re['c']):
N.save(re['c'], c_max)
N.save(re['phi'], phi_max)
N.save(re['theta'], theta_max)
N.save(re['psi'], psi_max)
return re
if os.path.isfile(re['c']) and os.path.isfile(re['phi']) and os.path.isfile(re['theta']) and os.path.isfile(re['psi']): return re
def scan(op):
if not os.path.isdir(op['out_dir']):
os.makedirs(op['out_dir'])
re = {'c': os.path.join(op['out_dir'], '%s-c.npy' % (op['id'],)),
'phi': os.path.join(op['out_dir'], '%s-phi.npy' % (op['id'],)),
'theta': os.path.join(op['out_dir'], '%s-theta.npy' % (op['id'],)),
'psi': os.path.join(op['out_dir'], '%s-psi.npy' % (op['id'],))}
if os.path.isfile(re['c']) and os.path.isfile(re['phi']) and os.path.isfile(re['theta']) and os.path.isfile(
re['psi']):
return re
t = N.load(op['template'])
# real space mask
tm = N.isfinite(t)
t_mean = t[tm].mean()
t[N.logical_not(tm)] = t_mean
tm = tm.astype(N.float)
v = N.load(op['map'])
print('map size', v.shape)
sys.stdout.flush()
diff_time_v = []
cur_time = time.time()
c_max = None
phi_max = None
theta_max = None
psi_max = None
for i, (phi, theta, psi) in enumerate(op['angles']):
tr = GR.rotate(t, angle=(phi, theta, psi), default_val=t_mean)
if op['mode'] == 'convolve':
# c = FC.convolve(v=v, t=tr)
c = SNF.convolve(input=v, weights=tr, mode='reflect')
elif op['mode'] == 'normalized-cor':
tmr = GR.rotate(tm, angle=(phi, theta, psi), default_val=0.0)
tr[tmr < 0.5] = float('NaN')
c = FNCC.cor(v=v, t=tr)
else:
raise Exception('mode')
if c_max is None:
c_max = c
phi_max = N.zeros(c.shape) + phi
theta_max = N.zeros(c.shape) + theta
psi_max = N.zeros(c.shape) + psi
else:
ind = (c > c_max)
c_max[ind] = c[ind]
phi_max[ind] = phi
theta_max[ind] = theta
psi_max[ind] = psi
diff_time = time.time() - cur_time
diff_time_v.append(diff_time)
remain_time = N.array(diff_time_v).mean() * (len(op['angles']) - i - 1)
print('angle', i, 'of', len(op['angles']), '; time used', diff_time, '; time remain', remain_time)
sys.stdout.flush()
cur_time = time.time()
if not os.path.isfile(re['c']):
N.save(re['c'], c_max)
N.save(re['phi'], phi_max)
N.save(re['theta'], theta_max)
N.save(re['psi'], psi_max)
return re
# generate simulated images
import aitom.model.util as TMU
v = TMU.generate_toy_model(dim_siz=32)
import aitom.geometry.ang_loc as TGAL
import aitom.geometry.rotate as TGR
import random
import numpy as N
# randomly rotate and translate v
loc_proportion = 0.1
loc_max = N.array(v.shape, dtype=float) * loc_proportion
angle = TGAL.random_rotation_angle_zyz()
loc_r = (N.random.random(3) - 0.5) * loc_max
vr = TGR.rotate(v, angle=angle, loc_r=loc_r, default_val=0.0)
#--------------------------------
# align vr against v
import aitom.align.util as TAU
al = TAU.align_vols_no_mask(v, vr)
print('rigid transform of alignment', al)
vr_i = TGR.rotate(
vr, angle=al['angle'], loc_r=al['loc'], default_val=0.0
) # rotate vr according to the alignment, expected to produce an image similiar to v
# save images of the slices of the corresponding 3D iamges for visual inspection
import aitom.image.io as TIIO
import aitom.image.vol.util as TIVU
# generate simulated images
import aitom.model.util as MU
v = MU.generate_toy_model(dim_siz=32)
import aitom.geometry.ang_loc as GAL
import aitom.geometry.rotate as GR
import numpy as N
# randomly rotate and translate v
loc_proportion = 0.1
loc_max = N.array(v.shape, dtype=float) * loc_proportion
angle = GAL.random_rotation_angle_zyz()
loc_r = (N.random.random(3) - 0.5) * loc_max
vr = GR.rotate(v, angle=angle, loc_r=loc_r, default_val=0.0)
# align vr against v
import aitom.align.fast.util as AFU
al = AFU.align_vols_no_mask(v, vr)
print('rigid transform of alignment', al)
# rotate vr according to the alignment, expected to produce an image similiar to v
vr_i = GR.rotate(vr, angle=al['angle'], loc_r=al['loc'], default_val=0.0)
# save images of the slices of the corresponding 3D iamges for visual inspection
import aitom.image.io as IIO
import aitom.image.vol.util as IVU
IIO.save_png(IVU.cub_img(v)['im'], "v.png")
