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)