def open_corrected_file(output_filename, shape, voxel_size): # Set the data type dtype = numpy.dtype(numpy.float32) # Create the extended header extended_header = numpy.zeros( shape=shape[0:2], dtype=mrcfile.dtypes.FEI1_EXTENDED_HEADER_DTYPE ) # Open the file outfile = mrcfile.new_mmap( output_filename, shape=(0, 0, 0, 0), mrc_mode=mrcfile.utils.mode_from_dtype(dtype), overwrite=True, ) # Some hacks to set the extended header without having to resize whole file. outfile._check_writeable() outfile._close_data() outfile._extended_header = extended_header outfile.header.nsymbt = extended_header.nbytes outfile.header.exttyp = "FEI1" outfile._open_memmap(dtype, shape) outfile.update_header_from_data() outfile.flush() # Set the voxel size outfile.voxel_size = voxel_size # Return the handle return outfile
def scale(origin_path, new_path, scale_rate=1): origin_mrc = mrcfile.mmap(origin_path, 'r') origin_shape = np.array(origin_mrc.data.shape) ox, oy, oz = origin_shape scale = SCALE_BASE**scale_rate new_shape = (origin_shape + scale - 1) // scale nx, ny, nz = new_shape new_mrc = mrcfile.new_mmap(new_path, (nz, nx, ny), mrc_mode=mrcfile.utils.mode_from_dtype( mrcfile.utils.data_dtype_from_header( origin_mrc.header)), overwrite=True) print('scaling data') for z in tqdm(range(nz)): sz = z * scale for x in range(nx): sx = x * scale for y in range(ny): sy = y * scale new_mrc.data[z, x, y] = np.mean( origin_mrc.data[sz:min(sz + scale, oz), sx:min(sx + scale, ox), sy:min(sy + scale, oy)]) print('updating header') new_mrc.update_header_stats() print('header updated') origin_mrc.close() new_mrc.close()
def __init__(self, filename, shape, pixel_size, dtype="uint8"): """ Initialise the writer Args: filename (str): The filename shape (tuple): The shape of the data pixel_size (float): The pixel size dtype (str): The data type """ # Get the dtype dtype = np.dtype(dtype) # Convert 32bit int and 64bit float if dtype == "int32": dtype = np.dtype(np.int16) elif dtype == "uint32": dtype = np.dtype(np.uint16) elif dtype == "float64": dtype = np.dtype(np.float32) elif dtype == "complex128": dtype = np.dtype(np.complex64) # Open the handle to the mrcfile self.handle = mrcfile.new_mmap( filename, shape=(0, 0, 0), mrc_mode=mrcfile.utils.mode_from_dtype(dtype), overwrite=True, ) # Setup the extended header extended_header = np.zeros(shape=shape[0], dtype=FEI_EXTENDED_HEADER_DTYPE) # Set the extended header self.handle._check_writeable() self.handle._close_data() self.handle._extended_header = extended_header self.handle.header.nsymbt = extended_header.nbytes self.handle.header.exttyp = "FEI1" self.handle._open_memmap(dtype, shape) self.handle.update_header_from_data() self.handle.flush() # Set the pixel size self.handle.voxel_size = pixel_size for i in range(self.handle.extended_header.shape[0]): self.handle.extended_header[i]["Pixel size X"] = pixel_size * 1e-10 self.handle.extended_header[i]["Pixel size Y"] = pixel_size * 1e-10 self.handle.extended_header[i]["Application"] = "RFI Simulation" # Set the data array self._data = self.handle.data self._angle = MrcFileWriter.AngleProxy(self.handle) self._position = MrcFileWriter.PositionProxy(self.handle) self._drift = MrcFileWriter.DriftProxy(self.handle) self._defocus = MrcFileWriter.DefocusProxy(self.handle)
def testUpdate(self): # Create a tmpdir in a context. Cleans up on exit. with tempfile.TemporaryDirectory() as tmpdir: # Create two filenames in our tmpdir mrcs_filepath = os.path.join(tmpdir, "test.mrc") files = (f"{mrcs_filepath}.1", f"{mrcs_filepath}.2") # Note below we will fix the time to avoid racy unit tests. epoch = datetime(1970, 1, 1).strftime("%Y-%m-%d %H:%M:%S") # The time is also packed into the label by mrcfile package. label = "{0:40s}{1:>39s} ".format("Created by aspire unit test", epoch) with mrcfile.new_mmap(files[0], shape=(self.n, self.n), mrc_mode=2, overwrite=True) as mrc: mrc.data[:, :] = self.a mrc.update_header_from_data() self.stats.update_header(mrc) mrc.header.time = epoch mrc.header.label[0] = label with mrcfile.new_mmap(files[1], shape=(self.n, self.n), mrc_mode=2, overwrite=True) as mrc: mrc.set_data(self.a.astype(np.float32)) mrc.header.time = epoch mrc.header.label[0] = label # Our homebrew and mrcfile files should now match to the bit. comparison = sha256sum(files[0]) == sha256sum(files[1]) # Expected hash: # 71355fa0bcd5b989ff88166962ea5d2b78ea032933bd6fda41fbdcc1c6d1a009 logging.debug(f"sha256(file0): {sha256sum(files[0])}") logging.debug(f"sha256(file1): {sha256sum(files[1])}") self.assertTrue(comparison)
def test_new_mmap(self): with mrcfile.new_mmap(self.temp_mrc_name, (3, 4, 5, 6), mrc_mode=2, fill=1.1) as mrc: assert repr(mrc) == ("MrcMemmap('{0}', mode='w+')".format( self.temp_mrc_name)) assert mrc.data.shape == (3, 4, 5, 6) assert np.all(mrc.data == 1.1) assert mrc.header.nx == 6 file_size = mrc._iostream.tell( ) # relies on flush() leaving stream at end assert file_size == mrc.header.nbytes + mrc.data.nbytes
def open_reconstruction_file(output_filename, shape, voxel_size): # Open the file outfile = mrcfile.new_mmap(output_filename, overwrite=True, mrc_mode=2, shape=shape) # Set the voxel size outfile.voxel_size = voxel_size # Return the handle return outfile
def generate_ctf( filename=None, image_size=None, pixel_size=1, defocus=0, num_defocus=None, step_defocus=None, spherical_aberration=2.7, astigmatism=0, astigmatism_angle=0, energy=300, defocus_spread=None, source_spread=0.02, chromatic_aberration=2.7, energy_spread=0.33e-6, current_spread=0.33e-6, acceleration_voltage_spread=0.8e-6, phase_shift=0, component="imag", envelope=True, recentre=False, ): """ Generate a CTF image Args: filename (str): The output filename image_size (tuple): The image size in pixels pixel_size (float): The pixel size (A) defocus (float): The defocus (df in A) spherical_aberration (float): The spherical aberration (Cs in mm) astigmatism (float): The astigmatism (A) astigmatism_angle (float): The astigmatism angle (deg) energy (float): The electron energy (keV) defocus_spread (float): The defocus spread (A) source_spread (float): The source spread (mrad) Cc (float): The chromatic aberrationa (mm) dEE (float): dE/E, the fluctuation in the electron energy dII (float): dI/I, the fluctuation in the lens current dVV (float): dV/V, the fluctuation in the acceleration voltage phase_shift (float): The phase shift (rad) - 0 without phase plate, pi/2 with phase plate component (str): The bit to plot (real, imag, abs) envelope (bool): Plot the envelope function """ # Convert some quantities source_spread = source_spread / 1e3 # mrad -> rad spherical_aberration = spherical_aberration * 1e7 # mm -> A chromatic_aberration = chromatic_aberration * 1e7 # mm -> A astigmatism_angle = astigmatism_angle * pi / 190 # deg -> rad energy = energy * 1e3 # keV -> eV # Compute the defocus spread if defocus_spread is None: defocus_spread = guanaco.detail.get_defocus_spread( Cc=chromatic_aberration, dEE=energy_spread, dII=current_spread, dVV=acceleration_voltage_spread, ) # Check the defocus if defocus is None: defocus = 0 if num_defocus is None or num_defocus == 0: num_defocus = 1 if step_defocus is None: step_defocus = 0 # Compute the wavelength wavelength = guanaco.detail.get_electron_wavelength(energy) # Init output ctf file handle = mrcfile.new_mmap( filename, shape=(num_defocus, image_size[0], image_size[1]), mrc_mode=mrcfile.utils.mode_from_dtype(numpy.dtype(numpy.complex64)), overwrite=True, ) # Set the voxel size handle.voxel_size = pixel_size # Loop through the defoci for i in range(num_defocus): # Set the defocus df = defocus + (i - 0.5 * (num_defocus - 1)) * step_defocus print("Computing CTF for defocus = %d A" % df) # Setup the CTF calculation ctf_calculator = guanaco.detail.CTF( l=wavelength, df=df, Cs=spherical_aberration, Ca=astigmatism, Pa=astigmatism_angle, dd=defocus_spread, theta_c=source_spread, phi=phase_shift, ) # Evaluate the ctf ctf = ctf_calculator.get_ctf(image_size[1], image_size[0], pixel_size) # Recentre if recentre: ctf = numpy.fft.fftshift(ctf) # Set the CTF handle.data[i, :, :] = ctf
def write_occupancy_multiscale_summary(image_resolution, dataset, model, model_input, gt, model_output, writer, total_steps, prefix='train_', output_mrc='test.mrc', skip=False, oversample=1.0, max_chunk_size=1024, mode='binary'): if skip: return model_input = dataset.get_eval_samples(oversample) print("Summary: Write occupancy multiscale summary...") # convert to cuda and add batch dimension tmp = {} for key, value in model_input.items(): if isinstance(value, torch.Tensor): tmp.update({key: value[None, ...]}) else: tmp.update({key: value}) model_input = tmp print("Summary: processing...") pred_occupancy = process_batch_in_chunks( model_input, model, max_chunk_size=max_chunk_size)['model_out']['output'] # get voxel idx for each coordinate coords = model_input['fine_abs_coords'].detach().cpu().numpy() voxel_idx = np.floor((coords + 1.) / 2. * (dataset.sidelength[0] * oversample)).astype(np.int32) voxel_idx = voxel_idx.reshape(-1, 3) # init a new occupancy volume display_occupancy = -1 * np.ones(image_resolution, dtype=np.float32) # assign predicted voxel occupancy values into the array pred_occupancy = pred_occupancy.reshape(-1, 1).detach().cpu().numpy() display_occupancy[voxel_idx[:, 0], voxel_idx[:, 1], voxel_idx[:, 2]] = pred_occupancy[..., 0] print(f"Summary: write MRC file {image_resolution}") if mode == 'hq': print("\tWriting float") with mrcfile.new_mmap(output_mrc, overwrite=True, shape=image_resolution, mrc_mode=2) as mrc: mrc.data[voxel_idx[:, 0], voxel_idx[:, 1], voxel_idx[:, 2]] = pred_occupancy[..., 0] elif mode == 'binary': print("\tWriting binary") with mrcfile.new_mmap(output_mrc, overwrite=True, shape=image_resolution) as mrc: mrc.data[voxel_idx[:, 0], voxel_idx[:, 1], voxel_idx[:, 2]] = pred_occupancy[..., 0] > 0 if writer is not None: print("Summary: Draw octtree") fig = dataset.octtree.draw() writer.add_figure(prefix + 'tiling', fig, global_step=total_steps) return display_occupancy
def save_star(image_source, starfile_filepath, batch_size=1024, save_mode=None, overwrite=False): """ Save an ImageSource to a STAR file + individual .mrcs files Note that .mrcs files are saved at the same location as the STAR file. :param image_source: The ImageSource object to save :param starfile_filepath: Path to STAR file where we want to save image_source :param batch_size: Batch size of images to query from the `ImageSource` object. Every `batch_size` rows, entries are written to STAR file, and the `.mrcs` files saved. :param save_mode: Whether to save all images in a single or multiple files in batch size. :param overwrite: Whether to overwrite any .mrcs files found at the target location. :return: None """ # TODO: Accessing protected member - provide a way to get a handle on the _metadata attribute. df = image_source._metadata.copy() # Drop any column that doesn't start with a *single* underscore df = df.drop([ str(col) for col in df.columns if not col.startswith('_') or col.startswith('__') ], axis=1) # Create a new column that we will be populating in the loop below df['_rlnImageName'] = '' with open(starfile_filepath, 'w') as f: if save_mode == 'single': # save all images into one single mrc file mrcs_filename = os.path.splitext( os.path.basename( starfile_filepath))[0] + f'_{0}_{image_source.n-1}.mrcs' mrcs_filepath = os.path.join(os.path.dirname(starfile_filepath), mrcs_filename) df['_rlnImageName'][0:image_source.n] = pd.Series([ '{0:06}@{1}'.format(j + 1, mrcs_filepath) for j in range(image_source.n) ]) with mrcfile.new_mmap(mrcs_filepath, shape=(image_source.n, image_source.L, image_source.L), mrc_mode=2, overwrite=overwrite) as mrc: for i_start in np.arange(0, image_source.n, batch_size): i_end = min(image_source.n, i_start + batch_size) num = i_end - i_start logger.info( f'Saving ImageSource[{i_start}-{i_end-1}] to {mrcs_filepath}' ) mrc.data[i_start:i_end, :, :] = np.swapaxes( image_source.images(start=i_start, num=num).data.astype('float32'), 0, 2) mrc.close() else: # save all images into multiple mrc files in batch size for i_start in np.arange(0, image_source.n, batch_size): i_end = min(image_source.n, i_start + batch_size) num = i_end - i_start mrcs_filename = os.path.splitext( os.path.basename( starfile_filepath))[0] + f'_{i_start}_{i_end-1}.mrcs' mrcs_filepath = os.path.join( os.path.dirname(starfile_filepath), mrcs_filename) logger.info( f'Saving ImageSource[{i_start}-{i_end-1}] to {mrcs_filepath}' ) im = image_source.images(start=i_start, num=num) im.save(mrcs_filepath, overwrite=overwrite) df['_rlnImageName'][i_start:i_end] = pd.Series([ '{0:06}@{1}'.format(j + 1, mrcs_filepath) for j in range(num) ]) # initial the star file object and save it starfile = StarFile(blocks=[StarFileBlock(loops=[df])]) starfile.save(f)
def stitchCapsomers(mrc_pentavalent, mrc_hexavalent, vector_pent, vector_hex, mapbox, angpix, output): """ Make ndarray to store the output map """ capsid = np.zeros((mapbox, mapbox, mapbox), dtype=np.float32) correction_mask = np.zeros_like(capsid, dtype=np.int) """ Take input """ vector_pentavalent = np.array( [vector_pent[0], vector_pent[1], vector_pent[2]]) vector_hexavalent = np.array([vector_hex[0], vector_hex[1], vector_hex[2]]) """ Grab the image values """ ndarray_pentavalent = openMRCfile(mrc_pentavalent) ndarray_hexavalent = openMRCfile(mrc_hexavalent) print("dtype of mrc input is", ndarray_pentavalent.dtype) """ Get box size for capsomers """ subbox_pentavalent = ndarray_pentavalent.shape[0] subbox_hexavalent = ndarray_hexavalent.shape[0] """ Get parameters for requested symmetry operation """ params_pent = AreaOfInterest('fivefold', vector_pentavalent) params_hex = AreaOfInterest('fullexpand', vector_hexavalent) """ Quaternions to return to original orientation """ pent_quat = Quaternion(params_pent.quatZ).inverse hex_quat = Quaternion(params_hex.quatZ).inverse """ Fetch the indices for non-redundant symmetry operations """ fivefold_indices = getUniqueVertices(vector_pent) fullexpand_indices = getUniqueVertices(vector_hex) """ Update to use only the unique symmetry operations """ params_pent.updateSymIndices(fivefold_indices) symops_pent = I1Quaternions[params_pent.symindices] symops_hex = I1Quaternions """ Faster debugging by reducing number of subvolumes relocated """ #symops_pent = symops_pent[0:2] #symops_hex = symops_hex[0:2] total_symops = int(len(symops_pent)) + int(len(symops_hex)) """ Make array to store the x,y,z of each capsomer wrt the capsid """ capsomer_centers = np.zeros( (total_symops, 3)) # Works for papillomavirus and polyomavirus """ Catch bad pentavalent vector, while still allowing for debug """ if (len(symops_pent) > 12): print( len(symops_pent), "symops for pentavalent suggests bad vector. Ensure ratio is 0 0.618 1" ) sys.exit() """ Store the centers. Rotate before converting to ZYX ordering """ for index, symop in enumerate(symops_pent, start=0): # papillomavirus and polyomavirus center = Quaternion(symop).inverse.rotate(vector_pentavalent) capsomer_centers[index] = np.array([(1 * center[2]), (-1 * center[1]), (-1 * center[0])]).astype(np.int) for index, symop in enumerate( symops_hex, start=len(symops_pent)): # papillomavirus and polyomavirus center = Quaternion(symop).inverse.rotate(vector_hexavalent) capsomer_centers[index] = np.array([(1 * center[2]), (-1 * center[1]), (-1 * center[0])]).astype(np.int) """ Adjust the pentavalent capsomers """ for index, symop in enumerate(symops_pent, start=0): capsomer, correction_mask = adjustCapsomer( ndarray_pentavalent, correction_mask, symop, pent_quat, vector_pentavalent, capsomer_centers, subbox_pentavalent, mapbox) capsid = np.add(capsid, capsomer) print("Pentavalent capsomer", index, "completed.") """ Adjust the hexavalent capsomers """ for index, symop in enumerate(symops_hex, start=0): capsomer, correction_mask = adjustCapsomer(ndarray_hexavalent, correction_mask, symop, hex_quat, vector_hexavalent, capsomer_centers, subbox_hexavalent, mapbox) capsid = np.add(capsid, capsomer) print("Hexavalent capsomer", index, "completed.") print("DEBUG: Correction_mask currently contains values", np.unique(correction_mask)) """ Cast to float32. Might be unneccessary """ capsid = capsid.astype(np.float32) """ correction_mask currently contains values >1 at overlaps """ """ should convert zeros to ones before dividing """ correction_mask_offset = correction_mask == 0 correction_mask = correction_mask.astype( np.int) + correction_mask_offset.astype(np.int) """ Correct the voxels that are overweighted. Avoid divide-by-zero error """ capsid = np.divide(capsid, correction_mask.astype(np.float32)) """ Write the output mrc file """ mrc = mrcfile.new_mmap(output, shape=(mapbox, mapbox, mapbox), mrc_mode=2, overwrite=True) mrc.set_data(capsid) mrc.voxel_size = angpix mrc.header.label[1] = str( 'Created using ISECC_recombine, Goetschius DJ 2020') mrc.close() return
\n_rlnAmplitudeContrast #10 \ \n_rlnDefocusAngle #11 \ \n_rlnCtfBfactor #12 \ \n_rlnPhaseShift #13 \ \n_rlnPixelSize #14 \ \n_rlnImageName #15 \ \n') # initiate spider file: if matlab: spiderFile = open(spiderOut, 'w') ########################## # create empty arrays: img_array = mrcfile.new_mmap(stackOut, shape=(snapshots, box, box), mrc_mode=2, overwrite=True) #mrc_mode 2: float32 if matlab: binary_array = np.memmap(binaryOut, dtype='float32', mode='w+', shape=(snapshots, box, box)) #, offset=4*box**2) ########################## # INITIATE MAIN LOOP: ########################## # import original stacks: idx = 0 #index for stack/stars img = 0 #index for every individual image in new stack for z in stackPaths: stack = mrcfile.open(z)
head=gethead(args.good_star) print(head) coordlist,namelist=readstar(args.good_star) #print(coordlist) ''' for name in slist: print(name,get_col(head,name)) ''' x=args.x y=args.y z=args.z new_tomo=np.zeros([x,y,z],dtype=np.float32) i=0 o=mrcfile.new_mmap(args.output,overwrite='True',shape=new_tomo.shape,mrc_mode=0) for subname in namelist: sub_open=mrcfile.open(subname,'r+') sub_data=sub_open.data sub_open.close() sub=np.float32(sub_data) draw_subtomo(new_tomo,sub,coordlist[i]) o.set_data(new_tomo) o.flush() o.close() #print(i,coords[0])
def save_images(self, starfile_filepath, filename_indices=None, batch_size=512, overwrite=False): """ Save an ImageSource to MRCS files Note that .mrcs files are saved at the same location as the STAR file. :param filename_indices: Filename list for save all images :param starfile_filepath: Path to STAR file where we want to save image_source :param batch_size: Batch size of images to query from the `ImageSource` object. if `save_mode` is not `single`, images in the same batch will save to one MRCS file. :param overwrite: Whether to overwrite any .mrcs files found at the target location. :return: None """ if filename_indices is None: # Generate filenames from metadata filename_indices = [ self._metadata["_rlnImageName"][i].split("@")[1] for i in range(self.n) ] # get the save_mode from the file names unique_filename = set(filename_indices) save_mode = None if len(unique_filename) == 1: save_mode = "single" if save_mode == "single": # Save all images into one single mrc file # First, construct name for mrc file fdir = os.path.dirname(starfile_filepath) mrcs_filepath = os.path.join(fdir, filename_indices[0]) # Open new MRC file with mrcfile.new_mmap( mrcs_filepath, shape=(self.n, self.L, self.L), mrc_mode=2, overwrite=overwrite, ) as mrc: stats = MrcStats() # Loop over source setting data into mrc file for i_start in np.arange(0, self.n, batch_size): i_end = min(self.n, i_start + batch_size) num = i_end - i_start logger.info( f"Saving ImageSource[{i_start}-{i_end-1}] to {mrcs_filepath}" ) datum = self.images(start=i_start, num=num).data.astype("float32") # Assign to mrcfile mrc.data[i_start:i_end] = datum # Accumulate stats stats.push(datum) # To be safe, explicitly set the header # before the mrc file context closes. mrc.update_header_from_data() # Also write out updated statistics for this mrc. # This should be an optimization over mrc.update_header_stats # for large arrays. stats.update_header(mrc) else: # save all images into multiple mrc files in batch size for i_start in np.arange(0, self.n, batch_size): i_end = min(self.n, i_start + batch_size) num = i_end - i_start mrcs_filepath = os.path.join( os.path.dirname(starfile_filepath), filename_indices[i_start]) logger.info( f"Saving ImageSource[{i_start}-{i_end-1}] to {mrcs_filepath}" ) im = self.images(start=i_start, num=num) im.save(mrcs_filepath, overwrite=overwrite)