Source code for xmipp3.protocols.protocol_angular_alignment_zernike3d

# **************************************************************************
# *
# * Authors:     Amaya Jimenez Moreno (
# *              David Herreros Calero (
# *
# * Unidad de  Bioinformatica of Centro Nacional de Biotecnologia , CSIC
# *
# * This program is free software; you can redistribute it and/or modify
# * it under the terms of the GNU General Public License as published by
# * the Free Software Foundation; either version 2 of the License, or
# * (at your option) any later version.
# *
# * This program is distributed in the hope that it will be useful,
# * but WITHOUT ANY WARRANTY; without even the implied warranty of
# * GNU General Public License for more details.
# *
# * You should have received a copy of the GNU General Public License
# * along with this program; if not, write to the Free Software
# * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
# * 02111-1307  USA
# *
# *  All comments concerning this program package may be sent to the
# *  e-mail address ''
# *
# **************************************************************************

import numpy as np

import pyworkflow.protocol.params as params
from pyworkflow.object import Integer
from pyworkflow.utils.path import moveFile
from pyworkflow import VERSION_2_0

from pwem.protocols import ProtAnalysis3D
import pwem.emlib.metadata as md
from pwem.emlib.image import ImageHandler
from pwem.constants import ALIGN_PROJ

from xmipp3.convert import (writeSetOfParticles, createItemMatrix,
from xmipp3.base import writeInfoField, readInfoField

[docs]class XmippProtAngularAlignmentZernike3D(ProtAnalysis3D): """ Protocol for flexible angular alignment based on Zernike3D basis. """ _label = 'angular align - Zernike3D' _lastUpdateVersion = VERSION_2_0 # --------------------------- DEFINE param functions -------------------------------------------- def _defineParams(self, form): form.addSection(label='Input') form.addHidden(params.USE_GPU, params.BooleanParam, default=True, label="Use GPU for execution", help="This protocol has both CPU and GPU implementation.\ Select the one you want to use.") form.addHidden(params.GPU_LIST, params.StringParam, default='0', expertLevel=params.LEVEL_ADVANCED, label="Choose GPU IDs", help="Add a list of GPU devices that can be used") form.addParam('inputParticles', params.PointerParam, label="Input particles", pointerClass='SetOfParticles') form.addParam('inputVolume', params.PointerParam, label="Input volume", pointerClass='Volume') form.addParam('inputVolumeMask', params.PointerParam, label="Input volume mask", pointerClass='VolumeMask', allowsNull=True) form.addParam('targetResolution', params.FloatParam, label="Target resolution (A)", default=8.0, help="In Angstroms, the images and the volume are rescaled so that this resolution is at " "2/3 of the Fourier spectrum.") form.addParam('l1', params.IntParam, default=3, label='Zernike Degree', expertLevel=params.LEVEL_ADVANCED, help='Degree Zernike Polynomials of the deformation=1,2,3,...') form.addParam('l2', params.IntParam, default=2, label='Harmonical Degree', expertLevel=params.LEVEL_ADVANCED, help='Degree Spherical Harmonics of the deformation=1,2,3,...') form.addParam('maxShift', params.FloatParam, default=-1, label='Maximum shift (px)', expertLevel=params.LEVEL_ADVANCED, help='Maximum shift allowed in pixels') form.addParam('maxAngular', params.FloatParam, default=5, label='Maximum angular change (degrees)', expertLevel=params.LEVEL_ADVANCED, help='Maximum angular change allowed (in degrees)') form.addParam('maxResolution', params.FloatParam, default=4.0, label='Maximum resolution (A)', expertLevel=params.LEVEL_ADVANCED, help='Maximum resolution (A)') form.addParam('regularization', params.FloatParam, default=0.005, label='Regularization', expertLevel=params.LEVEL_ADVANCED, help='Penalization to deformations (higher values penalize more the deformation).') form.addParam('ignoreCTF', params.BooleanParam, default=False, label='Ignore CTF?', expertLevel=params.LEVEL_ADVANCED, help="If true, volume projection won't be subjected to CTF corrections") form.addParam('optimizeAlignment', params.BooleanParam, default=True, label='Optimize alignment?', expertLevel=params.LEVEL_ADVANCED) form.addParallelSection(threads=1, mpi=8) def _createFilenameTemplates(self): """ Centralize how files are called """ myDict = { 'imgsFn': self._getExtraPath('input_particles.xmd'), 'fnVol': self._getExtraPath('input_volume.vol'), 'fnVolMask': self._getExtraPath('input_volume_mask.vol'), 'fnOut': self._getExtraPath('output_particles.xmd'), 'fnOutDir': self._getExtraPath() } self._updateFilenamesDict(myDict) # --------------------------- INSERT steps functions -------------------------------------------- def _insertAllSteps(self): self._createFilenameTemplates() self._insertFunctionStep("convertStep") self._insertFunctionStep("alignmentStep") self._insertFunctionStep("createOutputStep") # --------------------------- STEPS functions ---------------------------------------------------
[docs] def convertStep(self): imgsFn = self._getFileName('imgsFn') fnVol = self._getFileName('fnVol') fnVolMask = self._getFileName('fnVolMask') inputParticles = self.inputParticles.get() writeSetOfParticles(inputParticles, imgsFn) Xdim = inputParticles.getXDim() self.Ts = inputParticles.getSamplingRate() newTs = self.targetResolution.get() * 1.0 /3.0 self.newTs = max(self.Ts, newTs) self.newXdim = int(Xdim * self.Ts / newTs) writeInfoField(self._getExtraPath(), "sampling", md.MDL_SAMPLINGRATE, newTs) writeInfoField(self._getExtraPath(), "size", md.MDL_XSIZE, self.newXdim) if self.newXdim != Xdim: self.runJob("xmipp_image_resize", "-i %s -o %s --save_metadata_stack %s --fourier %d" % (imgsFn, self._getExtraPath('scaled_particles.stk'), self._getExtraPath('scaled_particles.xmd'), self.newXdim), numberOfMpi=1) moveFile(self._getExtraPath('scaled_particles.xmd'), imgsFn) ih = ImageHandler() ih.convert(self.inputVolume.get(), fnVol) Xdim = self.inputVolume.get().getDim()[0] if Xdim != self.newXdim: self.runJob("xmipp_image_resize", "-i %s --dim %d " % (fnVol, self.newXdim), numberOfMpi=1) if self.inputVolumeMask.get(): ih.convert(self.inputVolumeMask.get(), fnVolMask) if Xdim != self.newXdim: self.runJob("xmipp_image_resize", "-i %s --dim %d --interp nearest" % (fnVolMask, self.newXdim), numberOfMpi=1)
[docs] def alignmentStep(self): imgsFn = self._getFileName('imgsFn') fnVol = self._getFileName('fnVol') fnOut = self._getFileName('fnOut') fnVolMask = self._getFileName('fnVolMask') fnOutDir = self._getFileName('fnOutDir') Ts = readInfoField(self._getExtraPath(), "sampling", md.MDL_SAMPLINGRATE) params = ' -i %s --ref %s -o %s --optimizeDeformation --optimizeDefocus ' \ '--l1 %d --l2 %d --max_shift %f --max_angular_change %f --sampling %f ' \ ' --max_resolution %f --odir %s --resume --regularization %f' %\ (imgsFn, fnVol, fnOut, self.l1.get(), self.l2.get(), self.maxShift, self.maxAngular, Ts, self.maxResolution, fnOutDir, self.regularization.get()) if self.optimizeAlignment.get(): params += ' --optimizeAlignment' if self.ignoreCTF.get(): params += ' --ignoreCTF' if self.inputParticles.get().isPhaseFlipped(): params += ' --phaseFlipped' if self.inputVolumeMask.get(): params += ' --mask %s' % fnVolMask if self.useGpu.get(): params += ' --device %d' % self.getGpuList()[0] program = 'xmipp_cuda_angular_sph_alignment' self.runJob(program, params) else: program = 'xmipp_angular_sph_alignment' self.runJob(program, params, numberOfMpi=self.numberOfMpi.get())
[docs] def createOutputStep(self): from sklearn.manifold import TSNE Xdim = self.inputParticles.get().getXDim() self.Ts = self.inputParticles.get().getSamplingRate() newTs = self.targetResolution.get() * 1.0 /3.0 self.newTs = max(self.Ts, newTs) self.newXdim = int(Xdim * self.Ts / newTs) fnOut = self._getFileName('fnOut') mdOut = md.MetaData(fnOut) coeffMatrix = np.vstack(mdOut.getColumnValues(md.MDL_SPH_COEFFICIENTS)) X_tsne_1d = TSNE(n_components=1).fit_transform(coeffMatrix) X_tsne_2d = TSNE(n_components=2).fit_transform(coeffMatrix) newMdOut = md.MetaData() i = 0 for row in md.iterRows(mdOut): newRow = row newRow.setValue(md.MDL_SPH_TSNE_COEFF1D, float(X_tsne_1d[i, 0])) newRow.setValue(md.MDL_SPH_TSNE_COEFF2D, [float(X_tsne_2d[i, 0]), float(X_tsne_2d[i, 1])]) if self.newTs != self.Ts: coeffs = mdOut.getValue(md.MDL_SPH_COEFFICIENTS, row.getObjId()) correctionFactor = self.inputVolume.get().getDim()[0] / self.newXdim coeffs = [correctionFactor*coeff for coeff in coeffs] newRow.setValue(md.MDL_SPH_COEFFICIENTS, coeffs) newRow.addToMd(newMdOut) i += 1 newMdOut.write(fnOut) inputSet = self.inputParticles.get() partSet = self._createSetOfParticles() partSet.copyInfo(inputSet) partSet.setAlignmentProj() partSet.copyItems(inputSet, updateItemCallback=self._updateParticle, itemDataIterator=md.iterRows(fnOut, sortByLabel=md.MDL_ITEM_ID)) partSet.L1 = Integer(self.l1.get()) partSet.L2 = Integer(self.l2.get()) partSet.Rmax = Integer(self.inputVolume.get().getDim()[0] / 2) self._defineOutputs(outputParticles=partSet) self._defineTransformRelation(self.inputParticles, partSet)
# --------------------------- UTILS functions -------------------------------------------- def _updateParticle(self, item, row): setXmippAttributes(item, row, md.MDL_ANGLE_ROT, md.MDL_ANGLE_TILT, md.MDL_ANGLE_PSI, md.MDL_SHIFT_X, md.MDL_SHIFT_Y, md.MDL_FLIP, md.MDL_SPH_DEFORMATION, md.MDL_SPH_COEFFICIENTS, md.MDL_SPH_TSNE_COEFF1D, md.MDL_SPH_TSNE_COEFF2D) createItemMatrix(item, row, align=ALIGN_PROJ)
[docs] def getInputParticles(self): return self.inputParticles.get()
# ----------------------- VALIDATE functions ----------------------------------------
[docs] def validate(self): """ Try to find errors on define params. """ errors = [] l1 = self.l1.get() l2 = self.l2.get() if (l1 - l2) < 0: errors.append('Zernike degree must be higher than ' 'SPH degree.') return errors