def compute_full_ctf(D, Nimg, args):
freqs = np.arange(-D / 2, D / 2) / (args.Apix * D)
x0, x1 = np.meshgrid(freqs, freqs)
freqs = np.stack([x0.ravel(), x1.ravel()], axis=1)
if args.ctf_pkl: # todo: refator
params = pickle.load(open(args.ctf_pkl, 'rb'))
try:
assert len(params) == Nimg
except AssertionError:
log('Note that the input ctf.pkl contains ' + str(len(params)) +
' particles, but that you have only chosen to output the first '
+ str(Nimg) + ' particle')
params = params[:Nimg]
ctf = np.array([compute_ctf(freqs, *x[2:], args.b) for x in params])
ctf = ctf.reshape((Nimg, D, D))
df1 = np.array([x[2] for x in params])
df2 = np.array([x[3] for x in params])
dfa = np.array([x[4] for x in params])
dfp = np.array([x[8] for x in params])
df = np.stack([df1, df2, dfa, dfp], axis=1)
elif args.df_file:
raise NotImplementedError
# df = pickle.load(open(args.df_file,'rb'))
# try:
# assert len(df) == Nimg
# except AssertionError:
# log('Note that the input df.pkl contains '+ str(len(df)) + ' particles, but that you have only chosen to output the first ' + str(Nimg) + ' particle')
# df = df[:Nimg]
# ctf = np.array([compute_ctf(freqs, i, i, args.ang, args.kv, args.cs, args.wgh, args.ps, args.b) \
# for i in df])
# ctf = ctf.reshape((Nimg, D, D))
# df = np.stack([df,df], axis=1)
elif args.sample_df:
raise NotImplementedError
#
# df1 = np.random.normal(args.dfu,args.sample_df,Nimg)
# if args.no_astigmatism:
# assert args.dfv == args.dfu, "--dfu and --dfv must be the same"
# df2 = df1
# else:
# df2 = np.random.normal(args.dfv,args.sample_df,Nimg)
# ctf = np.array([compute_ctf(freqs, i, j, args.ang, args.kv, args.cs, args.wgh, args.ps, args.b) \
# for i, j in zip(df1, df2)])
# ctf = ctf.reshape((Nimg, D, D))
# df = np.stack([df1,df2], axis=1)
else:
ctf = compute_ctf(freqs, args.dfu, args.dfv, args.ang, args.kv,
args.cs, args.wgh, args.ps, args.b)
ctf = ctf.reshape((D, D))
df = np.stack([np.ones(Nimg) * args.dfu,
np.ones(Nimg) * args.dfv],
np.ones(Nimg) * args.ang,
np.ones(Nimg) * args.ps,
axis=1)
return ctf, df
def train(model,
lattice,
optim,
y,
rot,
trans=None,
ctf_params=None,
use_amp=False):
model.train()
optim.zero_grad()
B = y.size(0)
D = lattice.D
# reconstruct circle of pixels instead of whole image
mask = lattice.get_circular_mask(D // 2)
yhat = model(lattice.coords[mask] @ rot).view(B, -1)
if ctf_params is not None:
freqs = lattice.freqs2d[mask]
freqs = freqs.unsqueeze(0).expand(
B, *freqs.shape) / ctf_params[:, 0].view(B, 1, 1)
yhat *= ctf.compute_ctf(freqs, *torch.split(ctf_params[:, 1:], 1, 1))
y = y.view(B, -1)[:, mask]
if trans is not None:
y = lattice.translate_ht(y, trans.unsqueeze(1), mask).view(B, -1)
loss = F.mse_loss(yhat, y)
if use_amp:
with amp.scale_loss(loss, optim) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optim.step()
return loss.item()
def run_batch(model, lattice, y, yt, rot, tilt=None, ctf_params=None, yr=None):
use_tilt = yt is not None
use_ctf = ctf_params is not None
B = y.size(0)
D = lattice.D
if use_ctf:
freqs = lattice.freqs2d.unsqueeze(0).expand(
B, *lattice.freqs2d.shape) / ctf_params[:, 0].view(B, 1, 1)
c = ctf.compute_ctf(freqs, *torch.split(ctf_params[:, 1:], 1,
1)).view(B, D, D)
# encode
if yr is not None:
input_ = (yr, )
else:
input_ = (y, yt) if yt is not None else (y, )
if use_ctf:
input_ = (x * c.sign() for x in input_) # phase flip by the ctf
z_mu, z_logvar = _unparallelize(model).encode(*input_)
z = _unparallelize(model).reparameterize(z_mu, z_logvar)
# decode
mask = lattice.get_circular_mask(D // 2) # restrict to circular mask
y_recon = model(lattice.coords[mask] / lattice.extent / 2 @ rot,
z).view(B, -1)
if use_ctf: y_recon *= c.view(B, -1)[:, mask]
# decode the tilt series
if use_tilt:
y_recon_tilt = model(
lattice.coords[mask] / lattice.extent / 2 @ tilt @ rot, z)
if use_ctf: y_recon_tilt *= c.view(B, -1)[:, mask]
else:
y_recon_tilt = None
return z_mu, z_logvar, z, y_recon, y_recon_tilt, mask
def eval_z(model, lattice, data, batch_size, device, trans=None, use_tilt=False, ctf_params=None, use_real=False):
assert not model.training
z_mu_all = []
z_logvar_all = []
data_generator = DataLoader(data, batch_size=batch_size, shuffle=False)
for minibatch in data_generator:
ind = minibatch[-1]
y = minibatch[0].to(device)
yt = minibatch[1].to(device) if use_tilt else None
B = len(ind)
D = lattice.D
if ctf_params is not None:
freqs = lattice.freqs2d.unsqueeze(0).expand(B,*lattice.freqs2d.shape)/ctf_params[ind,0].view(B,1,1)
c = ctf.compute_ctf(freqs, *torch.split(ctf_params[ind,1:], 1, 1)).view(B,D,D)
if trans is not None:
y = lattice.translate_ht(y.view(B,-1), trans[ind].unsqueeze(1)).view(B,D,D)
if yt is not None: yt = lattice.translate_ht(yt.view(B,-1), trans[ind].unsqueeze(1)).view(B,D,D)
if use_real:
input_ = (torch.from_numpy(data.particles_real[ind]).to(device),)
else:
input_ = (y,yt) if yt is not None else (y,)
if ctf_params is not None:
assert not use_real, "Not implemented"
input_ = (x*c.sign() for x in input_) # phase flip by the ctf
z_mu, z_logvar = _unparallelize(model).encode(*input_)
z_mu_all.append(z_mu.detach().cpu().numpy())
z_logvar_all.append(z_logvar.detach().cpu().numpy())
z_mu_all = np.vstack(z_mu_all)
z_logvar_all = np.vstack(z_logvar_all)
return z_mu_all, z_logvar_all
def run_model(y):
'''Helper function'''
# reconstruct circle of pixels instead of whole image
mask = lattice.get_circular_mask(D // 2)
yhat = model(lattice.coords[mask] @ rot).view(B, -1)
if ctf_params is not None:
freqs = lattice.freqs2d[mask]
freqs = freqs.unsqueeze(0).expand(
B, *freqs.shape) / ctf_params[:, 0].view(B, 1, 1)
yhat *= ctf.compute_ctf(freqs,
*torch.split(ctf_params[:, 1:], 1, 1))
y = y.view(B, -1)[:, mask]
if trans is not None:
y = lattice.translate_ht(y, trans.unsqueeze(1), mask).view(B, -1)
return F.mse_loss(yhat, y)
def main(args):
assert args.o.endswith('.mrc')
t1 = time.time()
log(args)
if not os.path.exists(os.path.dirname(args.o)):
os.makedirs(os.path.dirname(args.o))
## set the device
use_cuda = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')
log('Use cuda {}'.format(use_cuda))
if use_cuda:
torch.set_default_tensor_type(torch.cuda.FloatTensor)
else:
log('WARNING: No GPUs detected')
# load the particles
if args.tilt is None:
data = dataset.LazyMRCData(args.particles,
norm=(0, 1),
invert_data=args.invert_data,
datadir=args.datadir)
tilt = None
else:
data = dataset.TiltMRCData(args.particles,
args.tilt,
norm=(0, 1),
invert_data=args.invert_data,
datadir=args.datadir)
tilt = torch.tensor(utils.xrot(args.tilt_deg).astype(np.float32))
D = data.D
Nimg = data.N
lattice = Lattice(D, extent=D // 2)
posetracker = PoseTracker.load(args.poses, Nimg, D, None, None)
if args.ctf is not None:
log('Loading ctf params from {}'.format(args.ctf))
ctf_params = ctf.load_ctf_for_training(D - 1, args.ctf)
ctf_params = torch.tensor(ctf_params)
else:
ctf_params = None
Apix = ctf_params[0, 0] if ctf_params is not None else 1
V = torch.zeros((D, D, D))
counts = torch.zeros((D, D, D))
mask = lattice.get_circular_mask(D // 2)
if args.ind:
iterator = pickle.load(open(args.ind, 'rb'))
elif args.first:
args.first = min(args.first, Nimg)
iterator = range(args.first)
else:
iterator = range(Nimg)
for ii in iterator:
if ii % 100 == 0: log('image {}'.format(ii))
r, t = posetracker.get_pose(ii)
ff = data.get(ii)
if tilt is not None:
ff, ff_tilt = ff # EW
ff = torch.tensor(ff)
ff = ff.view(-1)[mask]
if ctf_params is not None:
freqs = lattice.freqs2d / ctf_params[ii, 0]
c = ctf.compute_ctf(freqs, *ctf_params[ii, 1:]).view(-1)[mask]
ff *= c.sign()
if t is not None:
ff = lattice.translate_ht(ff.view(1, -1), t.view(1, 1, 2),
mask).view(-1)
ff_coord = lattice.coords[mask] @ r
add_slice(V, counts, ff_coord, ff, D)
# tilt series
if args.tilt is not None:
ff_tilt = torch.tensor(ff_tilt)
ff_tilt = ff_tilt.view(-1)[mask]
if ctf_params is not None:
ff_tilt *= c.sign()
if t is not None:
ff_tilt = lattice.translate_ht(ff_tilt.view(1, -1),
t.view(1, 1, 2), mask).view(-1)
ff_coord = lattice.coords[mask] @ tilt @ r
add_slice(V, counts, ff_coord, ff_tilt, D)
td = time.time() - t1
log('Backprojected {} images in {}s ({}s per image)'.format(
len(iterator), td, td / Nimg))
counts[counts == 0] = 1
V /= counts
V = fft.ihtn_center(V[0:-1, 0:-1, 0:-1].cpu().numpy())
mrc.write(args.o, V.astype('float32'), Apix=Apix)
