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)