From a1048f881359c9e46ba97ab51126154035ad7247 Mon Sep 17 00:00:00 2001 From: Jose Miguel de la Rosa Trevin Date: Wed, 2 Oct 2019 12:24:58 +0200 Subject: [PATCH] Added more options to alignment assignment protocol --- .../em/protocol/protocol_alignment_assign.py | 99 ++++++++++++++----- 1 file changed, 73 insertions(+), 26 deletions(-) diff --git a/pyworkflow/em/protocol/protocol_alignment_assign.py b/pyworkflow/em/protocol/protocol_alignment_assign.py index dad40093f5..7f28dae09c 100644 --- a/pyworkflow/em/protocol/protocol_alignment_assign.py +++ b/pyworkflow/em/protocol/protocol_alignment_assign.py @@ -26,7 +26,7 @@ from pyworkflow.em import ALIGN_2D -from pyworkflow.protocol.params import PointerParam +from pyworkflow.protocol.params import PointerParam, EnumParam from protocol_2d import ProtAlign2D @@ -36,21 +36,55 @@ class ProtAlignmentAssign(ProtAlign2D): between the two sets and multiply by the right factor the shifts. The particles with the alignment can also be a subset of the other images """ + ASSIGN_ALL = 0 + ASSIGN_ONLY_SHIFTS = 1 + ASSIGN_ONLY_ANGLES = 2 + _label = 'assign alignment' def _defineParams(self, form): form.addSection(label='Input') - form.addParam('inputParticles', PointerParam, pointerClass='SetOfParticles', + form.addParam('inputParticles', PointerParam, + pointerClass='SetOfParticles', label='Input particles', - help='Select the particles that you want to update the new alignment.') - form.addParam('inputAlignment', PointerParam, pointerClass='SetOfParticles', + help='Select the particles that you want to update the ' + 'new alignment.') + form.addParam('inputAlignment', PointerParam, + pointerClass='SetOfParticles', label="Input alignments", - help='Select the particles with alignment to be apply to the other particles.') + help='Select the particles with alignment to be apply ' + 'to the other particles.') + + form.addParam('assignMode', EnumParam, default=self.ASSIGN_ALL, + choices=['all', 'only shifts', 'only angles'], + display=EnumParam.DISPLAY_HLIST, + label='Assign mode', + help="Select what information do you want to assign " + "from the input alignment. \n" + " **all**: Both angles and shifts will be used.\n" + " **only shifts**: Only the shifts values will " + "be considered.\n" + " **only angles**: Only angular assigment will " + "used. This case is useful when the " + "'scipion extract-coordinates' protocol is used with " + "option to apply shifts, so only the angles should be " + "used after this operation.") form.addParallelSection(threads=0, mpi=0) -#--------------------------- INSERT steps functions -------------------------------------------- + # --------------------------- INSERT steps functions ---------------------- + def _initialize(self): + """ Some variables initialization. """ + parts = self.inputParticles.get() + alignment = self.inputAlignment.get() + self.scale = alignment.getSamplingRate() / parts.getSamplingRate() + # Compute the factor to modify the shifts + # it will be the scale of pixelSize or zero if only angles + self.factor = 0 + if self.assignMode != self.ASSIGN_ONLY_ANGLES: + self.factor = self.scale + self.mode = self.assignMode.get() def _insertAllSteps(self): """for each ctf insert the steps to compare it @@ -63,19 +97,24 @@ def _updateItem(self, item, row): that will be stored in the output Set. """ # Add alignment info from corresponding item on inputAlignment - inputAlignment = self.inputAlignment.get() - scale = inputAlignment.getSamplingRate()/self.inputParticles.get().getSamplingRate() + inputAlign = self.inputAlignment.get() - alignedParticle = inputAlignment[item.getObjId()] - # If alignment is found for this particle set the alignment info on the output particle, if not do not write that item + alignedParticle = inputAlign[item.getObjId()] + # If alignment is found for this particle set the alignment info on + # the output particle, if not do not write that item if alignedParticle is not None: alignment = alignedParticle.getTransform() - alignment.scaleShifts(scale) + m = alignment.getMatrix() + if self.mode == self.ASSIGN_ONLY_SHIFTS: + m[:3, :3] = 0 # Set to zero the angles part + # Always scale the shifts or set to zero if only angles + m[:3, 3] *= self.factor item.setTransform(alignment) else: item._appendItem = False def createOutputStep(self): + self._initialize() inputParticles = self.inputParticles.get() inputAlignment = self.inputAlignment.get() outputParticles = self._createSetOfParticles() @@ -83,33 +122,38 @@ def createOutputStep(self): outputParticles.setAlignment(inputAlignment.getAlignment()) outputParticles.copyItems(inputParticles, - updateItemCallback=self._updateItem) + updateItemCallback=self._updateItem) self._defineOutputs(outputParticles=outputParticles) self._defineSourceRelation(self.inputParticles, outputParticles) self._defineSourceRelation(self.inputAlignment, outputParticles) def _summary(self): + self._initialize() summary = [] if not hasattr(self, 'outputParticles'): summary.append("Output particles not ready yet.") else: - scale = self.inputAlignment.get().getSamplingRate()/self.inputParticles.get().getSamplingRate() - summary.append("Assigned alignment to %s particles from a total of %s." % (self.outputParticles.getSize(), self.inputParticles.get().getSize())) - if scale != 1: - summary.append("Applied scale of %s." % scale) + + summary.append("Assigned alignment to %s particles from a total " + "of %s." % (self.outputParticles.getSize(), + self.inputParticles.get().getSize())) + if self.scale != 1: + summary.append("Applied scale of %s." % self.scale) return summary def _methods(self): + self._initialize() methods = [] if not hasattr(self, 'outputParticles'): methods.append("Output particles not ready yet.") else: - scale = self.inputAlignment.get().getSamplingRate()/self.inputParticles.get().getSamplingRate() - methods.append("We assigned alignment to %s particles from %s and produced %s." - % (self.outputParticles.getSize(), self.getObjectTag('inputParticles'), self.getObjectTag('outputParticles'))) - if scale != 1: - methods.append("Applied scale factor of %s." % scale) + methods.append("We assigned alignment to %s particles from %s and " + "produced %s." % (self.outputParticles.getSize(), + self.getObjectTag('inputParticles'), + self.getObjectTag('outputParticles'))) + if self.scale != 1: + methods.append("Applied scale factor of %s." % self.scale) return methods def _validate(self): @@ -117,18 +161,21 @@ def _validate(self): is launched to be executed. It should return a list of errors. If the list is empty the protocol can be executed. """ - #check that input set of aligned particles do have 2D alignment - errors = [ ] + # Check that input set of aligned particles do have 2D alignment + errors = [] inputAlignmentSet = self.inputAlignment.get() if not inputAlignmentSet.hasAlignment(): - errors.append("Input alignment set should contains some kind of alignment (2D, 3D or Projection).") + errors.append("Input alignment set should contains some kind of " + "alignment (2D, 3D or Projection).") else: - # Just for consistency, check that the particles really contains Transform object + # Just for consistency, check that the particles really + # contains Transform object first = inputAlignmentSet.getFirstItem() alignment = first.getTransform() if alignment is None: errors.append('Inconsistency detected in *Input alignment* !!!') - errors.append('It has alignment: _%s_, but the alignment is missing!!!' % inputAlignmentSet.getAlignment()) + errors.append('It has alignment: _%s_, but the alignment is ' + 'missing!!!' % inputAlignmentSet.getAlignment()) # Add some errors if input is not valid return errors