def main(args):
mkbasedir(args.o)
warnexists(args.o)
assert (args.o.endswith('.mrcs') or args.o.endswith('mrc')), "Must specify output in .mrc(s) file format"
old = dataset.load_particles(args.mrcs, lazy=True, datadir=args.datadir)
oldD = old[0].get().shape[0]
assert args.D < oldD, f'New box size {args.D} must be smaller than original box size {oldD}'
assert args.D % 2 == 0, 'New box size must be even'
D = args.D
start = int(oldD/2 - D/2)
stop = int(oldD/2 + D/2)
if args.is_vol:
oldft = fft.htn_center(np.array([x.get() for x in old]))
log(oldft.shape)
newft = oldft[start:stop,start:stop,start:stop]
log(newft.shape)
new = fft.ihtn_center(newft).astype(np.float32)
log('Saving {}'.format(args.o))
mrc.write(args.o,new)
elif args.chunk is None:
new = []
for i in range(len(old)):
if i % 1000 == 0:
log(f'Processing image {i} of {len(old)}')
img = old[i]
oldft = fft.ht2_center(img.get()).astype(np.float32)
newft = oldft[start:stop, start:stop]
new.append(fft.ihtn_center(newft).astype(np.float32))
assert oldft[int(oldD/2),int(oldD/2)] == newft[int(D/2),int(D/2)]
new = np.asarray(new)
log(new.shape)
log('Saving {}'.format(args.o))
mrc.write(args.o,new)
else:
nchunks = len(old) // args.chunk + 1
for i in range(nchunks):
log('Processing chunk {}'.format(i))
out = '.{}'.format(i).join(os.path.splitext(args.o))
new = []
for img in old[i*args.chunk:(i+1)*args.chunk]:
oldft = fft.ht2_center(img.get()).astype(np.float32)
newft = oldft[start:stop, start:stop]
new.append(fft.ihtn_center(newft).astype(np.float32))
assert oldft[int(oldD/2),int(oldD/2)] == newft[int(D/2),int(D/2)]
new = np.asarray(new)
log(new.shape)
log('Saving {}'.format(out))
mrc.write(out,new)
def __init__(self, mrcfile, mrcfile_tilt, norm=None, keepreal=False, invert_data=False, ind=None, window=True, datadir=None):
# load untilted
particles_real = load_particles(mrcfile, False, datadir)
# load tilt series
particles_tilt_real = load_particles(mrcfile_tilt, False, datadir)
# filter dataset
if ind is not None:
particles_real = particles_real[ind]
particles_tilt_real = particles_tilt_real[ind]
N, ny, nx = particles_real.shape
assert ny == nx, "Images must be square"
assert ny % 2 == 0, "Image size must be even"
log('Loaded {} {}x{} images'.format(N, ny, nx))
assert particles_tilt_real.shape == (N, ny, nx), "Tilt series pair must have same dimensions as untilted particles"
log('Loaded {} {}x{} tilt pair images'.format(N, ny, nx))
# Real space window
if window:
m = window_mask(ny, .85, .99)
particles_real *= m
particles_tilt_real *= m
# compute HT
particles = np.asarray([fft.ht2_center(img) for img in particles_real]).astype(np.float32)
particles_tilt = np.asarray([fft.ht2_center(img) for img in particles_tilt_real]).astype(np.float32)
if invert_data:
particles *= -1
particles_tilt *= -1
# symmetrize HT
particles = fft.symmetrize_ht(particles)
particles_tilt = fft.symmetrize_ht(particles_tilt)
# normalize
if norm is None:
norm = [np.mean(particles), np.std(particles)]
norm[0] = 0
particles = (particles - norm[0])/norm[1]
particles_tilt = (particles_tilt - norm[0])/norm[1]
log('Normalized HT by {} +/- {}'.format(*norm))
self.particles = particles
self.particles_tilt = particles_tilt
self.norm = norm
self.N = N
self.D = particles.shape[1]
self.keepreal = keepreal
if keepreal:
self.particles_real = particles_real
self.particles_tilt_real = particles_tilt_real
def get(self, i):
img = self.particles[i].get()
if self.window is not None:
img *= self.window
img = fft.ht2_center(img).astype(np.float32)
if self.invert_data: img *= -1
img = fft.symmetrize_ht(img)
img = (img - self.norm[0])/self.norm[1]
return img
def estimate_normalization(self, n=1000):
n = min(n,self.N)
imgs = np.asarray([fft.ht2_center(self.particles[i].get()) for i in range(0,self.N, self.N//n)])
if self.invert_data: imgs *= -1
imgs = fft.symmetrize_ht(imgs)
norm = [np.mean(imgs), np.std(imgs)]
norm[0] = 0
log('Normalizing HT by {} +/- {}'.format(*norm))
return norm
def __init__(self, mrcfile, norm=None, keepreal=False, invert_data=False, ind=None, window=True, datadir=None):
particles_real = load_particles(mrcfile, False, datadir)
if ind is not None:
particles_real = particles_real[ind]
N, ny, nx = particles_real.shape
assert ny == nx, "Images must be square"
assert ny % 2 == 0, "Image size must be even"
log('Loaded {} {}x{} images'.format(N, ny, nx))
# Real space window
if window:
particles_real *= window_mask(ny, .85, .99)
# compute HT
particles = np.asarray([fft.ht2_center(img) for img in particles_real])
particles = particles.astype(np.float32)
if invert_data: particles *= -1
# symmetrize HT
particles = fft.symmetrize_ht(particles)
# normalize
if norm is None:
norm = [np.mean(particles), np.std(particles)]
norm[0] = 0
particles = (particles - norm[0])/norm[1]
log('Normalized HT by {} +/- {}'.format(*norm))
self.particles = particles
self.N = N
self.D = particles.shape[1] # ny + 1 after symmetrizing HT
self.norm = norm
self.keepreal = keepreal
if keepreal:
self.particles_real = particles_real
log('Normalized real space images by {}'.format(particles_real.std()))
self.particles_real /= particles_real.std()
import pickle
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
sys.path.insert(0,'../lib-python')
import fft
import models
import mrc
from lattice import Lattice
imgs,_,_ = mrc.parse_mrc('data/hand.mrcs')
img = imgs[0]
D = img.shape[0]
ht = fft.ht2_center(img)
ht = fft.symmetrize_ht(ht)
D += 1
lattice = Lattice(D)
model = models.FTSliceDecoder(D**2, D, 10,10,nn.ReLU)
coords = lattice.coords[...,0:2]/2
ht = torch.tensor(ht.astype(np.float32)).view(1,-1)
trans = torch.tensor([5.,10.]).view(1,1,2)
ht_shifted = lattice.translate_ht(ht, trans)
ht_np = ht_shifted.view(D,D).numpy()[0:-1, 0:-1]
img_shifted = fft.ihtn_center(ht_np)