def read_model(i, path, wedge_ang1, wedge_ang2, wedge_dir):
import aitom.image.vol.wedge.util as W
result = {'i':i}
f = mrc.open(path)
size = f.data.shape[0]
result['observed'] = fourier_transform(f.data)
mask = W.wedge_mask(size=f.data.shape, ang1=wedge_ang1, ang2=wedge_ang2, direction=wedge_dir)
result['mask'] = mask
return result
def single_average(subtom):
print('subtom_type=', type(subtom))
assert len(subtom) == 100
print('subtom[0]_type=', type(subtom[0]))
average = True
test_dir = './tmp/cls-test/' + str(uuid.uuid4()) # test dir
if os.path.exists(test_dir):
shutil.rmtree(test_dir)
os.makedirs(test_dir)
dj_file = os.path.join(test_dir, 'data.pickle')
img_db_file = os.path.join(test_dir, 'image.db')
v_num = 100 # the number of each class
v_dim_siz = 32
wedge_angle = 30
mask_id = str(uuid.uuid4())
dj = []
class_num = 1
for model_id in range(class_num):
for v_i in range(v_num):
ang_t = [_ for _ in N.random.random(3) * (N.pi * 2)]
# loc_t = TGA.random_translation(size=[v_dim_siz]*3, proportion=0.2)
loc_t = [0.0, 0.0, 0.0]
v_id = str(uuid.uuid4())
dj.append({
'subtomogram': v_id,
'mask': mask_id,
'angle': ang_t,
'loc': loc_t,
'model_id': model_id
})
AIF.pickle_dump(dj, dj_file)
sim_op = {
'model': {
'missing_wedge_angle': wedge_angle,
'titlt_angle_step': 1,
'SNR': 1000,
'band_pass_filter': False,
'use_proj_mask': False
},
'ctf': {
'pix_size': 1.0,
'Dz': -5.0,
'voltage': 300,
'Cs': 2.0,
'sigma': 0.4
}
}
img_db = TIDL.LSM(img_db_file)
index = 0
for d in dj:
img_db[d['subtomogram']] = subtom[index].astype(N.float)
# print(img_db[d['subtomogram']].shape)
index = index + 1
import aitom.image.vol.wedge.util as TIVWU
img_db[mask_id] = TIVWU.wedge_mask(size=[v_dim_siz] * 3, ang1=wedge_angle)
print('file generation complete')
out_dir = os.path.join(test_dir, 'out')
if os.path.exists(out_dir): shutil.rmtree(out_dir)
os.makedirs(out_dir)
from aitom.classify.align.simple_iterative.classify import randomize_orientation
from aitom.classify.align.simple_iterative.classify import export_avgs
if average:
import aitom.average.align.simple_iterative.average as avg
op = {}
op['option'] = {'pass_num': 20} # the number of iterations
op['data_checkpoint'] = os.path.join(out_dir, 'djs.pickle')
op['average'] = {}
op['average']['mask_count_threshold'] = 2
op['average']['checkpoint'] = os.path.join(out_dir, 'avgs.pickle')
dj = AIF.pickle_load(os.path.join(test_dir, 'data.pickle'))
img_db = TIDL.LSM(os.path.join(test_dir, 'image.db'), readonly=True)
randomize_orientation(dj)
avg.average(dj_init=dj, img_db=img_db, op=op)
export_avgs(AIF.pickle_load(os.path.join(out_dir, 'avgs.pickle')),
out_dir=os.path.join(out_dir, 'avgs-export'))
print('averaging done')
# visualization
# test_dir = './tmp/cls-test/'+str(uuid.uuid4()) # test dir
avgs = pickle_load(os.path.join(test_dir, 'out/avgs.pickle'))
out_dir = os.path.join(test_dir, 'image')
if os.path.exists(out_dir): shutil.rmtree(out_dir)
os.makedirs(out_dir)
for i in avgs.keys():
v = avgs[i]['v']
file_name = str(avgs[i]['pass_i']) + '_' + str(i) + '.png'
save_png(cub_img(v)['im'], os.path.join(out_dir, file_name))
print('images saved in', out_dir)
'Dz': -5.0,
'voltage': 300,
'Cs': 2.0,
'sigma': 0.4
}
}
img_db = TIDL.LSM(img_db_file)
index = 0
for d in dj:
img_db[d['subtomogram']] = subtom[index].astype(N.float)
# print(img_db[d['subtomogram']].shape)
index = index + 1
import aitom.image.vol.wedge.util as TIVWU
img_db[mask_id] = TIVWU.wedge_mask(size=[v_dim_siz] * 3, ang1=wedge_angle)
print('file generation complete')
out_dir = os.path.join(test_dir, 'out')
if os.path.exists(out_dir): shutil.rmtree(out_dir)
os.makedirs(out_dir)
from aitom.classify.align.simple_iterative.classify import randomize_orientation
from aitom.classify.align.simple_iterative.classify import export_avgs
if classify: # classification and averaging
import aitom.classify.align.simple_iterative.classify as clas
class_num = 2
op = {}
op['option'] = {'pass_num': 20} # the number of iterations
op['data_checkpoint'] = os.path.join(out_dir, 'djs.pickle')
op['dim_reduction'] = {}
op['dim_reduction']['pca'] = {
from aitom.image.vol.wedge.util import wedge_mask
def fourier_transform(v):
return fftshift(fftn(v))
image_db = LSM('data/aitom_demo_subtomograms.db') # Build a database file
with open('data/aitom_demo_subtomograms.pickle', 'rb') as f:
# Load demo subtomograms pickle file
subtomograms = pickle.load(f, encoding='iso-8859-1')
dj = [] # Key file for the database
m = wedge_mask([32, 32, 32], ang1=30, sphere_mask=True,
verbose=False) # Subtomogram missing wedge masks
for s in subtomograms['5T2C_data']:
# Save fourier transformed subtomogram data
v = fourier_transform(s)
s1 = str(uuid.uuid4())
s2 = str(uuid.uuid4())
# Save a subtomogram to the database according to a key
image_db[s1] = v
# Save a subtomogram mask to the database according to a key
image_db[s2] = m
dj.append({'v': s1, 'm': s2, 'id': '5T2C'})
for s in subtomograms['1KP8_data']:
v = fourier_transform(s)
