def subPixelPeak(inp, k=1, ignore_border=1):
from pytom.voltools import transform
is2d = len(inp.shape.squeeze()) == 2
if is2d:
inp2 = inp.squeeze()[ignore_border:-ignore_border,
ignore_border:-ignore_border]
scale = (k, k, 1)
out = xp.array(inp.squeeze(), dtype=xp.float32)
x, y = xp.array(xp.unravel_index(inp2.argmax(),
inp2.shape)) + ignore_border
out = xp.expand_dims(out, 2)
translation = [inp.shape[0] // 2 - x, inp.shape[1] // 2 - y, 0]
else:
inp2 = inp[ignore_border:-ignore_border, ignore_border:-ignore_border,
ignore_border:-ignore_border]
scale = (k, k, k)
out = xp.array(inp, dtype=xp.float32)
x, y, z = xp.array(xp.unravel_index(inp2.argmax(),
inp2.shape)) + ignore_border
translation = [
inp.shape[0] // 2 - x, inp.shape[1] // 2 - y, inp.shape[2] // 2 - z
]
zoomed = xp.zeros_like(out, dtype=xp.float32)
transform(out,
output=zoomed,
scale=scale,
translation=translation,
device='gpu:0',
interpolation='filt_bspline')
shifts = [x2, y2, z2] = xp.unravel_index(zoomed.argmax(), zoomed.shape)
return [zoomed[x2, y2, z2], shifts[3 - is2d]]
def normalised_cross_correlation(first, second, filter_mask=None, device=0):
"""
Do a cross correlation based on numpy
@param first: The first dataset (numpy 2D)
@type first: numpy array 2D
@param second: The second dataset (numpy 2D)
@type second: numpy array 2D
@param filter_mask: a filter which is used to filter image 'first'
@type filter_mask: numpy array 2D
@return: The cross correlation result
@returntype: numpy array 2D
@requires: the shape of first to be equal to the shape of second, and equal t the shape of the filter (if used of course)
"""
assert first.shape == second.shape
assert len(first.shape) == 2
if not (filter_mask is None): assert first.shape == filter_mask.shape
# if filter_mask is None:
ffirst = xp.fft.fft2(xp.array(first))
# else:
# ffirst = xp.fft.fftshift(xp.fft.fftshift(xp.fft.fftn(first)) * filter_mask)
ccmap = xp.real(
xp.fft.fftshift(
xp.fft.ifftn(xp.multiply(xp.fft.fftn(second),
xp.conj(ffirst))))) / first.size
return ccmap
def cut_from_projection(proj, center, size, device=2):
"""Cut out a subregion out from a 2D projection.
@param proj: 2D projection.
@param center: cutting center.
@param size: cutting size.
@return: subprojection.
"""
import pytom.voltools as vt
if len(proj.shape) > 2 and proj.shape[2] > 1:
raise Exception(
'We assume that projections come from a 3D object, thus your projection cannot be a 3D object itself'
)
import numpy as np
from scipy.ndimage import map_coordinates
try:
proj = proj.squeeze().get()
except:
proj = proj.squeeze()
#v =
# adjusted to python3
grid = np.mgrid[center[0] - size[0] // 2:center[0] + size[0] -
size[0] // 2 - 1:size[0] * 1j,
center[1] - size[1] // 2:center[1] + size[1] -
size[1] // 2 - 1:size[1] * 1j]
v = map_coordinates(proj, grid, order=2)
return xp.array(v)
def subPixelShifts(volume, scale=[1, 1, 1], cutout=20, shifts=True):
from pytom.voltools import transform
import numpy as np
volC = xp.array(
[volume.shape[0] // 2, volume.shape[1] // 2, volume.shape[2] // 2])
cx, cy, cz = center = find_coords_max_ccmap(volume)
x, y, z = volume.shape
if cx - cutout // 2 < 0 or cy - cutout // 2 < 0 or cz - cutout // 2 < 0 or cx + cutout // 2 >= x or cy + cutout // 2 >= y or cz + cutout // 2 >= z:
return [
volume.max(),
xp.array([cx - x // 2, cy - y // 2, cz - z // 2])
]
cutvol = volume[cx - cutout // 2:cx + cutout // 2,
cy - cutout // 2:cy + cutout // 2,
cz - cutout // 2:cz + cutout // 2]
nx, ny, nz = [int(np.ceil(cutout / s)) for s in scale]
nC = xp.array([nx // 2, ny // 2, nz // 2])
centered = xp.zeros((nx, ny, nz), dtype=xp.float32)
centered[nx // 2 - cutout // 2:nx // 2 + cutout // 2,
ny // 2 - cutout // 2:ny // 2 + cutout // 2,
nz // 2 - cutout // 2:nz // 2 + cutout // 2] = cutvol
interpolated_peak = xp.zeros_like(centered)
transform(centered,
scale=scale,
device=device,
interpolation='filt_bspline',
output=interpolated_peak)
interpol_center = find_coords_max_ccmap(interpolated_peak)
shifts = xp.array(center) - volC + (xp.array(interpol_center) -
xp.array(nC)) * xp.array(scale)
return [interpolated_peak.max(), shifts]
def subPixelShiftsImage(volume, scale=[1, 1, 1], cutout=20, shifts=True):
from pytom.voltools import transform
import numpy as np
if shifts:
volC = xp.array(
[vol.shape[0] // 2, volume.shape[1] // 2, vol.shape[2] // 2])
else:
volC = xp.array([0, 0, 0], dtype=xp.int32)
cx, cy, cz = center = find_coords_max_ccmap(volume)
cutvol = volume[cx - cutout // 2:cx + cutout // 2,
cy - cutout // 2:cy + cutout // 2, :]
nx, ny, nz = [int(np.ceil(cutout / s)) for s in scale]
nC = xp.array([nx // 2, ny // 2, 0])
centered = xp.zeros((nx, ny, 1), dtype=xp.float32)
centered[nx // 2 - cutout // 2:nx // 2 + cutout // 2,
ny // 2 - cutout // 2:ny // 2 + cutout // 2, :] = cutvol
interpolated_peak = transform(centered,
scale=scale,
device=device,
interpolation='filt_bspline')
interpol_center = find_coords_max_ccmap(interpolated_peak)
shifts = xp.array(center) - volC + (xp.array(interpol_center) -
xp.array(nC)) * xp.array(scale)
return [interpolated_peak.max(), shifts]
def exact_filter(tilt_angles, tiltAngle, sX, sY, sliceWidth=1, arr=[]):
"""
exactFilter: Generates the exact weighting function required for weighted backprojection - y-axis is tilt axis
Reference : Optik, Exact filters for general geometry three dimensional reconstuction, vol.73,146,1986.
@param tilt_angles: list of all the tilt angles in one tilt series
@param titlAngle: tilt angle for which the exact weighting function is calculated
@param sizeX: size of weighted image in X
@param sizeY: size of weighted image in Y
@return: filter volume
"""
import numpy as xp
# Calculate the relative angles in radians.
diffAngles = (xp.array(tilt_angles) - tiltAngle) * xp.pi / 180.
# Closest angle to tiltAngle (but not tiltAngle) sets the maximal frequency of overlap (Crowther's frequency).
# Weights only need to be calculated up to this frequency.
sampling = xp.min(xp.abs(diffAngles)[xp.abs(diffAngles) > 0.001])
crowtherFreq = min(sX // 2,
xp.int32(xp.ceil(sliceWidth / xp.sin(sampling))))
arrCrowther = xp.matrix(
xp.abs(xp.arange(-crowtherFreq, min(sX // 2, crowtherFreq + 1))))
# Calculate weights
wfuncCrowther = 1. / (xp.clip(
1 - xp.array(xp.matrix(xp.abs(xp.sin(diffAngles))).T * arrCrowther)**2,
0, 2)).sum(axis=0)
# Create full with weightFunc
wfunc = xp.ones((sX, sY), dtype=xp.float32)
# row_stack is not implemented in cupy
weightingFunc = xp.column_stack(([
(wfuncCrowther),
] * (sY))).T
wfunc[:, sX // 2 - crowtherFreq:sX // 2 +
min(sX // 2, crowtherFreq + 1)] = weightingFunc
return wfunc
def design_fsc_filter(fsc, fildim=None, fsc_criterion=0.143):
"""
design spectral filter to weight by SNR of frequency
@param fsc: input fsc
@type fsc: 1-d list
@param fildim: filter dimension
@type fildim: int
@return: filter
@rtype: list
@author: FF
"""
from math import sqrt
from pytom.basic.resolution import getResolutionBandFromFSC
if not fildim:
fildim = len(fsc)
nband = len(fsc)
if fsc_criterion != 0.0:
resolutionBand = getResolutionBandFromFSC(fsc, criterion=fsc_criterion)
smooth = max(resolutionBand / 5, 2)
else:
resolutionBand = len(fsc)
smooth = 1
#print "filter: fsc_criterion %2.3f, resolutionBand %d" % (fsc_criterion, resolutionBand)
# filter by sqrt(FSC)
fsc_fil = len(fsc) * [0.]
for ii in range(0, len(fsc)):
fscval = fsc[ii]
if fscval > 0.:
#fsc_fil[ii] = sqrt(fsc[ii])
if ii <= resolutionBand:
fsc_fil[ii] = sqrt(fsc[ii])
elif (ii > resolutionBand) and (ii <= resolutionBand + smooth):
fsc_fil[ii] = sqrt(fsc[ii]) * ((
(resolutionBand + smooth) - ii) /
smooth)**2 # squared filter
else:
fsc_fil[ii] = 0.
else:
fsc_fil[ii] = 0.
#design filter
fil = xp.array(fildim * [0.])
if nband != len(fil):
shrinkfac = 1. / (len(fil) / nband)
for ii in range(len(fil) - 1):
# linear resample fsc if nband ~= size(fil)
if nband != len(fil):
ilow = int(xp.floor(shrinkfac * ii))
dlow = shrinkfac * ii - ilow
ihi = ilow + 1
dhi = ihi - shrinkfac * ii
fil[ii] = fsc_fil[ilow] * dhi + fsc_fil[ihi] * dlow
else:
fil[ii] = fsc_fil[ii]
return fil
def read_em(filename, order='F', keepnumpy=False, deviceID=None):
"""Read EM file. Now only support read the type float32 on little-endian machines.
@param filename: file name to read.
@return The data from the EM file in ndarray
"""
f = open(filename, 'r')
try:
dt_header = np.dtype('int32')
header = np.fromfile(f, dt_header, 128)
x = header[1]
y = header[2]
z = header[3]
# read the data type
dt = int(hex(header[0])[2])
default_type = False
if dt == 1: # byte
raise Exception("Data type not supported yet!")
elif dt == 2: # short
dt_data = np.dtype('<i2')
elif dt == 4: # long
dt_data = np.dtype('<i4')
elif dt == 5: # float32
dt_data = np.dtype('<f4') # little-endian, float32
default_type = True
elif dt == 8: # float complex
raise Exception("Data type not supported yet!")
elif dt == 9: # double
dt_data = np.dtype('<f8')
elif dt == 10: # double complex
raise Exception("Data type not supported yet!")
else:
raise Exception("Data type not supported yet!")
v = np.fromfile(f, dt_data, x * y * z)
finally:
f.close()
if keepnumpy:
volume = np.array(v.reshape((x, y, z), order=order),
dtype='float32').copy() # fortran-order array
else: # if the input data is not the default type, convert
if not deviceID == None:
id = int(deviceID.split(":")[1])
xp.cuda.Device(id).use()
volume = xp.array(v.reshape((x, y, z), order=order),
dtype='float32').copy() # fortran-order array
return volume
def add_transformed_particle_to_sum(particle,
sum,
rotation=None,
rotation_order='rzxz',
translation=None,
scale=None):
from pytom.voltools import transform
rot = transform(particle,
rotation=rotation,
translation=translation,
scale=scale,
rotation_order=rotation_order,
device=device,
interpolation='filt_bspline')
sum += xp.array(rot)
return sum
def read_mrc(filename, order='F', keepnumpy=False, deviceID=None):
import numpy as np
f = open(filename, 'r')
try:
dt_header = np.dtype('int32')
header = np.fromfile(f, dt_header, 256)
x = header[0]
y = header[1]
z = header[2]
if header[23]:
extended_header = np.fromfile(f, np.float32, header[23] // 4)
# read the data type
dt = header[3]
default_type = False
if dt == 0: # byte
dt_data = np.dtype('int8')
elif dt == 1: # short
dt_data = np.dtype('int16')
elif dt == 2: # long
dt_data = np.dtype('float32') # little-endian, float32
default_type = True
elif dt == 4: # float32
dt_data = np.dtype('complex32')
elif dt == 6: # float complex
dt_data = np.dtype('uint16')
else:
raise Exception("Data type not supported yet!")
v = np.fromfile(f, dt_data, x * y * z)
finally:
f.close()
if keepnumpy:
volume = np.array(v.reshape((x, y, z), order=order),
dtype='float32').copy() # fortran-order array
else:
if not deviceID == None:
id = int(deviceID.split(":")[1])
xp.cuda.Device(id).use()
volume = xp.array(v.reshape((x, y, z), order=order),
dtype=xp.float32).copy() # fortran-order array
return volume
def find_sub_pixel_max_value_voltools(inp, k=.1, ignore_border=50):
"""
To find the highest point in a 2D array, with subpixel accuracy based on 1D spline interpolation.
The algorithm is based on a matlab script "tom_peak.m"
@param inp: A 2D numpy array containing the data points.
@type inp: numpy array 2D
@param k: The smoothing factor used in the spline interpolation, must be 1 <= k <= 5.
@type k: int
@return: A list of all points of maximal value in the structure of tuples with the x position, the y position and
the value.
@returntype: list
"""
# assert len(ixp.shape) == 2
# assert isinstance(k, int) and 1 <= k <= 5
import cupy
import numpy as xp
from scipy.interpolate import InterpolatedUnivariateSpline
from pytom.voltools import transform
inp2 = inp[ignore_border:-ignore_border, ignore_border:-ignore_border]
out = cupy.array(inp, dtype=cupy.float32)
x, y = xp.array(xp.unravel_index(inp2.argmax(),
inp2.shape)) + ignore_border
out = cupy.expand_dims(out, 2)
zoomed = cupy.zeros_like(out, dtype=cupy.float32)
transform(out,
output=zoomed,
scale=(k, k, 1),
translation=[inp.shape[0] // 2 - x, inp.shape[1] // 2 - y, 0],
device='gpu:0',
interpolation='filt_bspline')
z = zoomed.get().squeeze()
x2, y2 = xp.unravel_index(z.argmax(), z.shape)
return [
x + (x2 - z.shape[0] // 2) * k - z.shape[0] // 2,
y + (y2 - z.shape[1] // 2) * k - z.shape[0] // 2, z, x, y, x2, y2
]
def FSC(volume1,
volume2,
numberBands=None,
mask=None,
verbose=False,
filename=None,
num_procs=1):
"""
FSC - Calculates the Fourier Shell Correlation for two volumes
@param volume1: volume one
@type volume1: L{pytom_volume.vol}
@param volume2: volume two
@type volume2: L{pytom_volume.vol}
@param numberBands: number of shells for FSC
@type numberBands: int
@param filename: write FSC to ascii file if specified
@type filename: string
@return: Returns a list of cc values
@author: Thomas Hrabe
@rtype: list[floats]
"""
from pytom.tompy.correlation import bandCC
from pytom.basic.structures import Mask
from pytom.tompy.io import read, write
from pytom.tompy.tools import volumesSameSize
import time
t = time.time()
if not volumesSameSize(volume1, volume2):
raise RuntimeError('Volumes must have the same size!')
numberBands = volume1.shape[0] // 2 if numberBands is None else numberBands
if not mask is None:
if mask.__class__ == xp.array([0]).__class__:
volume1 = volume1 * mask
volume2 = volume2 * mask
elif mask.__class__ == Mask:
mask = mask.getVolume()
volume1 = volume1 * mask
volume2 = volume2 * mask
elif mask.__class__ == str:
mask = read(mask)
volume1 = volume1 * mask
volume2 = volume2 * mask
else:
raise RuntimeError(
'FSC: Mask must be a volume OR a Mask object OR a string path to a mask!'
)
fscResult = []
band = [-1, -1]
increment = int(volume1.shape[0] / 2 * 1 / numberBands)
import time
fvolume1 = xp.fft.fftn(volume1)
fvolume2 = xp.fft.fftn(volume2)
for n, i in enumerate(range(0, volume1.shape[0] // 2, increment)):
band[0] = i
band[1] = i + increment
if verbose:
print('Band : ', band)
res = bandCC(fvolume1, fvolume2, band, verbose)
if i == 0 and increment == 1:
#force a 1 for correlation of the zero frequency
res = 1
if verbose:
print('Correlation ', res)
fscResult.append(res)
if filename:
f = open(filename, 'w')
for item in fscResult:
f.write("%s\n" % item)
f.close()
return fscResult
def determineResolution(fsc,
resolutionCriterion,
verbose=False,
randomizedFSC=None,
gpu=False):
"""
determineResolution: Determines frequency and band where correlation drops below the resolutionCriterion. Uses linear interpolation between two positions
@param fsc: The fsc list determined by L{pytom.basic.correlation.FSC}
@param resolutionCriterion: A value between 0 and 1
@return: [resolution,interpolatedBand,numberBands]
@author: Thomas Hrabe
@todo: Add test!
"""
fsc = xp.array(fsc)
numberBands = len(fsc)
band = numberBands
if randomizedFSC is None:
randomizedFSC = xp.ones_like(fsc) * (fsc.min() - 0.1)
for i in range(numberBands):
if fsc[i] < resolutionCriterion and fsc[i] > randomizedFSC[i]:
band = i - 1 #select the band that is still larger than criterion
break
if verbose:
print('Band detected at ', band)
if band == -1:
raise RuntimeError(
"Please check your resolution criterion or you FSC!")
elif band < numberBands:
fsc1 = fsc[band]
fsc2 = fsc[band + 1]
rfsc1 = randomizedFSC[band]
rfsc2 = randomizedFSC[band + 1]
try:
if fsc2 < rfsc2:
interpolatedBand = (fsc1 - rfsc1) / (rfsc2 - rfsc1 + fsc1 -
fsc2)
pass
else:
interpolatedBand = (resolutionCriterion - fsc1) / (fsc2 -
fsc1) + band
except ZeroDivisionError:
interpolatedBand = band
else:
interpolatedBand = band
if verbose:
print('Band interpolated to ', interpolatedBand)
resolution = (interpolatedBand + 1) / float(numberBands)
if resolution < 0:
resolution = 1
interpolatedBand = numberBands
print(
'Warning: PyTom determined a resolution < 0 for your data. Please check "mass" in data is positive or negative for all cubes.'
)
print('Warning: Setting resolution to 1 and ', interpolatedBand)
print('')
return [resolution, interpolatedBand, numberBands]
def subPixelMax3D(volume,
k=.01,
ignore_border=50,
interpolation='filt_bspline',
plan=None,
profile=True,
num_threads=1024,
zoomed=None,
fast_sum=None,
max_id=None):
"""
Function to find the highest point in a 3D array, with subpixel accuracy using cubic spline interpolation.
@param inp: A 3D numpy/cupy array containing the data points.
@type inp: numpy/cupy array 3D
@param k: The interpolation factor used in the spline interpolation, k < 1 is zoomed in, k>1 zoom out.
@type k: float
@return: A list of maximal value in the interpolated volume and a list of x position, the y position and
the value.
@returntype: list
"""
from pytom.voltools import transform
from pytom.tompy.io import write
ox, oy, oz = volume.shape
ib = ignore_border
cropped_volume = volume[ib:ox - ib, ib:oy - ib,
ib:oz - ib].astype(xp.float32)
if profile:
stream = xp.cuda.Stream.null
t_start = stream.record()
# x,y,z = xp.array(maxIndex(cropped_volume)) + ignore_border
x, y, z = xp.array(
xp.unravel_index(cropped_volume.argmax(),
cropped_volume.shape)) + ignore_border
dx, dy, dz = volume.shape
translation = [dx // 2 - x, dy // 2 - y, dz // 2 - z]
if profile:
t_end = stream.record()
t_end.synchronize()
time_took = xp.cuda.get_elapsed_time(t_start, t_end)
print(f'initial find max time: \t{time_took:.3f}ms')
t_start = stream.record()
b = border = max(0, int(volume.shape[0] // 2 - 4 / k))
zx, zy, zz = volume.shape
out = volume[b:zx - b, b:zy - b, b:zz - b]
transform(out,
output=zoomed,
scale=(k, k, k),
device=device,
translation=translation,
interpolation=interpolation)
if profile:
t_end = stream.record()
t_end.synchronize()
time_took = xp.cuda.get_elapsed_time(t_start, t_end)
print(f'transform finished in \t{time_took:.3f}ms')
t_start = stream.record()
nblocks = int(xp.ceil(zoomed.size / num_threads / 2))
argmax((
nblocks,
1,
), (num_threads, 1, 1), (zoomed, fast_sum, max_id, zoomed.size),
shared_mem=8 * num_threads)
x2, y2, z2 = xp.unravel_index(max_id[fast_sum.argmax()], zoomed.shape)
peakValue = zoomed[x2][y2][z2]
peakShift = [
x + (x2 - zoomed.shape[0] // 2) * k - volume.shape[0] // 2,
y + (y2 - zoomed.shape[1] // 2) * k - volume.shape[1] // 2,
z + (z2 - zoomed.shape[2] // 2) * k - volume.shape[2] // 2
]
if profile:
t_end = stream.record()
t_end.synchronize()
time_took = xp.cuda.get_elapsed_time(t_start, t_end)
print(f'argmax finished in \t{time_took:.3f}ms')
t_start = stream.record()
print()
return [peakValue, peakShift]
def averageGPU(particleList,
averageName,
showProgressBar=False,
verbose=False,
createInfoVolumes=False,
weighting=False,
norm=False,
gpuId=None,
profile=True):
"""
average : Creates new average from a particleList
@param particleList: The particles
@param averageName: Filename of new average
@param verbose: Prints particle information. Disabled by default.
@param createInfoVolumes: Create info data (wedge sum, inverted density) too? False by default.
@param weighting: apply weighting to each average according to its correlation score
@param norm: apply normalization for each particle
@return: A new Reference object
@rtype: L{pytom.basic.structures.Reference}
@author: Thomas Hrabe
@change: limit for wedgeSum set to 1% or particles to avoid division by small numbers - FF
"""
import time
from pytom.tompy.io import read, write, read_size
from pytom.tompy.filter import bandpass as lowpassFilter, rotateWeighting, applyFourierFilter, applyFourierFilterFull, create_wedge
from pytom.voltools import transform, StaticVolume
from pytom.basic.structures import Reference
from pytom.tompy.normalise import mean0std1
from pytom.tompy.tools import volumesSameSize, invert_WedgeSum, create_sphere
from pytom.tompy.transform import fourier_full2reduced, fourier_reduced2full
from cupyx.scipy.fftpack.fft import fftn as fftnP
from cupyx.scipy.fftpack.fft import ifftn as ifftnP
from cupyx.scipy.fftpack.fft import get_fft_plan
from pytom.tools.ProgressBar import FixedProgBar
from multiprocessing import RawArray
import numpy as np
import cupy as xp
if not gpuId is None:
device = f'gpu:{gpuId}'
xp.cuda.Device(gpuId).use()
else:
print(gpuId)
raise Exception('Running gpu code on non-gpu device')
print(device)
cstream = xp.cuda.Stream()
if profile:
stream = xp.cuda.Stream.null
t_start = stream.record()
# from pytom.tools.ProgressBar import FixedProgBar
from math import exp
import os
if len(particleList) == 0:
raise RuntimeError('The particle list is empty. Aborting!')
if showProgressBar:
progressBar = FixedProgBar(0, len(particleList), 'Particles averaged ')
progressBar.update(0)
numberAlignedParticles = 0
# pre-check that scores != 0
if weighting:
wsum = 0.
for particleObject in particleList:
wsum += particleObject.getScore().getValue()
if wsum < 0.00001:
weighting = False
print("Warning: all scores have been zero - weighting not applied")
import time
sx, sy, sz = read_size(particleList[0].getFilename())
wedgeInfo = particleList[0].getWedge().convert2numpy()
print('angle: ', wedgeInfo.getWedgeAngle())
wedgeZero = xp.fft.fftshift(
xp.array(wedgeInfo.returnWedgeVolume(sx, sy, sz, True).get(),
dtype=xp.float32))
# wedgeZeroReduced = fourier_full2reduced(wedgeZero)
wedge = xp.zeros_like(wedgeZero, dtype=xp.float32)
wedgeSum = xp.zeros_like(wedge, dtype=xp.float32)
print('init texture')
wedgeText = StaticVolume(xp.fft.fftshift(wedgeZero),
device=device,
interpolation='filt_bspline')
newParticle = xp.zeros((sx, sy, sz), dtype=xp.float32)
centerX = sx // 2
centerY = sy // 2
centerZ = sz // 2
result = xp.zeros((sx, sy, sz), dtype=xp.float32)
fftplan = get_fft_plan(wedge.astype(xp.complex64))
n = 0
total = len(particleList)
# total = int(np.floor((11*1024**3 - mempool.total_bytes())/(sx*sy*sz*4)))
# total = 128
#
#
# particlesNP = np.zeros((total, sx, sy, sz),dtype=np.float32)
# particles = []
# mask = create_sphere([sx,sy,sz], sx//2-6, 2)
# raw = RawArray('f', int(particlesNP.size))
# shared_array = np.ctypeslib.as_array(raw)
# shared_array[:] = particlesNP.flatten()
# procs = allocateProcess(particleList, shared_array, n, total, wedgeZero.size)
# del particlesNP
if profile:
t_end = stream.record()
t_end.synchronize()
time_took = xp.cuda.get_elapsed_time(t_start, t_end)
print(f'startup time {n:5d}: \t{time_took:.3f}ms')
t_start = stream.record()
for particleObject in particleList:
rotation = particleObject.getRotation()
rotinvert = rotation.invert()
shiftV = particleObject.getShift()
# if n % total == 0:
# while len(procs):
# procs =[proc for proc in procs if proc.is_alive()]
# time.sleep(0.1)
# print(0.1)
# # del particles
# # xp._default_memory_pool.free_all_blocks()
# # pinned_mempool.free_all_blocks()
# particles = xp.array(shared_array.reshape(total, sx, sy, sz), dtype=xp.float32)
# procs = allocateProcess(particleList, shared_array, n, total, size=wedgeZero.size)
# #pinned_mempool.free_all_blocks()
# #print(mempool.total_bytes()/1024**3)
particle = read(particleObject.getFilename(), deviceID=device)
#particle = particles[n%total]
if norm: # normalize the particle
mean0std1(particle) # happen inplace
# apply its wedge to
#particle = applyFourierFilter(particle, wedgeZeroReduced)
#particle = (xp.fft.ifftn( xp.fft.fftn(particle) * wedgeZero)).real
particle = (ifftnP(fftnP(particle, plan=fftplan) * wedgeZero,
plan=fftplan)).real
### create spectral wedge weighting
wedge *= 0
wedgeText.transform(
rotation=[rotinvert[0], rotinvert[2], rotinvert[1]],
rotation_order='rzxz',
output=wedge)
#wedge = xp.fft.fftshift(fourier_reduced2full(create_wedge(30, 30, 21, 42, 42, 42, rotation=[rotinvert[0],rotinvert[2], rotinvert[1]])))
# if analytWedge:
# # > analytical buggy version
# wedge = wedgeInfo.returnWedgeVolume(sx, sy, sz, True, rotinvert)
# else:
# # > FF: interpol bugfix
# wedge = rotateWeighting(weighting=wedgeInfo.returnWedgeVolume(sx, sy, sz, True), rotation=[rotinvert[0], rotinvert[2], rotinvert[1]])
# # < FF
# # > TH bugfix
# # wedgeVolume = wedgeInfo.returnWedgeVolume(wedgeSizeX=sizeX, wedgeSizeY=sizeY, wedgeSizeZ=sizeZ,
# # humanUnderstandable=True, rotation=rotinvert)
# # wedge = rotate(volume=wedgeVolume, rotation=rotinvert, imethod='linear')
# # < TH
### shift and rotate particle
newParticle *= 0
transform(particle,
output=newParticle,
rotation=[-rotation[1], -rotation[2], -rotation[0]],
center=[centerX, centerY, centerZ],
translation=[-shiftV[0], -shiftV[1], -shiftV[2]],
device=device,
interpolation='filt_bspline',
rotation_order='rzxz')
#write(f'trash/GPU_{n}.em', newParticle)
# print(rotation.toVector())
# break
result += newParticle
wedgeSum += xp.fft.fftshift(wedge)
# if showProgressBar:
# numberAlignedParticles = numberAlignedParticles + 1
# progressBar.update(numberAlignedParticles)
if n % total == 0:
if profile:
t_end = stream.record()
t_end.synchronize()
time_took = xp.cuda.get_elapsed_time(t_start, t_end)
print(f'total time {n:5d}: \t{time_took:.3f}ms')
t_start = stream.record()
cstream.synchronize()
n += 1
print('averaged particles')
###apply spectral weighting to sum
result = lowpassFilter(result, high=sx / 2 - 1, sigma=0)
# if createInfoVolumes:
write(averageName[:len(averageName) - 3] + '-PreWedge.em', result)
write(averageName[:len(averageName) - 3] + '-WedgeSumUnscaled.em',
fourier_full2reduced(wedgeSum))
wedgeSumINV = invert_WedgeSum(wedgeSum,
r_max=sx // 2 - 2.,
lowlimit=.05 * len(particleList),
lowval=.05 * len(particleList))
wedgeSumINV = wedgeSumINV
#print(wedgeSum.mean(), wedgeSum.std())
if createInfoVolumes:
write(averageName[:len(averageName) - 3] + '-WedgeSumInverted.em',
xp.fft.fftshift(wedgeSumINV))
result = applyFourierFilterFull(result, xp.fft.fftshift(wedgeSumINV))
# do a low pass filter
result = lowpassFilter(result, sx / 2 - 2, (sx / 2 - 1) / 10.)[0]
write(averageName, result)
if createInfoVolumes:
resultINV = result * -1
# write sign inverted result to disk (good for chimera viewing ... )
write(averageName[:len(averageName) - 3] + '-INV.em', resultINV)
newReference = Reference(averageName, particleList)
return newReference
def euclidian_distance(pos1, pos2):
return xp.linalg.norm(xp.array(pos1) - xp.array(pos2))
def volumesSameSize(v0, v1):
if len(v0.shape) != len(v1.shape):
return False
return xp.array([
abs(v0.shape[pos] - v1.shape[pos]) == 0 for pos in range(len(v0.shape))
]).sum() == len(v0.shape)
def alignImagesUsingAlignmentResultFile(alignmentResultsFile,
weighting=None,
lowpassFilter=0.9,
binning=1,
circleFilter=False):
import os
from pytom.basic.files import read as readCVol
from pytom_numpy import vol2npy, npy2vol
from pytom.gui.guiFunctions import fmtAR, headerAlignmentResults, datatype, datatypeAR, loadstar
from pytom.reconstruction.reconstructionStructures import Projection, ProjectionList
from pytom.tompy.io import read, write, read_size
from pytom.tompy.tools import taper_edges, create_circle
from pytom.tompy.filter import circle_filter, ramp_filter, exact_filter
import pytom.voltools as vt
from pytom.gpu.initialize import xp, device
print("Create aligned images from alignResults.txt")
alignmentResults = loadstar(alignmentResultsFile, dtype=datatypeAR)
imageList = alignmentResults['FileName']
tilt_angles = alignmentResults['TiltAngle']
imdim = read_size(imageList[0], 'x')
if binning > 1:
imdim = int(float(imdim) / float(binning) + .5)
else:
imdim = imdim
sliceWidth = imdim
if (weighting != None) and (float(weighting) < -0.001):
weightSlice = xp.fft.fftshift(ramp_filter(imdim, imdim))
if circleFilter:
circleFilterRadius = imdim // 2
circleSlice = xp.fft.fftshift(
circle_filter(imdim, imdim, circleFilterRadius))
else:
circleSlice = xp.ones((imdim, imdim))
# design lowpass filter
if lowpassFilter:
if lowpassFilter > 1.:
lowpassFilter = 1.
print("Warning: lowpassFilter > 1 - set to 1 (=Nyquist)")
# weighting filter: arguments: (()dimx, dimy), cutoff radius, sigma
lpf = xp.fft.fftshift(
create_circle((imdim, imdim),
lowpassFilter * (imdim // 2),
sigma=0.4 * lowpassFilter * (imdim // 2)))
projectionList = ProjectionList()
for n, image in enumerate(imageList):
atx = alignmentResults['AlignmentTransX'][n] / binning
aty = alignmentResults['AlignmentTransY'][n] / binning
rot = alignmentResults['InPlaneRotation'][n]
mag = 1 / alignmentResults['Magnification'][n]
projection = Projection(imageList[n],
tiltAngle=tilt_angles[n],
alignmentTransX=atx,
alignmentTransY=aty,
alignmentRotation=rot,
alignmentMagnification=mag)
projectionList.append(projection)
stack = xp.zeros((imdim, imdim, len(imageList)), dtype=xp.float32)
phiStack = xp.zeros((1, 1, len(imageList)), dtype=xp.float32)
thetaStack = xp.zeros((1, 1, len(imageList)), dtype=xp.float32)
offsetStack = xp.zeros((1, 2, len(imageList)), dtype=xp.float32)
for (ii, projection) in enumerate(projectionList):
print(f'Align {projection._filename}')
image = read(str(projection._filename))[::binning, ::binning].squeeze()
if lowpassFilter:
image = xp.abs((xp.fft.ifftn(xp.fft.fftn(image) * lpf)))
tiltAngle = projection._tiltAngle
# normalize to contrast - subtract mean and norm to mean
immean = image.mean()
image = (image - immean) / immean
# smoothen borders to prevent high contrast oscillations
image = taper_edges(image, imdim // 30)[0]
# transform projection according to tilt alignment
transX = projection._alignmentTransX / binning
transY = projection._alignmentTransY / binning
rot = float(projection._alignmentRotation)
mag = float(projection._alignmentMagnification)
inputImage = xp.expand_dims(image, 2).copy()
outputImage = xp.zeros_like(inputImage, dtype=xp.float32)
vt.transform(inputImage.astype(xp.float32),
rotation=[0, 0, rot],
rotation_order='rxyz',
output=outputImage,
device=device,
translation=[transX, transY, 0],
scale=[mag, mag, 1],
interpolation='filt_bspline')
image = outputImage.squeeze()
# smoothen once more to avoid edges
image = taper_edges(image, imdim // 30)[0]
# analytical weighting
if (weighting != None) and (weighting < 0):
# image = (ifft(complexRealMult(fft(image), w_func)) / (image.sizeX() * image.sizeY() * image.sizeZ()))
image = xp.fft.ifftn(
xp.fft.fftn(image) * weightSlice * circleSlice)
elif (weighting != None) and (weighting > 0):
weightSlice = xp.fft.fftshift(
exact_filter(tilt_angles, tiltAngle, imdim, imdim, sliceWidth))
image = xp.fft.ifftn(
xp.fft.fftn(image) * weightSlice * circleSlice)
thetaStack[0, 0, ii] = int(round(projection.getTiltAngle()))
offsetStack[0, :, ii] = xp.array([
int(round(projection.getOffsetX())),
int(round(projection.getOffsetY()))
])
stack[:, :, ii] = image
arrays = []
for fname, arr in (('stack.mrc', stack), ('offsetStack.mrc', offsetStack),
('thetaStack.mrc', thetaStack), ('phiStack.mrc',
phiStack)):
if 'gpu' in device:
arr = arr.get()
import numpy as np
res = npy2vol(np.array(arr, dtype='float32', order='F'), 3)
arrays.append(res)
#
# write('stack.mrc', stack)
# stack = readCVol('stack.mrc')
# write('offsetstack.mrc', offsetStack)
# offsetStack = readCVol('offsetstack.mrc')
# write('thetastack.mrc', thetaStack)
# thetaStack = readCVol('thetastack.mrc')
# write('phistack.mrc', phiStack)
# phiStack = readCVol('phistack.mrc')
#
# os.remove('stack.mrc')
# os.remove('offsetstack.mrc')
# os.remove('thetastack.mrc')
# os.remove('psistack.mrc')
return arrays
def profile2FourierVol(profile, dim=None, reduced=False):
"""
create Volume from 1d radial profile, e.g., to modulate signal with \
specific function such as CTF or FSC. Simple linear interpolation is used\
for sampling.
@param profile: profile
@type profile: 1-d L{pytom_volume.vol} or 1-d python array
@param dim: dimension of (cubic) output
@type dim: L{int}
@param reduced: If true reduced Fourier representation (N/2+1, N, N) is generated.
@type reduced: L{bool}
@return: 3-dim complex volume with spherically symmetrical profile
@rtype: L{pytom_volume.vol}
@author: FF
"""
if dim is None:
try:
dim = [
2 * profile.shape[0],
] * 3
except:
dim = [
2 * len(profile),
] * 3
is3D = (len(dim) == 3)
nx, ny = dim[:2]
if reduced:
if is3D:
nz = int(dim[2] // 2) + 1
else:
ny = int(ny // 2) + 1
else:
if is3D:
nz = dim[2]
try:
r_max = profile.shape[0] - 1
except:
r_max = len(profile) - 1
if len(dim) == 3:
if reduced:
X, Y, Z = xp.meshgrid(xp.arange(-nx // 2, nx // 2 + nx % 2),
xp.arange(-ny // 2, ny // 2 + ny % 2),
xp.arange(0, nz))
else:
X, Y, Z = xp.meshgrid(xp.arange(-nx // 2, nx // 2 + nx % 2),
xp.arange(-ny // 2, ny // 2 + ny % 2),
xp.arange(-nz // 2, nz // 2 + nz % 2))
R = xp.sqrt(X**2 + Y**2 + Z**2)
else:
if reduced:
X, Y = xp.meshgrid(xp.arange(-nx // 2, ny // 2 + ny % 2),
xp.arange(0, ny))
else:
X, Y = xp.meshgrid(xp.arange(-nx // 2, nx // 2 + nx % 2),
xp.arange(-ny // 2, ny // 2 + ny % 2))
R = xp.sqrt(X**2 + Y**2)
IR = xp.floor(R).astype(xp.int64)
valIR_l1 = IR.copy()
valIR_l2 = valIR_l1 + 1
val_l1, val_l2 = xp.zeros_like(X, dtype=xp.float64), xp.zeros_like(
X, dtype=xp.float64)
l1 = R - IR.astype(xp.float32)
l2 = 1 - l1
try:
profile = xp.array(profile)
except:
import numpy
profile = xp.array(numpy.array(profile))
for n in xp.arange(r_max):
val_l1[valIR_l1 == n] = profile[n]
val_l2[valIR_l2 == n + 1] = profile[n + 1]
val_l1[IR == r_max] = profile[n + 1]
val_l2[IR == r_max] = profile[n + 1]
val_l1[R > r_max] = 0
val_l2[R > r_max] = 0
fkernel = l2 * val_l1 + l1 * val_l2
if reduced:
fkernel = xp.fft.fftshift(fkernel, axes=(0, 1))
else:
fkernel = xp.fft.fftshift(fkernel)
return fkernel
def polish_particles(particle_list_filename,
projection_directory,
averaged_subtomogram,
binning,
offset,
projections,
tilt_angles,
fsc_path='',
peak_border=75,
outputDirectory='./',
create_graphics=False,
number_of_particles=0,
verbose=False,
gpuID=-1):
"""
To polish a particle list based on (an) initial subtomogram(s).
:param particle_list_filename: the filename of the particlelist
:type particle_list_filename: str
:param projection_directory: the directory of the projections
:type projection_directory: str
:param averaged_subtomogram: to give a path to an averaged subtomogram to be used instead of subtomograms of all
particles separately
:type averaged_subtomogram: str
:param binning: the binning factor used
:type binning: int
:param offset: the offset used (x, y, z)
:type offset: list(int, int, int)
:param projections: a list with filenames of projections
:type projections: list(str)
:param tilt_angles: the list of tiltangles used
:type tilt_angles: list(int)
:param create_graphics: to create plots of major parts of the algorithm, mainly used for debugging
and initial creation
:type create_graphics: bool
:param number_of_particles: to use a subset of the particles for the particle polishing
:type number_of_particles: int
:param skip_alignment: skips the alignment phase, does not do particle polishing
:type skip_alignment: bool
:return: nothing, it writes everything to disk
:returntype: void
"""
assert number_of_particles == -1 or number_of_particles > 0
assert binning > 0
assert vol_size > 0
assert vol_size % 2 == 0
assert isinstance(projections, list)
assert isinstance(vol_size, int)
assert isinstance(binning, int)
assert isinstance(offset, list) and len(offset) == 3
assert isinstance(offset[0], int) and isinstance(
offset[1], int) and isinstance(offset[2], int)
assert isinstance(tilt_angles, list)
assert isinstance(particle_list_filename, str)
assert isinstance(projection_directory, str)
assert isinstance(create_graphics, bool)
assert isinstance(averaged_subtomogram, str)
assert isinstance(number_of_particles, int)
assert isinstance(skip_alignment, bool)
import os, time
from pytom.tompy.io import read_size, read
from pytom.gui.guiFunctions import fmtLAR, headerLocalAlignmentResults, LOCAL_ALIGNMENT_RESULTS
import pytom.voltools as vt
# load particle list
from pytom.basic.structures import ParticleList
particlelist = ParticleList()
particlelist.fromXMLFile(particle_list_filename)
particle_list_name = os.path.splitext(
os.path.basename(str(particle_list_filename)))[0]
if number_of_particles > 0:
particlelist = particlelist[:number_of_particles]
if verbose:
print(len(particlelist))
print("{:s}> Creating the input array".format(gettime()))
dimz = read_size(particlelist[0].getPickPosition().getOriginFilename(),
'z') * binning
vol_size = 200
input_to_processes = []
data = {}
for projectioname in projections:
data[projectioname] = xp.array(read(projectioname))
#
template1 = read(averaged_subtomogram, order='F')
#template1 = mrcfile.open(averaged_subtomogram,permissive=True).data.copy().T.copy()
template = vt.StaticVolume(template1,
interpolation='filt_bspline',
device=device)
# output = []
results_file = os.path.join(outputDirectory,
f"resultsPolish_{particle_list_name}.txt")
results = []
for particle_number, particle in enumerate(particlelist):
rot = (particle.getRotation().getZ1(), particle.getRotation().getX(),
particle.getRotation().getZ2())
# loop over tiltrange, take patch and cross correlate with reprojected subtomogram
for img, ang in zip(projections, tilt_angles):
pick_position = particle.getPickPosition().toVector()
result = run_single_tilt_angle(template, ang, offset, vol_size,
pick_position, rot,
particle.getFilename(),
particle_number, binning, data,
create_graphics, fsc_path, dimz,
peak_border, img,
averaged_subtomogram)
results.append(tuple(result))
try:
np.savetxt(results_file,
np.array(results, dtype=LOCAL_ALIGNMENT_RESULTS),
fmt=fmtLAR,
header=headerLocalAlignmentResults)
except Exception as e:
print(e)
for res in results:
print('{:7d} {:15.3f} {:15.3f} {:15.3f} {:15.3f} {:15.10f} {:s}'.
format(*res))
break
if verbose: print("{:s}> Ran the processes".format(gettime()))