# **************************************************************************
# *
# * Authors: J.M. De la Rosa Trevin (delarosatrevin@scilifelab.se) [1]
# *
# * [1] SciLifeLab, Stockholm University
# *
# * 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 3 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
# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# * 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 'scipion@cnb.csic.es'
# *
# **************************************************************************
import enum
import os
from datetime import datetime
from collections import OrderedDict
import pyworkflow.object as pwobj
import pyworkflow.protocol.params as params
import pyworkflow.protocol as pwprot
import pyworkflow.utils as pwutils
import pwem.objects as emobj
from pwem.protocols import ProtParticles
[docs]class ProtParticlePickingOutput(enum.Enum):
""" Possible outputs for particle picking protocols
"""
outputCoordinates = emobj.SetOfCoordinates
[docs]class ProtParticlePicking(ProtParticles):
_possibleOutputs = ProtParticlePickingOutput
OUTPUT_PREFIX = ProtParticlePickingOutput.outputCoordinates.name
def _defineParams(self, form):
form.addSection(label='Input')
form.addParam('inputMicrographs', params.PointerParam,
pointerClass='SetOfMicrographs',
label=pwutils.Message.LABEL_INPUT_MIC, important=True,
help='Select the SetOfMicrographs to be used during '
'picking.')
# -------------------------- INFO functions -------------------------------
[docs] def getSummary(self, coordSet):
summary = []
summary.append("Number of particles picked: %s" % coordSet.getSize())
summary.append("Particle size: %s" % coordSet.getBoxSize())
return "\n".join(summary)
[docs] def getMethods(self, output):
msg = 'User picked %d particles ' % output.getSize()
msg += 'with a particle size of %s.' % output.getBoxSize()
return msg
def _methods(self):
methodsMsgs = []
if self.getInputMicrographs() is None:
return ['Input micrographs not available yet.']
methodsMsgs.append("Input micrographs %s of size %d."
% (self.getObjectTag(self.getInputMicrographs()),
self.getInputMicrographs().getSize()))
if self.getOutputsSize() >= 1:
for key, output in self.iterOutputAttributes(emobj.SetOfCoordinates):
msg = self.getMethods(output)
methodsMsgs.append("%s: %s" % (self.getObjectTag(output), msg))
else:
methodsMsgs.append(pwutils.Message.TEXT_NO_OUTPUT_CO)
return methodsMsgs
def _summary(self):
summary = []
if self.getInputMicrographs() is not None:
summary.append("Number of input micrographs: %d"
% self.getInputMicrographs().getSize())
if self.getOutputsSize() >= 1:
for key, output in self.iterOutputAttributes():
summary.append("*%s:* \n %s " % (key, output.getObjComment()))
else:
summary.append(pwutils.Message.TEXT_NO_OUTPUT_CO)
return summary
[docs] def getInputMicrographsPointer(self):
return self.inputMicrographs
[docs] def getInputMicrographs(self):
return self.getInputMicrographsPointer().get()
def _getCoords(self, CoordClass):
result = None
for _, attr in self.iterOutputAttributes(CoordClass):
result = attr # Get the last output that is SetOfCoordinates or so
return result
[docs] def getCoords(self):
return self._getCoords(emobj.SetOfCoordinates)
[docs] def getCoordsTiltPair(self):
from pwem.objects.data_tiltpairs import CoordinatesTiltPair
return self._getCoords(CoordinatesTiltPair)
def _createOutput(self, outputDir):
micSet = self.getInputMicrographs()
suffix = self.__getOutputSuffix()
outputName = self.OUTPUT_PREFIX + suffix
coordSet = self._createSetOfCoordinates(micSet, suffix)
self.readSetOfCoordinates(outputDir, coordSet)
coordSet.setObjComment(self.getSummary(coordSet))
outputs = {outputName: coordSet}
self._defineOutputs(**outputs)
self._defineSourceRelation(self.getInputMicrographsPointer(), coordSet)
[docs] def createOutputStep(self):
self._createOutput(self._getExtraPath())
[docs] def readSetOfCoordinates(self, workingDir, coordSet):
pass
[docs] def getCoordsDir(self):
pass
def __getOutputSuffix(self):
""" Get the name to be used for a new output.
For example: outputCoordinates7.
It should take into account previous outputs
and number with a higher value.
"""
maxCounter = -1
for attrName, _ in self.iterOutputAttributes(emobj.SetOfCoordinates):
suffix = attrName.replace(self.OUTPUT_PREFIX, '')
try:
counter = int(suffix)
except Exception as ex:
counter = 1 # when there is not number assume 1
maxCounter = max(counter, maxCounter)
return str(maxCounter+1) if maxCounter > 0 else '' # empty if not output
[docs] def registerCoords(self, coordsDir):
""" This method is usually inherited by all Pickers
and it is used from the Java picking GUI to register
a new SetOfCoordinates when the user click on +Particles button.
"""
from pwem import Domain
suffix = self.__getOutputSuffix()
outputName = self.OUTPUT_PREFIX + suffix
readSetOfCoordinates = Domain.importFromPlugin('xmipp3.convert',
'readSetOfCoordinates')
inputset = self.getInputMicrographs()
# micrographs are the input set if protocol is not finished
outputset = self._createSetOfCoordinates(inputset, suffix=suffix)
readSetOfCoordinates(coordsDir, outputset.getMicrographs(), outputset)
summary = self.getSummary(outputset)
outputset.setObjComment(summary)
outputs = {outputName: outputset}
self._defineOutputs(**outputs)
# Using a pointer to define the relations is more robust to scheduling
# and id changes between the protocol run.db and the main project
# database. The pointer defined below points to the outputset object
self._defineSourceRelation(self.getInputMicrographsPointer(),
pwobj.Pointer(value=self,
extended=outputName))
self._store()
[docs]class ProtParticlePickingAuto(ProtParticlePicking):
""" A derived class from ProtParticlePicking to differentiate those
picking protocols that works in automatic way, i.e., that are not
interactive and are also good candidates to be run in streaming. """
def _insertAllSteps(self):
self.initialIds = self._insertInitialSteps()
self.micDict = OrderedDict()
pwutils.makeFilePath(self._getAllDone())
micDict, self.streamClosed = self._loadInputList()
pickMicIds = self._insertNewMicsSteps(micDict.values())
self._insertFinalSteps(pickMicIds)
def _insertInitialSteps(self):
""" Override this function to insert some steps before the
picking micrograph steps.
Should return a list of ids of the initial steps. """
return []
def _insertFinalSteps(self, micSteps):
""" Override this function to insert some steps after the
picking micrograph steps.
Receive the list of step ids of the picking steps. """
self._insertFunctionStep(self.createOutputStep.__name__,
prerequisites=micSteps, wait=True, needsGPU=False)
def _getPickArgs(self):
""" Should be implemented in sub-classes to define the argument
list that should be passed to the picking step function.
"""
return []
def _insertPickMicrographStep(self, mic, prerequisites, *args):
""" Basic method to insert a picking step for a given micrograph. """
micStepId = self._insertFunctionStep('pickMicrographStep',
mic.getMicName(), *args,
prerequisites=prerequisites)
return micStepId
[docs] def pickMicrographStep(self, micName, *args):
""" Step function that will be common for all picking protocols.
It will take care of re-building the micrograph object from the micDict
argument and perform any conversion if needed. Then, the function
_pickMicrograph will be called, that should be implemented by each
picking protocol.
"""
mic = self.micDict[micName]
micDoneFn = self._getMicDone(mic)
micFn = mic.getFileName()
if self.isContinued() and os.path.exists(micDoneFn):
self.info("Skipping micrograph: %s, seems to be done" % micFn)
return
# Clean old finished files
pwutils.cleanPath(micDoneFn)
self.info("Picking micrograph: %s " % micFn)
self._pickMicrograph(mic, *args)
# Mark this mic as finished
open(micDoneFn, 'w').close()
def _pickMicrograph(self, mic, *args):
""" This function should be implemented by subclasses in order
to picking the given micrograph. """
pass
# Group of functions to pick several micrographs if the batch size is
# defined. In some programs it might be more efficient to pick many
# at once and not one by one
def _insertPickMicrographListStep(self, micList, prerequisites, *args):
""" Basic method to insert a picking step for a given micrograph. """
micNameList = [mic.getMicName() for mic in micList]
micStepId = self._insertFunctionStep('pickMicrographListStep',
micNameList, *args,
prerequisites=prerequisites)
return micStepId
[docs] def pickMicrographListStep(self, micNameList, *args):
micList = []
for micName in micNameList:
mic = self.micDict[micName]
micDoneFn = self._getMicDone(mic)
micFn = mic.getFileName()
if self.isContinued() and os.path.exists(micDoneFn):
self.info("Skipping micrograph: %s, seems to be done" % micFn)
else:
# Clean old finished files
pwutils.cleanPath(micDoneFn)
self.info("Picking micrograph: %s " % micFn)
micList.append(mic)
self._pickMicrographList(micList, *args)
for mic in micList:
# Mark this mic as finished
open(self._getMicDone(mic), 'w').close()
def _pickMicrographList(self, micList, *args):
""" This function can be implemented by subclasses if it is a more
efficient way to pick many micrographs at once.
Default implementation will just call the _pickMicrograph
"""
for mic in micList:
self._pickMicrograph(mic, *args)
# --------------------------- UTILS functions ----------------------------
def _stepsCheck(self):
# To allow streaming picking we need to detect:
# 1) new micrographs ready to be picked
# 2) new output coordinates that have been produced and add then
# to the output set.
self._checkNewInput()
self._checkNewOutput()
def _insertNewMicsSteps(self, inputMics):
""" Insert steps to process new mics (from streaming)
Params:
insertedDict: contains already processed mics
inputMics: input mics set to be check
"""
return self._insertNewMics(inputMics,
lambda mic: mic.getMicName(),
self._insertPickMicrographStep,
self._insertPickMicrographListStep,
*self._getPickArgs())
def _loadSet(self, inputSet, SetClass, getKeyFunc):
""" Load a given input set if their items are not already present
in the self.micDict.
This can be used to load new micrographs for picking as well as
new CTF (if used) in streaming.
"""
setFn = inputSet.getFileName()
self.debug("Loading input db: %s" % setFn)
updatedSet = SetClass(filename=setFn)
updatedSet.loadAllProperties()
newItemDict = OrderedDict()
for item in updatedSet:
micKey = getKeyFunc(item)
if micKey not in self.micDict:
newItemDict[micKey] = item.clone()
streamClosed = updatedSet.isStreamClosed()
updatedSet.close()
self.debug("Closed db.")
return newItemDict, streamClosed
def _loadMics(self, micSet):
return self._loadSet(micSet, emobj.SetOfMicrographs,
lambda mic: mic.getMicName())
def _loadCTFs(self, ctfSet):
return self._loadSet(ctfSet, emobj.SetOfCTF,
lambda ctf: ctf.getMicrograph().getMicName())
def _loadInputList(self):
""" Load the input set of micrographs that are ready to be picked. """
return self._loadMics(self.getInputMicrographs())
def _checkNewInput(self):
# Check if there are new micrographs to process from the input set
localFile = self.getInputMicrographs().getFileName()
now = datetime.now()
self.lastCheck = getattr(self, 'lastCheck', now)
mTime = datetime.fromtimestamp(os.path.getmtime(localFile))
self.debug('Last check: %s, modification: %s'
% (pwutils.prettyTime(self.lastCheck),
pwutils.prettyTime(mTime)))
# If the input micrographs.sqlite have not changed since our last check,
# it does not make sense to check for new input data
if self.lastCheck > mTime and hasattr(self, 'listOfMics'):
return None
self.lastCheck = now
# Open input micrographs.sqlite and close it as soon as possible
micDict, self.streamClosed = self._loadInputList()
newMics = micDict.values()
outputStep = self._getFirstJoinStep()
if newMics:
fDeps = self._insertNewMicsSteps(newMics)
if outputStep is not None:
outputStep.addPrerequisites(*fDeps)
self.updateSteps()
def _checkNewOutput(self):
if getattr(self, 'finished', False):
return
# Load previously done items (from text file)
doneList = self._readDoneList()
# Check for newly done items
listOfMics = self.micDict.values()
nMics = len(listOfMics)
newDone = [m for m in listOfMics
if m.getObjId() not in doneList and self._isMicDone(m)]
# Update the file with the newly done mics
# or exit from the function if no new done mics
self.debug('_checkNewOutput: ')
self.debug(' listOfMics: %s, doneList: %s, newDone: %s'
% (nMics, len(doneList), len(newDone)))
allDone = len(doneList) + len(newDone)
# We have finished when there is not more input mics (stream closed)
# and the number of processed mics is equal to the number of inputs
self.finished = self.streamClosed and allDone == nMics
streamMode = pwobj.Set.STREAM_CLOSED if self.finished else pwobj.Set.STREAM_OPEN
self.debug(' streamMode: %s newDone: %s' % (streamMode,
not(newDone == [])))
if newDone:
newDoneUpdated = self._updateOutputCoordSet(newDone, streamMode)
self._writeDoneList(newDoneUpdated)
elif not self.finished:
# If we are not finished and no new output have been produced
# it does not make sense to proceed and updated the outputs
# so we exit from the function here
# Maybe it would be good idea to take a snap to avoid
# so much IO if this protocol does not have much to do now
if allDone == nMics:
self._streamingSleepOnWait()
return
self.debug(' finished: %s ' % self.finished)
self.debug(' self.streamClosed (%s) AND' % self.streamClosed)
self.debug(' allDone (%s) == len(self.listOfMics (%s)'
% (allDone, nMics))
if self.finished: # Unlock createOutputStep if finished all jobs
self._updateStreamState(streamMode)
outputStep = self._getFirstJoinStep()
if outputStep and outputStep.isWaiting():
outputStep.setStatus(pwprot.STATUS_NEW)
def _micIsReady(self, mic):
""" Function to check if a micrograph (although reported done)
is ready to update the coordinates from it. An practical use of this
function will be for protocols that need to wait for the CTF of that
micrograph to be ready as well.
"""
return True
[docs] def readCoordsFromMics(self, outputDir, micDoneList, outputCoords):
""" This method should be implemented in subclasses to read
the coordinates from a given list of micrographs.
"""
pass
def _updateOutputCoordSet(self, micList, streamMode):
micDoneList = [mic for mic in micList if self._micIsReady(mic)]
# Do no proceed if there is not micrograph ready
if not micDoneList:
return []
outputName = 'outputCoordinates'
outputDir = self.getCoordsDir()
outputCoords = getattr(self, outputName, None)
# If there are not outputCoordinates yet, it means that is the first
# time we are updating output coordinates, so we need to first create
# the output set
firstTime = outputCoords is None
if firstTime:
micSetPtr = self.getInputMicrographsPointer()
outputCoords = self._createSetOfCoordinates(micSetPtr)
else:
outputCoords.enableAppend()
self.info("Reading coordinates from mics: %s" % ','.join([mic.strId() for mic in micList]))
self.readCoordsFromMics(outputDir, micDoneList, outputCoords)
self.debug(" _updateOutputCoordSet Stream Mode: %s " % streamMode)
outputCoords.setObjComment(self.getSummary(outputCoords))
self._updateOutputSet(outputName, outputCoords, streamMode)
if firstTime:
self._defineSourceRelation(self.getInputMicrographsPointer(),
outputCoords)
return micDoneList
def _updateStreamState(self, streamMode):
outputName = 'outputCoordinates'
outputCoords = getattr(self, outputName, None)
# If there are not outputCoordinates yet, it means that is the first
# time we are updating output coordinates, so we need to first create
# the output set
firstTime = outputCoords is None
if firstTime:
micSetPtr = self.getInputMicrographsPointer()
outputCoords = self._createSetOfCoordinates(micSetPtr)
else:
outputCoords.enableAppend()
self.debug(" _updateStreamState Stream Mode: %s " % streamMode)
self._updateOutputSet(outputName, outputCoords, streamMode)
def _getMicDone(self, mic):
return self._getExtraPath('DONE', 'mic_%06d.TXT' % mic.getObjId())
def _isMicDone(self, mic):
""" A mic is done if the marker file exists. """
return os.path.exists(self._getMicDone(mic))
def _getAllDone(self):
return self._getExtraPath('DONE', 'all.TXT')
def _readDoneList(self):
""" Read from a text file the id's of the items that have been done. """
doneFile = self._getAllDone()
doneList = []
# Check what items have been previously done
if os.path.exists(doneFile):
with open(doneFile) as f:
doneList += [int(line.strip()) for line in f]
return doneList
def _writeDoneList(self, micList):
""" Write to a text file the items that have been done. """
doneFile = self._getAllDone()
if not os.path.exists(doneFile):
pwutils.makeFilePath(doneFile)
with open(doneFile, 'a') as f:
for mic in micList:
f.write('%d\n' % mic.getObjId())
def _getFirstJoinStepName(self):
# This function will be used for streaming, to check which is
# the first function that need to wait for all micrographs
# to have completed, this can be overwritten in subclasses
# (eg in Xmipp 'sortPSDStep')
return self.createOutputStep.__name__
def _getFirstJoinStep(self):
for s in self._steps:
if s.funcName == self._getFirstJoinStepName():
return s
return None
[docs] def createOutputStep(self):
# Not really required now
# self._createOutput(self._getExtraPath())
pass
[docs]class ProtPickingDifference(ProtParticlePicking):
"""
Protocol to compute the difference between a reference SetOfParticles and
a another set (usually a negative reference).
The output will be a SetOfCoordinates with the particles in the reference
input that are not close to coordinates in the negative input.
"""
_label = 'picking difference'
def _defineParams(self, form):
form.addSection(label='Input')
form.addParam('inputCoordinates', params.PointerParam,
pointerClass='SetOfCoordinates',
label="Input coordinates",
help='Select the reference set of coordinates.')
form.addParam('negativeCoordinates', params.PointerParam,
pointerClass='SetOfCoordinates',
label="Negative coordinates",
help='Negative coordinates that will help to exclude '
'coordinates from the input set. ')
form.addParam('differenceRadius', params.IntParam, default=10,
label="Radius (px)",
help="Distance radius to consider that two coordinates "
"close enough. If a particle in the input reference"
"set have any close particle in the negative set, "
"it will not be included in the output set. ")
form.addParallelSection(threads=0, mpi=0)
# ------------------------- INSERT steps functions -------------------------
def _insertAllSteps(self):
self._insertFunctionStep("createOutputStep",
self.inputCoordinates.getObjId(),
self.negativeCoordinates.getObjId(),
self.differenceRadius.get())
def _summary(self):
summary = []
return summary
def _methods(self):
return []
[docs] def createOutputStep(self, inputId, negId, radius):
inputCoords = self.inputCoordinates.get()
negCoords = self.negativeCoordinates.get()
inputMics = inputCoords.getMicrographs()
outputCoords = self._createSetOfCoordinates(inputMics)
outputCoords.setBoxSize(inputCoords.getBoxSize())
r2 = radius * radius # to avoid computing sqrt when comparing distances
def _discard(coord, negPos):
""" Find if there is a close negative particle to this one. """
x, y = coord.getPosition()
for nx, ny in negPos:
if abs((x - nx) * (y - ny)) < r2:
return True
return False # Far enough from all negative coordinates
# Read all consensus particles
for mic in inputMics:
negPos = [(c.getX(), c.getY())
for c in negCoords.iterCoordinates(mic)]
for coord in inputCoords.iterCoordinates(mic):
if not _discard(coord, negPos):
outputCoords.append(coord)
# Set output
self._defineOutputs(outputCoordinates=outputCoords)
self._defineTransformRelation(self.inputCoordinates,
self.outputCoordinates)
self._defineSourceRelation(self.negativeCoordinates,
self.outputCoordinates)