def EM(img_data, K, iteration, path, snapshot_interval=5, reg=False, use_voronoi=True):
np.seterr(all='ignore')
X = get_image_db(img_data['db_path'])
dj = img_data['dj']
N = len(dj)
n_x, n_y, n_z = X[dj[0]['v']].shape
theta = dict()
theta['N'] = N
theta['J'] = n_x * n_y * n_z
theta['n'] = n_x
theta['K'] = K
theta['xi'] = theta['n'] # Proportional to the radius of the image
theta['A'] = np.zeros([K, n_x, n_y, n_z], dtype=np.complex128) # We need to initialize this later
theta['alpha'] = np.ones([K], dtype=np.float_) / K
theta['trans_list'] = None
theta['predictions'] = np.zeros([N])
# Print relavent information
print("Running model based alignment: N=%d, K=%d, dimensions=(%d,%d,%d)" % (N, K, n_x, n_y, n_z))
if reg:
print("With regularization")
else:
print("Without regularization")
if use_voronoi:
print("With voronoi weights")
else:
print("Without voronoi weights")
# Regularization
reg_step = (float(N) / K ** 2) / 2
theta['theta_reg'] = 5 * reg_step if reg else 0
# Sample K random data points from the set to initialize A
indices = np.random.permutation(N)
num_models = [0 for _ in range(K)]
k = 0
for i in range(N):
theta['A'][k] += X[dj[indices[i]]['v']] * X[dj[indices[i]]['m']]
num_models[k] += 1
k = (k+1) % K
for k in range(K):
theta['A'][k] /= num_models[k]
# Get a random A_k and a random X_i
# And calculate sum_j to get sigma_sq
k = np.random.randint(K)
i = np.random.randint(N)
sum_j = np.sum(np.square(np.absolute(theta['A'][k] - X[dj[i]['v']]) * X[dj[i]['m']]))
theta['sigma_sq'] = sum_j / theta['J']
print("Sigma_sq initialized to %d" % theta['sigma_sq'])
checkpoint_dir = os.path.join(path, 'checkpoints')
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
interval = snapshot_interval
for i in range(iteration):
checkpoint_file = os.path.join(checkpoint_dir, '%08d.pickle'%(i))
if os.path.exists(checkpoint_file):
checkpoint_data = TIF.pickle_load(checkpoint_file)
theta = checkpoint_data['theta']
continue
if (i % interval == 0):
output_images(theta, i, path=path)
print("Running iteration %d" % (i+1))
# Update alpha before updating A
compute_trans_list(theta=theta, img_data=img_data, use_voronoi=use_voronoi)
alpha = update_alpha(img_data=img_data, theta=theta, use_voronoi=use_voronoi)
print("Alpha updated! Alpha = ", end=' ')
print(alpha.tolist())
sigma_sq = update_sigma(img_data=img_data, theta=theta, reg=reg, use_voronoi=use_voronoi)
print("Sigma updated! Sigma^2 = ", end=' ')
print(sigma_sq)
xi = update_xi(img_data=img_data, theta=theta, use_voronoi=use_voronoi)
print("Xi updated! Xi = ", end=' ')
print(xi)
A = update_a(img_data = img_data, theta=theta, alpha=alpha, reg=reg, use_voronoi=use_voronoi)
print("A updated! Average intensity of A = ", end=' ')
print(np.average(A, (1, 2, 3)))
theta['alpha'] = alpha
theta['sigma_sq'] = sigma_sq
theta['xi'] = xi
theta['A'] = A
# Since we changed the models A, the list of optimal transforms
# needs to be re-calculated
theta['trans_list'] = None
theta['pred'] = None
# Decrease the regularization coefficient
if reg and theta['theta_reg'] > 0:
theta['theta_reg'] -= reg_step
theta['theta_reg'] = max(0, theta['theta_reg'])
try:
assert not os.path.exists(checkpoint_file)
except:
raise Exception("Checkpoint file already exists!")
TIF.pickle_dump({'theta':theta}, checkpoint_file)
print_prediction_results(theta, img_data)
output_images(theta, iteration, path=path)
print("Prediction from model: ", end=' ')
print(theta['predictions'])
return theta
/shared/opt/local/img/em/et/util/situs/Situs_2.7.2/bin/pdb2vol
/opt/local/img/em/et/util/situs/Situs_2.7.2/bin/pdb2vol
spacing_s: is the voxel spacing in Anstron
resolution_s: the image resolution in Anstron
pdb_dir: the directory that contains pdb files. Could be your own directory or /shared/shared/data/pdb or /shared/data/pdb
out_file: the output file that contains converted density maps
'''
op = {
'situs_pdb2vol_program':
'/shared/opt/local/img/em/et/util/situs/Situs_2.7.2/bin/pdb2vol',
'spacing_s': [10.0],
'resolution_s': [10.0],
'pdb_dir': '/shared/data/pdb',
'out_file': 'situs_maps.pickle'
}
# convert to density maps, save to situs_maps.pickle
import aitom.tomominer.structure.pdb.situs_pdb2vol__batch as TSPS
TSPS.batch_processing(op)
'''
The density maps in situs_maps.pickle have different sizes for different structures,
use resize_center_batch_dict() to change them into same size
'''
import aitom.tomominer.io.file as TIF
ms = TIF.pickle_load('situs_maps.pickle')
ms = {_: ms[_][10.0][10.0]['map'] for _ in ms}
import aitom.tomominer.image.vol.util as TIVU
ms = TIVU.resize_center_batch_dict(vs=ms, cval=0.0)