diff --git a/CHANGES.txt b/CHANGES.txt index 188d374..becf7de 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -1,3 +1,23 @@ +v3.3.1: + Users: + - Small fixes regarding the register of the tilt series acquisition attributes. + - Fix the template matching output register reported by Ricardo Righetto (thanks!). + - Fix the classic subtomo refine transformation matrices assignment when generating the output. + - Fix the particles order in the registered output of protocol extract particles from TS. + - Protocol 'extract particles from TS' is now semi-streamified. + - Protocol 'tilt series alignment and tomo reconstruction': + * Fix the sampling rate of the set of tilt series interpolated. + * Fix the generation of tomograms when not requested. + * The interpolated tilt series is now generated applying the transformation matrix in Scipion convention to the + non-interpolated tilt series. + * Input and output convert steps improved. + * Remove form parameter 'Step between tilts' as it can now be read from the acquisition. + * Add param to correct the XDrift. + * Some other changes in the protocol form. + - Protocol 'average subtomograms' adapted to work with EMAN 3D particles generated from 3D particles (the ones + resulting when using the protocol 'extract particles from TS') + Developers: + - Test for TS alignment & tomo reconstruction updated to the new TCL. v3.3.0: Users: - Prevent viewers from locking the screen focus. diff --git a/emantomo/__init__.py b/emantomo/__init__.py index 76a99d5..918375f 100644 --- a/emantomo/__init__.py +++ b/emantomo/__init__.py @@ -22,7 +22,7 @@ # * e-mail address 'scipion@cnb.csic.es' # * # ************************************************************************** -__version__ = "3.3.0" +__version__ = "3.3.1" _logo = "eman2_logo.png" _references = ['GALAZMONTOYA2015279', 'BELL201625'] _url = "https://github.com/scipion-em/scipion-em-emantomo" diff --git a/emantomo/convert/convert.py b/emantomo/convert/convert.py index 694c173..4bb6416 100644 --- a/emantomo/convert/convert.py +++ b/emantomo/convert/convert.py @@ -449,12 +449,18 @@ def jsons2SetCoords3D(protocol, setTomograms, outPath): coord3DSet.setSamplingRate(setTomograms.getSamplingRate()) first = True for tomo in setTomograms.iterItems(): - outFile = '*%s_info.json' % pwutils.removeBaseExt(tomo.getFileName().split("__")[0]) + # File search by tsId + outFile = '*%s_info.json' % tomo.getTsId() pattern = join(outPath, outFile) files = glob.glob(pattern) if not files or not isfile(files[0]): - continue + # File search by filename + outFile = '*%s_info.json' % pwutils.removeBaseExt(tomo.getFileName().split("__")[0]) + pattern = join(outPath, outFile) + files = glob.glob(pattern) + if not files or not isfile(files[0]): + continue jsonFnbase = files[0] jsonBoxDict = loadJson(jsonFnbase) diff --git a/emantomo/objects.py b/emantomo/objects.py index f56d721..5d33915 100644 --- a/emantomo/objects.py +++ b/emantomo/objects.py @@ -73,6 +73,19 @@ def getProjsFrom2dStack(self, shrink=1): (shrink * img.attrs['EMAN.ptcl_source_coord'][:-1]).tolist()) return projsList + def get3dCoordsFrom3dStack(self, scaleFactor=1, invertZ=False): + """Generate a list of elements in which each element is a list containing the following data: + [xCoord, yCoord, zCoord] read from the HDF stack header. + """ + coordsList = [] + for particleId in range(len(self._imgObjList)): + img = self._imgObjList[str(particleId)] + coords = (scaleFactor * img.attrs['EMAN.ptcl_source_coord']).tolist() + if invertZ: + coords[-1] = -1 * coords[-1] + coordsList.append(coords) + return coordsList + def getSamplingRate(self): """Reads the sampling rate from the HDF header""" return self._imgObjList['0'].attrs['EMAN.apix_x'] diff --git a/emantomo/protocols/protocol_average_subtomos.py b/emantomo/protocols/protocol_average_subtomos.py index bb93e15..fa556c4 100644 --- a/emantomo/protocols/protocol_average_subtomos.py +++ b/emantomo/protocols/protocol_average_subtomos.py @@ -29,13 +29,13 @@ import glob import os import shutil -from os.path import join, basename +from os.path import join, basename, abspath from pyworkflow import BETA from pwem.protocols import EMProtocol from pyworkflow.protocol import PointerParam, FloatParam, StringParam, BooleanParam, LEVEL_ADVANCED -from pyworkflow.utils import Message, makePath, removeBaseExt, replaceExt +from pyworkflow.utils import Message, makePath, removeBaseExt, replaceExt, createLink from ..constants import SPT_00_DIR, INPUT_PTCLS_LST, THREED_01, SYMMETRY_HELP_MSG, SUBTOMOGRAMS_DIR from ..convert import writeSetOfSubTomograms, refinement2Json @@ -43,6 +43,7 @@ from tomo.protocols import ProtTomoBase from tomo.objects import AverageSubTomogram +from ..objects import EmanParticle, EmanSetOfParticles class OutputsAverageSubtomos(enum.Enum): @@ -118,7 +119,13 @@ def _initialize(self): def convertInputStep(self): inSubtomos = self.inputSetOfSubTomogram.get() - writeSetOfSubTomograms(inSubtomos, self.hdfSubtomosDir) + if type(inSubtomos) is EmanSetOfParticles: + # If True, the 3D particles HDF stacks will already exist because they have been extracted with EMAN pppt + hdfStacks = inSubtomos.getUniqueValues(EmanParticle.STACK_3D_HDF) + [createLink(abspath(hdfStack), self._getExtraPath(SUBTOMOGRAMS_DIR, basename(hdfStack))) for + hdfStack in hdfStacks] + else: + writeSetOfSubTomograms(inSubtomos, self.hdfSubtomosDir) refinement2Json(self, inSubtomos) # Generate a virtual stack of particle represented by a .lst file, as expected by EMAN diff --git a/emantomo/protocols/protocol_base.py b/emantomo/protocols/protocol_base.py index a82fa46..ced43d8 100644 --- a/emantomo/protocols/protocol_base.py +++ b/emantomo/protocols/protocol_base.py @@ -24,6 +24,7 @@ # * # ************************************************************************** import glob +import logging import re from os.path import join, abspath, basename @@ -44,6 +45,9 @@ from tomo.protocols import ProtTomoBase from tomo.utils import getNonInterpolatedTsFromRelations + +logger = logging.getLogger(__name__) + IN_TS = 'inputTS' IN_COORDS = 'inputCoordinates' IN_CTF = 'inputCTF' @@ -70,9 +74,7 @@ def convertTsStep(self, mdObj): # The converted TS must be unbinned, because EMAN will read the sampling rate from its header. This is why # the TS associated to the CTF is the one considered first. Later, when generating the json, the TS alignment # parameters are read from the introduced TS and the shifts are scaled to at the unbinned scale - # if mdObj.ctf: - # inTsFName = mdObj.ctf.getTiltSeries().getFirstItem().getFileName() - # else: + logger.info(f'Converting TS {mdObj.tsId} into HDF...') inTsFName = mdObj.ts.getFirstItem().getFileName() sRate = mdObj.ts.getSamplingRate() self.convertOrLink(inTsFName, mdObj.tsId, TS_DIR, sRate) diff --git a/emantomo/protocols/protocol_clip_tomograms.py b/emantomo/protocols/protocol_clip_tomograms.py index 24e2d2b..3654983 100644 --- a/emantomo/protocols/protocol_clip_tomograms.py +++ b/emantomo/protocols/protocol_clip_tomograms.py @@ -62,7 +62,7 @@ def _defineParams(self, form): label="Input tomograms", important=True) - line = form.addLine('New center coordinates [pix.] (opt.)', + line = form.addLine('New center coordinates [px] (opt.)', help='For each empty coordinate, the correspònding original center coordinate ' 'will be used.') line.addParam('xc', IntParam, @@ -75,7 +75,7 @@ def _defineParams(self, form): label="cz", allowsNull=True) - line = form.addLine('New dimensions [pix.] (opt.)', + line = form.addLine('New dimensions [px] (opt.)', help='For each empty dimension, the original corresponding dimension will be used.') line.addParam('xDim', IntParam, label="dx", diff --git a/emantomo/protocols/protocol_extraction_from_ts.py b/emantomo/protocols/protocol_extraction_from_ts.py index f7a8fc7..0899de6 100644 --- a/emantomo/protocols/protocol_extraction_from_ts.py +++ b/emantomo/protocols/protocol_extraction_from_ts.py @@ -25,8 +25,10 @@ # * # ************************************************************************** import glob +import logging from enum import Enum from os.path import join, exists, basename, abspath +import numpy as np from emantomo import Plugin from emantomo.constants import GROUP_ID, PARTICLES_3D_DIR, PARTICLES_DIR, TOMOBOX, TOMOGRAMS_DIR, TS_DIR from emantomo.objects import EmanHdf5Handler, EmanSetOfParticles, EmanParticle, EmanMetaData @@ -35,14 +37,18 @@ getPresentPrecedents from pwem.convert.headers import fixVolume from pwem.objects import Transform -from pyworkflow.protocol import PointerParam, FloatParam, LEVEL_ADVANCED, GE, LE, GT, IntParam, BooleanParam +from pyworkflow.protocol import PointerParam, FloatParam, LEVEL_ADVANCED, GE, LE, GT, IntParam, BooleanParam, \ + STEPS_PARALLEL from pyworkflow.utils import Message, replaceExt, removeExt, createLink from emantomo.convert import coords2Json, ts2Json, ctfTomo2Json -from tomo.constants import TR_EMAN +from tomo.constants import TR_EMAN, SCIPION from tomo.objects import SetOfLandmarkModels, LandmarkModel, Coordinate3D, SetOfSubTomograms, SetOfCoordinates3D, \ SetOfMeshes +logger = logging.getLogger(__name__) + + class outputObjects(Enum): subtomograms = EmanSetOfParticles projected2DCoordinates = SetOfLandmarkModels @@ -56,6 +62,7 @@ class EmanProtTSExtraction(ProtEmantomoBase): def __init__(self, **kwargs): super().__init__(**kwargs) + self.stepsExecutionMode = STEPS_PARALLEL self.groupIds = None self.emanDict = None self.sphAb = None @@ -64,15 +71,17 @@ def __init__(self, **kwargs): self.projectionsDict = {} self.currentBoxSize = None self.currentSRate = None + self.fiducialSize = None # --------------------------- DEFINE param functions ---------------------- def _defineParams(self, form): form.addSection(label=Message.LABEL_INPUT) form.addParam(IN_SUBTOMOS, PointerParam, - label="Coordinates", + label="Coordinates or 3D particles", important=True, - pointerClass='SetOfCoordinates3D, EmanSetOfParticles', - help='The corresponding tomograms data will be accessed from the provided coordinates.') + pointerClass='SetOfCoordinates3D, SetOfSubTomograms', + help='The corresponding tomograms data will be accessed from the provided coordinates or the ' + 'coordinates associated to the 3D particles.') form.addParam(IN_CTF, PointerParam, label="CTF tomo series", pointerClass='SetOfCTFTomoSeries', @@ -93,7 +102,7 @@ def _defineParams(self, form): form.addParam(IN_BOXSIZE, IntParam, allowsNull=False, important=True, - label='Box size unbinned (pix.)', + label='Box size unbinned (px)', help='The subtomograms are extracted as a cubic box of this size.') form.addParam('shrink', FloatParam, default=1, @@ -139,18 +148,22 @@ def _defineParams(self, form): # --------------------------- INSERT steps functions ---------------------- def _insertAllSteps(self): + pIds = [] mdObjDict = self._initialize() - self._insertFunctionStep(self.createExtractionEmanPrjStep, mdObjDict) for mdObj in mdObjDict.values(): - self._insertFunctionStep(self.convertTsStep, mdObj) - self._insertFunctionStep(self.writeData2JsonFileStep, mdObjDict) - self._insertFunctionStep(self.extractParticlesStep, mdObjDict) - for mdObj in mdObjDict.values(): - self._insertFunctionStep(self.convertOutputStep, mdObj) - self._insertFunctionStep(self.createOutputStep, mdObjDict) + createPrjId = self._insertFunctionStep(self.createExtractionEmanPrjStep, mdObj) + convertInId = self._insertFunctionStep(self.convertTsStep, mdObj, prerequisites=createPrjId) + writeJsonId = self._insertFunctionStep(self.writeData2JsonFileStep, mdObj, prerequisites=convertInId) + extractId = self._insertFunctionStep(self.extractParticlesStep, mdObj, prerequisites=writeJsonId) + convertOutId = self._insertFunctionStep(self.convertOutputStep, mdObj, prerequisites=extractId) + createOutputId = self._insertFunctionStep(self.createOutputStep, mdObj, prerequisites=convertOutId) + pIds.extend([createPrjId, convertInId, writeJsonId, extractId, convertOutId, createOutputId]) + self._insertFunctionStep(self._closeOutputSet, prerequisites=pIds) # --------------------------- STEPS functions ----------------------------- def _initialize(self): + logger.info('Initializing...') + self.createInitEmanPrjDirs() inParticles = getattr(self, IN_SUBTOMOS).get() inCtfSet = getattr(self, IN_CTF).get() inTsSet = self.getAttrib(IN_TS) @@ -168,36 +181,39 @@ def _initialize(self): uniqueTomoValsDict = getFromPresentObjects(coords, [Coordinate3D.TOMO_ID_ATTR, GROUP_ID]) self.groupIds = uniqueTomoValsDict[GROUP_ID] self.emanDict = genEmanGrouping(self.groupIds) - # Calculate the scale factor - self.scaleFactor = inParticles.getSamplingRate() / inTsSet.getSamplingRate() + # Calculate the scale factor for coords and shift scaling: according to EMAN's behavior, it is the rate between + # the sampling rate of the input coordinates and the TS unbinned sampling rate multiplied by the binning factor + # introduced + self.scaleFactor = inParticles.getSamplingRate() / (inTsSet.getSamplingRate() * self.shrink.get()) + # The fiducial size is the diameter in nm + self.fiducialSize = 0.1 * self.getAttrib(IN_BOXSIZE) * self.getAttrib(IN_TS).getSamplingRate() / 2 # Generate the md object data units as a dict return self.genMdObjDict(inTsSet, inCtfSet, coords) - def createExtractionEmanPrjStep(self, mdObjDict): + def createExtractionEmanPrjStep(self, mdObj): + logger.info(f'Creating the project for tsIf = {mdObj.tsId}...') # Create project dir structure - self.createInitEmanPrjDirs() tomogramsDir = self.getTomogramsDir() # TS must be in HDF (in EMAN's native code it is searched in hardcoded HDF format based on the tomogram # basename), so it will be converted later. However, the tomograms can be linked in MRC without any problem - for mdObj in mdObjDict.values(): - tomoFName = mdObj.inTomo.getFileName() - createLink(tomoFName, join(tomogramsDir, basename(tomoFName))) + tomoFName = mdObj.inTomo.getFileName() + createLink(tomoFName, join(tomogramsDir, basename(tomoFName))) - def writeData2JsonFileStep(self, mdObjDict): - for mdObj in mdObjDict.values(): - mode = 'a' if exists(mdObj.jsonFile) else 'w' - coords2Json(mdObj, self.emanDict, self.groupIds, self.getBoxSize(), doFlipZ=self.doFlipZInTomo.get(), - mode=mode) - ts2Json(mdObj, mode='a') - ctfTomo2Json(mdObj, self.sphAb, self.voltage, mode='a') + def writeData2JsonFileStep(self, mdObj): + logger.info(f'Writing the json files for TS {mdObj.tsId}...') + mode = 'a' if exists(mdObj.jsonFile) else 'w' + coords2Json(mdObj, self.emanDict, self.groupIds, self.getBoxSize(), doFlipZ=self.doFlipZInTomo.get(), + mode=mode) + ts2Json(mdObj, mode='a') + ctfTomo2Json(mdObj, self.sphAb, self.voltage, mode='a') - def extractParticlesStep(self, mdObjDict): - # if self.isReExtraction: - # self.buildEmanSets(outAliPath=None) + def extractParticlesStep(self, mdObj): + logger.info(f'Extracting the particles from TS {mdObj.tsId}...') program = Plugin.getProgram('e2spt_extract.py') - self.runJob(program, self._genExtractArgs(mdObjDict), cwd=self._getExtraPath()) + self.runJob(program, self._genExtractArgs(mdObj), cwd=self._getExtraPath()) def convertOutputStep(self, mdObj): + logger.info(f'Converting and unstacking the 3D particles from TS {mdObj.tsId} into MRC...') fName = self._getEmanFName(mdObj.tsId) # Unstack the 3d particles HDF stack into individual MRC files stack3d = join(PARTICLES_3D_DIR, fName) @@ -209,74 +225,65 @@ def convertOutputStep(self, mdObj): # outFile = replaceExt(inFile, 'mrc') # self.convertBetweenHdfAndMrc(inFile, outFile) - def createOutputStep(self, mdObjDict): - inCoordsPointer = getattr(self, IN_SUBTOMOS) - inCoords = inCoordsPointer.get() - inTsPointer = getattr(self, IN_TS) - inTs = inTsPointer.get() - subtomoSet = EmanSetOfParticles.create(self._getPath(), template='emanParticles%s.sqlite') - if self.inParticles: - subtomoSet.copyInfo(inCoords) - else: - subtomoSet.setCoordinates3D(inCoordsPointer) - subtomoSet.setSamplingRate(self.currentSRate) - + def createOutputStep(self, mdObj): + tsId = mdObj.tsId + logger.info(f'Creating the outputs corresponding to TS {tsId}...') + subtomoSet = self.getOutSetOfParticles() # Generate the fiducial model (for data visualization purpose) - fiducialModelGaps = SetOfLandmarkModels.create(self._getPath(), suffix='Gaps') - fiducialModelGaps.copyInfo(inTs) - fiducialModelGaps.setSetOfTiltSeries(inTsPointer) - # The fiducial size is the diameter in nm - fiducialSize = 0.1 * self.getAttrib(IN_BOXSIZE) * self.getAttrib(IN_TS).getSamplingRate() / 2 - - nTotalParticles = 0 - for tsId, mdObj in mdObjDict.items(): - _, _, nImgs = mdObj.ts.getDimensions() - coords = mdObj.particles if self.inParticles else mdObj.coords - subtomoFiles = sorted(glob.glob(self._getExtraPath(PARTICLES_3D_DIR, '%s*.mrc' % tsId))) - stack2d = glob.glob(self._getExtraPath(PARTICLES_DIR, '%s*.hdf' % tsId))[0] - stack3d = glob.glob(self._getExtraPath(PARTICLES_3D_DIR, '%s*.hdf' % tsId))[0] - # Get the projections - tsProjections = self.getProjections(mdObj) - landmarkModelGaps = self.createLandmarkModelGaps(mdObj.ts, fiducialSize) - - particleCounter = 0 - for coord, subtomoFile in zip(coords, subtomoFiles): - self.fillLandmarkModel(landmarkModelGaps, tsProjections, nTotalParticles, nImgs) - subtomogram = EmanParticle() - transform = Transform() - subtomogram.setFileName(subtomoFile) - subtomogram.setLocation(particleCounter, subtomoFile) # The index is stored to be used as - # the position of the particle in the 3d HDF stack (to handle subsets - LST file gen.) - subtomogram.setSamplingRate(self.currentSRate) - subtomogram.setVolName(mdObj.tsId) - subtomogram.setCoordinate3D(coord) - M = coord.getMatrix() - shift_x = M[0, 3] - shift_y = M[1, 3] - shift_z = M[2, 3] - transform.setMatrix(M) - transform.setShifts(self.scaleFactor * shift_x, - self.scaleFactor * shift_y, - self.scaleFactor * shift_z) - subtomogram.setTransform(transform, convention=TR_EMAN) - # Fill EmanParticle own attributes - subtomogram.setInfoJson(mdObj.jsonFile) - subtomogram.setTsHdf(self._getExtraPath(mdObj.tsHdfName)) - subtomogram.setStack2dHdf(stack2d) - subtomogram.setStack3dHdf(stack3d) - subtomogram.setAbsIndex(nTotalParticles) - subtomoSet.append(subtomogram) - particleCounter += 1 - nTotalParticles += 1 - fiducialModelGaps.append(landmarkModelGaps) - - # Define outputs and relations - self._defineOutputs(**{self._possibleOutputs.subtomograms.name: subtomoSet, - self._possibleOutputs.projected2DCoordinates.name: fiducialModelGaps}) - self._defineSourceRelation(getattr(self, IN_SUBTOMOS), subtomoSet) - self._defineSourceRelation(getattr(self, IN_CTF), subtomoSet) - self._defineSourceRelation(inTsPointer, subtomoSet) - self._defineSourceRelation(inTsPointer, fiducialModelGaps) + fiducialModelGaps = self.getOutSetIfFiducials() + + # Get the projections + tsProjections = self.getProjections(mdObj) + landmarkModelGaps = self.createLandmarkModelGaps(mdObj.ts) + + # Because the order is changed when the particles are processed with EMAN, if they have transformation matrix, + # there will be a data mismatch. To avoid it, The EMAN coordinates will be read from the 3d HDF stack, then the + # corresponding coordinate will be located among the input set, so the transformation is correctly matched, and + # finally, the corresponding output unstacked MRC file will be located and added to the output particle. + emanCoords = self.getEmanCoordsFromHdfStack(tsId) + + _, _, nImgs = mdObj.ts.getDimensions() + inCoords = mdObj.particles if self.inParticles else mdObj.coords + subtomoFiles = sorted(glob.glob(self._getExtraPath(PARTICLES_3D_DIR, '%s*.mrc' % tsId))) + stack2d = glob.glob(self._getExtraPath(PARTICLES_DIR, '%s*.hdf' % tsId))[0] + stack3d = glob.glob(self._getExtraPath(PARTICLES_3D_DIR, '%s*.hdf' % tsId))[0] + + nTotalParticles = mdObj.processingInd + particleCounter = 0 + for emanCoord, subtomoFile in zip(emanCoords, subtomoFiles): + self.fillLandmarkModel(landmarkModelGaps, tsProjections, particleCounter, nImgs) + subtomogram = EmanParticle() + transform = Transform() + subtomogram.setFileName(subtomoFile) + subtomogram.setLocation(particleCounter, subtomoFile) # The index is stored to be used as + # the position of the particle in the 3d HDF stack (to handle subsets - LST file gen.) + subtomogram.setSamplingRate(self.currentSRate) + subtomogram.setVolName(mdObj.tsId) + scipionCoord, inCoords = self.getEmanMatchingCoord(emanCoord, inCoords) + subtomogram.setCoordinate3D(scipionCoord) + M = scipionCoord.getMatrix() + shift_x = M[0, 3] + shift_y = M[1, 3] + shift_z = M[2, 3] + transform.setMatrix(M) + transform.setShifts(self.scaleFactor * shift_x, + self.scaleFactor * shift_y, + self.scaleFactor * shift_z) + subtomogram.setTransform(transform, convention=TR_EMAN) + # Fill EmanParticle own attributes + subtomogram.setInfoJson(mdObj.jsonFile) + subtomogram.setTsHdf(self._getExtraPath(mdObj.tsHdfName)) + subtomogram.setStack2dHdf(stack2d) + subtomogram.setStack3dHdf(stack3d) + subtomogram.setAbsIndex(nTotalParticles) + subtomoSet.append(subtomogram) + particleCounter += 1 + nTotalParticles += 1 + subtomoSet.write() + self._store(subtomoSet) + fiducialModelGaps.append(landmarkModelGaps) + fiducialModelGaps.write() + self._store(fiducialModelGaps) # --------------------------- INFO functions ----------------------------------- def _warnings(self): @@ -314,6 +321,7 @@ def genMdObjDict(self, inTsSet, inCtfSet, coords): # Split all the data into subunits (EmanMetaData objects) referred to the same tsId mdObjDict = {} + procInd = 0 for tomoId, ts in sorted(tsIdsDict.items()): # Sorting the items before the iteration based on the tsId (key) will ensure the alphabetical order in # all the data generated, preventing data mismatching between the alignment files generated by EMAN in @@ -331,20 +339,22 @@ def genMdObjDict(self, inTsSet, inCtfSet, coords): ctf=ctfIdsDict[tomoId], coords=iCoords, particles=iParticles, - jsonFile=genJsonFileName(self.getInfoDir(), tomoId)) + jsonFile=genJsonFileName(self.getInfoDir(), tomoId), + processingInd=procInd) + procInd += len(iCoords) if not mdObjDict: raise Exception("There isn't any common tilt series among the coordinates, the CTF tomo series and " "the tilt series chosen.") - return mdObjDict + return dict(sorted(mdObjDict.items())) - def _genExtractArgs(self, mdObjDict): + def _genExtractArgs(self, mdObj): # align2dFile = self.getNewAliFile(is3d=False) align3dFile = self.getNewAliFile() - tomoFiles = [join(TOMOGRAMS_DIR, basename(mdObj.inTomo.getFileName())) for mdObj in mdObjDict.values()] + tomoFiles = join(TOMOGRAMS_DIR, basename(mdObj.inTomo.getFileName())) args = [ # The particles can be read using the input tomograms or directly the particles contained in the ali3d file # in clase it exists - f'--jsonali={align3dFile}' if exists(self._getExtraPath(align3dFile)) else f'{" ".join(tomoFiles)}', + f'--jsonali={align3dFile}' if exists(self._getExtraPath(align3dFile)) else f'{tomoFiles}', f'--boxsz_unbin={self.getBoxSize()}', f'--shrink={self.shrink.get():.2f}', f'--maxtilt={self.maxTilt.get()}', @@ -393,14 +403,14 @@ def getProjections(self, mdObj): # lists of lists, as expected return projList - def createLandmarkModelGaps(self, ts, fiducialSize): + def createLandmarkModelGaps(self, ts): tsId = ts.getTsId() landmarkModelGapsFilePath = self._getExtraPath(str(tsId) + "_gaps.sfid") landmarkModelGaps = LandmarkModel(tsId=tsId, tiltSeriesPointer=ts, fileName=landmarkModelGapsFilePath, modelName=None, - size=fiducialSize, + size=self.fiducialSize, applyTSTransformation=False) landmarkModelGaps.setTiltSeries(ts) return landmarkModelGaps @@ -409,9 +419,76 @@ def createLandmarkModelGaps(self, ts, fiducialSize): def fillLandmarkModel(landmarkModelGaps, tsProjections, particle3dInd, nImgs): ind = particle3dInd * nImgs for i in range(ind, ind + nImgs): - projection = tsProjections[i] - tiltId = projection[1] + 1 - partId = particle3dInd + 1 - xCoord = round(projection[2]) - yCoord = round(projection[3]) - landmarkModelGaps.addLandmark(xCoord, yCoord, tiltId, partId, 0, 0) \ No newline at end of file + try: + # Found cases in which some 2d tilt images seem to have been excluded, e.g., TS with 58 images, 819 + # coordinates, and the stack 2d is of size 48251, which is not divisible by 59 (48251 / 59 = 817.8135). + # But in the end it's not critical as this 2d landmark models are only used for visualization and + # result checking purposes + projection = tsProjections[i] + tiltId = projection[1] + 1 + partId = particle3dInd + 1 + xCoord = round(projection[2]) + yCoord = round(projection[3]) + landmarkModelGaps.addLandmark(xCoord, yCoord, tiltId, partId, 0, 0) + except Exception as e: + continue + + def getOutSetOfParticles(self): + outParticles = getattr(self, self._possibleOutputs.subtomograms.name, None) + if outParticles: + outParticles.enableAppend() + else: + inTsPointer = getattr(self, IN_TS) + inParticlesPointer = getattr(self, IN_SUBTOMOS) + outParticles = EmanSetOfParticles.create(self._getPath(), template='emanParticles%s.sqlite') + if isinstance(inParticlesPointer.get(), SetOfSubTomograms): + outParticles.copyInfo(inParticlesPointer.get()) + else: + # The input are coordinates + outParticles.setCoordinates3D(inParticlesPointer) + outParticles.setSamplingRate(self.currentSRate) + # Define outputs and relations + self._defineOutputs(**{self._possibleOutputs.subtomograms.name: outParticles}) + self._defineSourceRelation(getattr(self, IN_SUBTOMOS), outParticles) + self._defineSourceRelation(getattr(self, IN_CTF), outParticles) + self._defineSourceRelation(inTsPointer, outParticles) + + return outParticles + + def getOutSetIfFiducials(self): + outFiducials = getattr(self, self._possibleOutputs.projected2DCoordinates.name, None) + if outFiducials: + outFiducials.enableAppend() + else: + inTsPointer = getattr(self, IN_TS) + inTs = inTsPointer.get() + outFiducials = SetOfLandmarkModels.create(self._getPath(), suffix='Gaps') + outFiducials.copyInfo(inTs) + outFiducials.setSetOfTiltSeries(inTsPointer) + # Define outputs and relations + self._defineOutputs(**{self._possibleOutputs.projected2DCoordinates.name: outFiducials}) + self._defineSourceRelation(inTsPointer, outFiducials) + return outFiducials + + def getEmanCoordsFromHdfStack(self, tsId): + fName = self._getEmanFName(tsId) + # Unstack the 3d particles HDF stack into individual MRC files + stack3d = abspath(join(self.getStack3dDir(), fName)) + eh = EmanHdf5Handler(stack3d) + return eh.get3dCoordsFrom3dStack(invertZ=self.doFlipZInTomo.get()) + + def getEmanMatchingCoord(self, emanCoord, scipionCoords): + def euclideanDist(coord1, coord2): + return np.linalg.norm(np.array(coord1) * self.scaleFactor - np.array(coord2)) + + distances = [euclideanDist(scipionCoord.getPosition(SCIPION), emanCoord) for scipionCoord in scipionCoords] + minDistInd = np.argmin(distances) + logger.info(f'MIN DIST = {distances[minDistInd]}') + matchingCoord = scipionCoords[minDistInd] + scaledCoords = np.array(matchingCoord.getPosition(SCIPION)) * self.scaleFactor + matchingCoord.setPosition(scaledCoords[0], scaledCoords[1], scaledCoords[2], SCIPION) + # Remove the matching coord from the list, so in an iterative project, the list becomes smaller with each + # iteration, being more efficient + reducedList = scipionCoords[:minDistInd] + scipionCoords[minDistInd+1:] + return matchingCoord, reducedList + diff --git a/emantomo/protocols/protocol_refine_multi_new.py b/emantomo/protocols/protocol_refine_multi_new.py index 5912cbf..a75e200 100644 --- a/emantomo/protocols/protocol_refine_multi_new.py +++ b/emantomo/protocols/protocol_refine_multi_new.py @@ -127,7 +127,7 @@ def _defineParams(self, form): form.addParam('maxShift', IntParam, default=-1, condition='doAlignment', - label='Maximum shift (pix.)', + label='Maximum shift (px)', help='If set to -1, it will be estimated as maxShift=boxSize/6.') form.addSection(label='Extra params') diff --git a/emantomo/protocols/protocol_refine_new.py b/emantomo/protocols/protocol_refine_new.py index 4978462..32be304 100644 --- a/emantomo/protocols/protocol_refine_new.py +++ b/emantomo/protocols/protocol_refine_new.py @@ -148,7 +148,7 @@ def _defineParams(self, form): form.addParam('maxShift', IntParam, default=-1, condition='doLocalRefine', - label='Maximum shift (pix.)', + label='Maximum shift (px)', help='If set to -1, it will be estimated as maxShift= boxSize/6.') group = form.addGroup('Motion correction', condition='doLocalRefine') group.addParam('smooth', IntParam, diff --git a/emantomo/protocols/protocol_template_matching.py b/emantomo/protocols/protocol_template_matching.py index 74e28c7..2398515 100644 --- a/emantomo/protocols/protocol_template_matching.py +++ b/emantomo/protocols/protocol_template_matching.py @@ -141,6 +141,9 @@ def templateMatchingStep(self): def createOutputStep(self): jsons2SetCoords3D(self, self.inTomos, self.getInfoDir()) + # Throw an exception if no coordinates were registered + if len(getattr(self, 'coordinates', '')) == 0: + raise Exception('ERROR!!! No coordintes were registered.') # --------------------------- UTILS functions ----------------------------- def _genTempMatchArgs(self): diff --git a/emantomo/protocols/protocol_tomo_subtomogram_refinement.py b/emantomo/protocols/protocol_tomo_subtomogram_refinement.py index 38652b7..be86810 100644 --- a/emantomo/protocols/protocol_tomo_subtomogram_refinement.py +++ b/emantomo/protocols/protocol_tomo_subtomogram_refinement.py @@ -39,8 +39,9 @@ import pwem.constants as emcts from pyworkflow.utils import createLink from tomo.protocols import ProtTomoBase -from tomo.objects import AverageSubTomogram, SetOfSubTomograms +from tomo.objects import AverageSubTomogram, SetOfSubTomograms, Coordinate3D from ..constants import SUBTOMOGRAMS_DIR, SPT_00_DIR, SYMMETRY_HELP_MSG +from ..utils import getPresentPrecedents class EmanTomoRefinementOutputs(enum.Enum): @@ -150,33 +151,38 @@ def _defineParams(self, form): # --------------- INSERT steps functions ---------------- def _insertAllSteps(self): - self.whole3dStack = self._getExtraPath(SUBTOMOGRAMS_DIR, 'whole3dstack.hdf') + self._initialize() self._insertFunctionStep(self.convertInputStep) self._insertFunctionStep(self.refinementSubtomogram) self._insertFunctionStep(self.convertOutputStep) self._insertFunctionStep(self.createOutputStep) # --------------- STEPS functions ----------------------- - def convertInputStep(self): - storePath = self._getExtraPath(SUBTOMOGRAMS_DIR) - pwutils.makePath(storePath) + def _initialize(self): + self.whole3dStack = self._getExtraPath(SUBTOMOGRAMS_DIR, 'whole3dstack.hdf') if self.useAlign.get(): self.alignType = emcts.ALIGN_3D else: self.alignType = emcts.ALIGN_NONE - writeSetOfSubTomograms(self.inputSetOfSubTomogram.get(), storePath, alignType=self.alignType) - fileList = glob(join(storePath, '*.hdf')) - # Fix the sampling rate as it might be set wrong + def convertInputStep(self): inSubtomos = self.inputSetOfSubTomogram.get() + storePath = self._getExtraPath(SUBTOMOGRAMS_DIR) + pwutils.makePath(storePath) + writeSetOfSubTomograms(inSubtomos, storePath, alignType=self.alignType) + + # Fix the sampling rate as it might be set wrong sampling_rate = inSubtomos.getSamplingRate() program = emantomo.Plugin.getProgram('e2proc3d.py') - for file in fileList: + for file in sorted(glob(join(storePath, '*.hdf'))): # It's critical that the HDF stack keep the sqlite order + # to later generate the output correctly in terms of transformation matrix assignment + # Bypass the problematic lst generation with this classical approach by making one HDF file containing all # the particles from all the tomograms args = "--apix %f %s %s --append" % (sampling_rate, file, self.whole3dStack) self.runJob(program, args) remove(file) + # The same with the reference volume self.refFileIn = self.inputRef.get().getFileName() self.refFileOut = self._getExtraPath('refVol.hdf') diff --git a/emantomo/protocols/protocol_ts_align_and_tomo_rec.py b/emantomo/protocols/protocol_ts_align_and_tomo_rec.py index 11610f2..6fc060d 100644 --- a/emantomo/protocols/protocol_ts_align_and_tomo_rec.py +++ b/emantomo/protocols/protocol_ts_align_and_tomo_rec.py @@ -35,6 +35,7 @@ from emantomo.objects import EmanMetaData, EmanHdf5Handler from emantomo.protocols.protocol_base import ProtEmantomoBase, IN_TS from emantomo.utils import genJsonFileName, getPresentTsIdsInSet +from pwem import ALIGN_2D, ALIGN_NONE from pwem.convert.headers import fixVolume from pwem.emlib.image import ImageHandler from pwem.objects import Transform @@ -66,14 +67,14 @@ class EmanProtTsAlignTomoRec(ProtEmantomoBase): """ This protocol wraps *e2tomogram.py* EMAN2 program. - Tomogram tiltseries alignment. + Tilt series alignment and tomogram reconstruction. Tomograms are not normally reconstructed at full resolution, generally limited to 1k x 1k or 2k x 2k, but the tilt-series are aligned at full resolution. For high resolution subtomogram averaging, the raw tilt-series data is used, based on coordinates from particle picking in the downsampled tomograms. On a typical workstation reconstruction takes about 4-5 minutes per tomogram. """ - _label = 'TS align & tomo rec' + _label = 'Alignment and reconstruction' _possibleOutputs = outputObjects def __init__(self, **kwargs): @@ -109,20 +110,27 @@ def _defineParams(self, form): 'generated in HDF and then converted into MRC. The HDF files are deleted by default to ' 'save storage.') - form.addSection(label='TS alignment') + form.addSection(label='Tilt series alignment') form.addParam('doAlignment', BooleanParam, default=True, label='Align the tilt series?') form.addParam('genInterpolatedTs', BooleanParam, - label='Generated interpolated TS?', + label='Generated the interpolated tilt series?', default=False, condition=alignCond, help='If set to Yes, an additional set of tilt series with the transformation matrix applied ' - 'will be generated. It can be used to check if the alignment was correctly calculated.') + 'will be generated. It can be used to check if the alignment was correctly calculated.\n\n' + '*NOTE*: Depending on the combination of parameters "Align the tilt series" and ' + '"Reconstruct the tomogram", the interpolated tilt series generated may be ' + 'different in terms of size and sampling rate.\n' + 'The interpolated tilt series generated when the tomogram is requested to be reconstructed ' + 'is different from the one generated if only align TS is selected (It is part of ' + 'the reconstruction functionality in the first case, and a temp file part of the alignment ' + 'step in the second case).') form.addParam('nLandmarks', IntParam, default=20, condition=alignCond, - label='Number of landmarks to use') + label='No. landmarks to use') form.addParam('pkKeep', FloatParam, default=0.9, label='Fraction of landmarks to keep in the tracking', @@ -140,24 +148,23 @@ def _defineParams(self, form): 'tracking iterations.') form.addParam('boxSizeTrk', IntParam, default=32, - label='Box size of the particles for tracking (pix.)', + label='Box size of the particles for tracking (px)', condition=alignCond, + expertLevel=LEVEL_ADVANCED, help='It may be helpful to use a larger one for fiducial-less cases.') form.addParam('letEmanEstimateTilts', BooleanParam, - default=True, + default=False, label='Should Eman estimate the tilt angles?', condition=alignCond, + expertLevel=LEVEL_ADVANCED, help='If set to No, a .tlt file will be generated containing the tilt angles read from Scipion ' 'imported tilt series metadata and passed to Eman. Default=True will let Eman to estimate ' 'the tilt angles, following the Eman original default behavior.') - form.addParam('tltStep', FloatParam, - default=2, - label='Step between tilts', - condition=f'{alignCond} and letEmanEstimateTilts') form.addParam('zeroId', IntParam, default=-1, label='Index of the center tilt', - condition=f'{alignCond} and letEmanEstimateTilts') + condition=f'{alignCond} and letEmanEstimateTilts', + help='If set -1, it will be estimated by EMAN.') form.addParam('tiltAxisAngle', FloatParam, allowsNull=True, label='Tilt axis angle', @@ -167,7 +174,7 @@ def _defineParams(self, form): 'contained in the metadata, it will be estimated by EMAN.') form.addParam('writeIntermediateRes', BooleanParam, default=False, - condition='%s and not genInterpolatedTs' % alignCond, + condition=alignCond, expertLevel=LEVEL_ADVANCED, label='Write intermediate results?', help='They will be generated always the interpolated tilt series are requested.') @@ -181,19 +188,22 @@ def _defineParams(self, form): default=SIZE_1K, choices=OUT_TOMO_SIZE_CHOICES, condition=recCond, - label='Size of output tomograms') + label='Size of the output tomograms') form.addParam('nIters', StringParam, default='2,1,1,1', condition=recCond, - label='No. iterations for bin8, bin4, bin2 images') + expertLevel=LEVEL_ADVANCED, + label='No. iterations for 500, 1K, 2K, and 8K images') form.addParam('clipz', IntParam, default=-1, condition=recCond, - label='Thickness (pix.)', + label='Thickness (px)', help='Z thickness of the final tomogram output. default is -1, (5/16 of tomogram length).') form.addParam('tltkeep', FloatParam, default=0.9, label='Fraction of tilts to keep in the reconstruction', + expertLevel=LEVEL_ADVANCED, + condition=recCond, validators=[GT(0), LE(1)]) form.addParam('bytile', BooleanParam, default=True, @@ -211,7 +221,13 @@ def _defineParams(self, form): default=False, label='Correct rotation', condition=recCond, + expertLevel=LEVEL_ADVANCED, help='Correct for global rotation and position sample flat in tomogram.') + form.addParam('correctXDrift', BooleanParam, + default=False, + condition=recCond, + label='Correct drifting along the X axis?', + expertLevel=LEVEL_ADVANCED) form.addParam('filterres', FloatParam, default=40, condition=recCond, @@ -239,7 +255,7 @@ def _insertAllSteps(self): for mdObj in mdObjList: self._insertFunctionStep(self.convertTsStep, mdObj) self._insertFunctionStep(self.writeData2JsonFileStep, mdObj) - self._insertFunctionStep(self.emanStep, mdObj.tsId) + self._insertFunctionStep(self.emanStep, mdObj) self._insertFunctionStep(self.convertOutputStep, mdObj) self._insertFunctionStep(self.createOutputStep, mdObj) self._insertFunctionStep(self.closeOutputSetsStep) @@ -269,20 +285,20 @@ def _initialize(self): counter += 1 return mdObjList - def convertTsStep(self, mObj): - ts = mObj.ts - inFile = ts.getFirstItem().getFileName() - outFile = mObj.tsId + getExt(inFile) - createLink(inFile, self._getExtraPath(outFile)) - program = emantomo.Plugin.getProgram("e2import.py") - cmd = [ - f'{outFile}', - f'--apix={ts.getSamplingRate():.2f}', - '--import_tiltseries', - f'--compressbits={self.compressBits.get()}', - '--importation=copy' - ] - self.runJob(program, ' '.join(cmd), cwd=self._getExtraPath()) + # def convertTsStep(self, mObj): + # ts = mObj.ts + # inFile = ts.getFirstItem().getFileName() + # outFile = mObj.tsId + getExt(inFile) + # createLink(inFile, self._getExtraPath(outFile)) + # program = emantomo.Plugin.getProgram("e2import.py") + # cmd = [ + # f'{outFile}', + # f'--apix={ts.getSamplingRate():.2f}', + # '--import_tiltseries', + # f'--compressbits={self.compressBits.get()}', + # '--importation=copy' + # ] + # self.runJob(program, ' '.join(cmd), cwd=self._getExtraPath()) def writeData2JsonFileStep(self, mdObj): # Generate initial json file required by EMAN for the reconstruction considering a previous alignment @@ -290,30 +306,33 @@ def writeData2JsonFileStep(self, mdObj): # This only applies to the reconstruction ts2Json(mdObj, mode='w') - def emanStep(self, tsId): + def emanStep(self, mdObj): + tsId = mdObj.tsId program = emantomo.Plugin.getProgram("e2tomogram.py") cmd = [self.getCommonArgs(tsId)] if self.doAlignment.get(): - cmd.append(self._getAlignArgs(tsId)) + cmd.append(self._getAlignArgs(mdObj)) if self.doRec.get(): cmd.append(self._getRecArgs()) + else: + cmd.append('--dryrun') self.runJob(program, ' '.join(cmd), cwd=self._getExtraPath()) def convertOutputStep(self, mdObj): + # Create a link that will be the non-interpolated TS output file name + ts = mdObj.ts tsId = mdObj.tsId - hdfTs = self.getCurrentHdfFile(TS_DIR, tsId) - hdfTomo = self.getCurrentHdfFile(TOMOGRAMS_DIR, tsId) - filesToConvert = [hdfTs, hdfTomo] - if self.genInterpolatedTs.get(): - tsHdfInterpFinalLoc = self.getTsInterpFinalLoc(tsId) - createLink(self._getInterpTsFile(mdObj), tsHdfInterpFinalLoc) - filesToConvert.append(tsHdfInterpFinalLoc) - for hdfFile in filesToConvert: - eh = EmanHdf5Handler(hdfFile) + inFile = ts.getFirstItem().getFileName() + outFile = self.getOutTsNonInterpFName(tsId) + createLink(inFile, outFile) + # Convert the tomogram from HDF format into MRC + if self.doRec.get(): + outFName = self.getOutTomoFName(mdObj.tsId) + hdfTomo = self.getCurrentHdfFile(TOMOGRAMS_DIR, mdObj.tsId) + eh = EmanHdf5Handler(hdfTomo) extraArgs = '--apix %.3f ' % eh.getSamplingRate() - convertBetweenHdfAndMrc(self, hdfFile, replaceExt(hdfFile, 'mrc'), extraArgs=extraArgs) - # Fix the converted tomogram file headerTomograms - fixVolume(replaceExt(hdfTomo, 'mrc')) + convertBetweenHdfAndMrc(self, hdfTomo, outFName, extraArgs=extraArgs) + fixVolume(outFName) def createOutputStep(self, mdObj): # TS alignment @@ -321,19 +340,19 @@ def createOutputStep(self, mdObj): jsonDict = loadJson(mdObj.jsonFile) aliLoss = jsonDict[ALI_LOSS] alignParams = jsonDict[TLT_PARAMS] - self._createOutputTs(mdObj, alignParams, aliLoss) + tsNonInterp = self._createOutputTs(mdObj, alignParams, aliLoss) if self.genInterpolatedTs.get(): - self._createOutputTsInterpolated(mdObj, alignParams) + self._createOutputTsInterpolated(mdObj, alignParams, tsNonInterp) # Tomogram reconstruction if self.doRec.get(): tsId = mdObj.tsId eh = EmanHdf5Handler(self.getCurrentHdfFile(TOMOGRAMS_DIR, tsId)) tomogram = Tomogram() - tomogram.setFileName(self._getOutTomoName(tsId)) + tomogram.setFileName(self.getOutTomoFName(tsId)) tomogram.setSamplingRate(eh.getSamplingRate()) tomogram.setTsId(tsId) - acq = mdObj.ts.getAcquisition() + acq = mdObj.ts.getAcquisition().clone() if self._doUpdateTiltAxisAng: acq.setTiltAxisAngle(self._tiltAxisAngle) tomogram.setAcquisition(acq) @@ -342,8 +361,8 @@ def createOutputStep(self, mdObj): outTomoSet = self.getOutputSetOfTomograms() outTomoSet.append(tomogram) outTomoSet.update(tomogram) - - self._store() + outTomoSet.write() + self._store(outTomoSet) def closeOutputSetsStep(self): outPutDict = {} @@ -412,7 +431,7 @@ def getTAxisAngle(self): self._tiltAxisAngle = ts.getAcquisition().getTiltAxisAngle() return self._tiltAxisAngle - def _getAlignArgs(self, tsId): + def _getAlignArgs(self, mdObj): tAx = self.getTAxisAngle() args = [ f'--npk={self.nLandmarks.get()}', @@ -423,11 +442,11 @@ def _getAlignArgs(self, tsId): ] if self.letEmanEstimateTilts.get(): # Introduce angular step and zero id - args.append(f'--tltstep={self.tltStep.get():.2f}') + args.append(f'--tltstep={mdObj.ts.getAcquisition().getStep()}') args.append(f'--zeroid={self.zeroId.get()}') else: - args.append(f'--rawtlt={join(TLT_DIR, tsId + ".tlt")}') - if not self.writeIntermediateRes.get() and not self.genInterpolatedTs.get(): + args.append(f'--rawtlt={join(TLT_DIR, mdObj.tsId + ".tlt")}') + if not self.writeIntermediateRes.get(): args.append('--notmp') if tAx: args.append(f'--tltax={tAx:.2f}') @@ -436,27 +455,26 @@ def _getAlignArgs(self, tsId): def getOutputSetOfTs(self, interpolated=False): inTs = self.inputTS.get() - outTs = self.getTiltSeries(interpolated=interpolated) - if outTs: - outTs.enableAppend() - tiltSeries = outTs + outTsSet = self.getTiltSeries(interpolated=interpolated) + if outTsSet: + outTsSet.enableAppend() else: + suffix = "interpolated" if interpolated else "" + outTsSet = SetOfTiltSeries.create(self._getPath(), template='tiltseries', suffix=suffix) + outTsSet.copyInfo(inTs) if interpolated: - suffix = 'interpolated' - x, y, z, _ = ImageHandler().getDimensions(glob.glob(self._getExtraPath(TS_DIR, '*_interp.mrc'))[0]) - dims = (x, y, z) + alignment = ALIGN_NONE + acq = inTs.getAcquisition() + acq.setTiltAxisAngle(0.) + outTsSet.setAcquisition(acq) else: - suffix = '' - dims = inTs.getDim() + alignment = ALIGN_2D - tiltSeries = SetOfTiltSeries.create(self._getPath(), template='tiltseries', suffix=suffix) - tiltSeries.copyInfo(self.inputTS.get()) - tiltSeries.setDim(dims) - tiltSeries.setSamplingRate(inTs.getSamplingRate()) - tiltSeries.setStreamState(Set.STREAM_OPEN) - self.setTiltSeries(tiltSeries, interpolated=interpolated) + outTsSet.setAlignment(alignment) + outTsSet.setStreamState(Set.STREAM_OPEN) + self.setTiltSeries(outTsSet, interpolated=interpolated) - return tiltSeries + return outTsSet def getTiltSeries(self, interpolated=False): if interpolated: @@ -472,39 +490,59 @@ def setTiltSeries(self, tsSet, interpolated=False): def _createOutputTs(self, mdObj, alignParams, aliLoss): outTsSet = self.getOutputSetOfTs() - tiltSeries = self._createCurrentOutTs(mdObj.ts) + tiltSeries = self._createCurrentOutTs(mdObj) outTsSet.append(tiltSeries) + outFileName = self.getOutTsNonInterpFName(mdObj.tsId) for idx, ti in enumerate(mdObj.ts.iterItems()): - outTi = TiltImage() - outTi.copyInfo(ti, copyId=True) - outTi.setIndex(ti.getIndex()) - outTi.setFileName(self._getExtraPath(TS_DIR, mdObj.tsId + '.mrc')) + outTi = ti.clone() + transform = Transform() + outTi.setFileName(outFileName) curerntAlignParams = alignParams[idx] setattr(outTi, EMAN_ALI_LOSS, Float(aliLoss[idx])) - self._genTrMatrixFromEmanAlign(outTi, curerntAlignParams) + # Transformation matrix + transformMatrix = self._genTrMatrixFromEmanAlign(curerntAlignParams) + transform.setMatrix(transformMatrix) + outTi.setTransform(transform) + tiltAngleRefined = curerntAlignParams[3] + outTi.setTiltAngle(tiltAngleRefined) + offTiltAngle = curerntAlignParams[4] + setattr(outTi, EMAN_OFF_TILT_AXIS, Float(offTiltAngle)) # Extended parameter # Append the current tilt image to the corresponding tilt series tiltSeries.append(outTi) outTsSet.update(tiltSeries) - return outTsSet + outTsSet.write() + self._store(outTsSet) + return tiltSeries - def _createOutputTsInterpolated(self, mdObj, alignParams): + def _createOutputTsInterpolated(self, mdObj, alignParams, tsNonInterp): outTsSet = self.getOutputSetOfTs(interpolated=True) - tiltSeries = self._createCurrentOutTs(mdObj.ts, interpolated=True) + tiltSeries = self._createCurrentOutTs(mdObj, interpolated=True) outTsSet.append(tiltSeries) - for idx, ti in enumerate(mdObj.ts.iterItems()): - outTi = TiltImage() - finalName = self.getTsInterpFinalLoc(mdObj.tsId).replace('.hdf', '.mrc') - outTi.copyInfo(ti, copyId=True) - outTi.setTsId(outTi.getTsId()) - outTi.setIndex(ti.getIndex()) + # Apply the transformation for the non-interpolated tilt-series + rotationAngle = tsNonInterp.getAcquisition().getTiltAxisAngle() + doSwap = 45 < abs(rotationAngle) < 135 + finalName = tsNonInterp.getFirstItem().getFileName().replace('.mrc', '_interpolated.mrc') + tsNonInterp.applyTransform(finalName, swapXY=doSwap) + for idx, ti in enumerate(mdObj.ts): + outTi = ti.clone() outTi.setFileName(finalName) - setattr(outTi, EMAN_OFF_TILT_AXIS, Float(alignParams[idx][4])) # Off tilt axis angle, extended parameter - outTi.setTiltAngle(alignParams[idx][3]) # Refined tilt angle + # Manage the acquisition + acq = outTi.getAcquisition() + acq.setTiltAxisAngle(0.) # 0 because it is aligned + ti.setAcquisition(acq) + # Set the refined tilt angles + curerntAlignParams = alignParams[idx] + tiltAngleRefined = curerntAlignParams[3] + outTi.setTiltAngle(tiltAngleRefined) tiltSeries.append(outTi) - outTsSet.setAlignment2D() + acq = tiltSeries.getAcquisition() + acq.setTiltAxisAngle(0.) # 0 because the TS is aligned + tiltSeries.setAcquisition(acq) outTsSet.update(tiltSeries) + outTsSet.write() + self._store(outTsSet) return outTsSet @staticmethod @@ -517,29 +555,38 @@ def genUpdateTsInterpOrigin(dims, samplingRate): 0) return origin - def _createCurrentOutTs(self, ts, interpolated=False): - tiltSeries = TiltSeries() + def _createCurrentOutTs(self, mdObj, interpolated=False): + ts = mdObj.ts + tsId = ts.getTsId() + tiltSeries = TiltSeries(tsId=tsId) + tiltSeries.copyInfo(ts) if interpolated: + tiltSeries.setInterpolated(True) + tiltSeries.setAlignment(ALIGN_NONE) # Avoid the copyInfo and set the origin manually to get the squared TS dims and origin correctly stored - tsId = ts.getTsId() - sr = ts.getSamplingRate() - tiltSeries.setSamplingRate(sr) - tiltSeries.setAcquisition(ts.getAcquisition()) - tiltSeries.setTsId(tsId) - x, y, z, _ = ImageHandler().getDimensions(self._getExtraPath(TS_DIR, tsId + '.mrc')) - dims = (x, y, z) - tiltSeries.setDim(dims) - tiltSeries.setOrigin(self.genUpdateTsInterpOrigin(dims, sr)) - else: - tiltSeries.copyInfo(ts) - if self._doUpdateTiltAxisAng: acq = ts.getAcquisition() - acq.setTiltAxisAngle(self._tiltAxisAngle) + acq.setTiltAxisAngle(0.) tiltSeries.setAcquisition(acq) + # x, y, z, _ = ImageHandler().getDimensions(self._getExtraPath(TS_DIR, tsId + '.mrc')) + # dims = (x, y, z) + # tiltSeries.setDim(dims) + # tiltSeries.setOrigin(self.genUpdateTsInterpOrigin(dims, ts.getSamplingRate())) + # # Set the interpolated TS sampling rate + # tiltSeries.setSamplingRate(interpSRate) + else: + tiltSeries.setAlignment2D() + if self._doUpdateTiltAxisAng: + acq = ts.getAcquisition() + acq.setTiltAxisAngle(self._tiltAxisAngle) + tiltSeries.setAcquisition(acq) return tiltSeries def _getInterpTsFile(self, mdObj): - matchingFiles = glob.glob(self._getExtraPath('tomorecon_%02d' % mdObj.processingInd, INTERP_TS + '.hdf')) + # NOTE: The interpolated TS generated when the tomogram is requested to be reconstructed is different from + # the one generated if only align TS is selected (Seems to be part of the reconstruction functionality in the + # first case, and a temp file part of the alignment step in the second case). + interpBName = INTERP_TS if self.doRec.get() else 'tltseries_transali' + matchingFiles = glob.glob(self._getExtraPath('tomorecon_%02d' % mdObj.processingInd, interpBName + '.hdf')) return max(matchingFiles, key=getctime) # --------------------------- reconstruct tomograms UTILS functions ---------------------------- @@ -570,6 +617,9 @@ def _getRecArgs(self): if self.extrapad.get(): args.append('--extrapad') + if self.correctXDrift.get(): + args.append('--xdrift') + return ' '.join(args) # def getFinalSampligRate(self): @@ -582,8 +632,8 @@ def _getRecArgs(self): # self._finalSamplingRate = binning * tsSet.getSamplingRate() # return self._finalSamplingRate - def _getOutTomoName(self, tsId): - return glob.glob(self._getExtraPath(TOMOGRAMS_DIR, '%s*.mrc' % tsId))[0] + # def _getOutTomoName(self, tsId): + # return glob.glob(self._getExtraPath(TOMOGRAMS_DIR, '%s*.mrc' % tsId))[0] def getOutputSetOfTomograms(self): outTomograms = self.getTomograms() @@ -614,15 +664,13 @@ def getTsInterpFinalLoc(self, tsId): return self._getExtraPath(TS_DIR, tsId + '_interp.hdf') @staticmethod - def _genTrMatrixFromEmanAlign(outTi, alignParams): - transform = Transform() + def _genTrMatrixFromEmanAlign(alignParams): matrix = np.eye(3) rotMatrix = np.eye(2) # tlt_params: (N, 5) list, where N is the number of tilt in the tilt series. The columns are translation # along x,y axis, and rotation around z, y, x axis in the EMAN coordinates. The translation is in unbinned # pixels, and rotation is in degrees rotAngle = -np.deg2rad(alignParams[2]) - tiltAngleRefined = alignParams[3] offTiltAngle = alignParams[4] tx = -alignParams[0] ty = -alignParams[1] @@ -637,10 +685,7 @@ def _genTrMatrixFromEmanAlign(outTi, alignParams): matrix[1, 0] = rotMatrix[1, 0] matrix[0, 2] = shiftsImod[0] matrix[1, 2] = shiftsImod[1] - transform.setMatrix(matrix) - outTi.setTransform(transform) - outTi.setTiltAngle(tiltAngleRefined) - setattr(outTi, EMAN_OFF_TILT_AXIS, Float(offTiltAngle)) # Extended parameter + return matrix def doAutoclipXY(self): """Code behavior expected for EMAN's option autoclipxy to be automatic. @@ -649,3 +694,11 @@ def doAutoclipXY(self): ih = ImageHandler() x, y, _, _ = ih.getDimensions(self.inTsSet.getFirstItem().getFirstItem().getFileName()) return True if x != y else False + + def getOutTsNonInterpFName(self, tsId): + return join(self.getTsDir(), tsId + '.mrc') + + def getOutTomoFName(self, tsId): + return join(self.getTomogramsDir(), tsId + '.mrc') + + diff --git a/emantomo/tests/test_protocols_sta_pppt.py b/emantomo/tests/test_protocols_sta_pppt.py index e6fde4f..ec25e9f 100644 --- a/emantomo/tests/test_protocols_sta_pppt.py +++ b/emantomo/tests/test_protocols_sta_pppt.py @@ -36,19 +36,22 @@ from pyworkflow.utils import magentaStr from reliontomo.protocols import ProtImportCoordinates3DFromStar from tomo.constants import TR_EMAN -from tomo.objects import TomoAcquisition, Coordinate3D, SetOfCoordinates3D +from tomo.objects import Coordinate3D, SetOfCoordinates3D from tomo.protocols import ProtImportTs, ProtImportTomograms, ProtImportTsCTF from tomo.protocols.protocol_import_ctf import ImportChoice from tomo.protocols.protocol_import_tomograms import OUTPUT_NAME from tomo.tests import RE4_STA_TUTO, DataSetRe4STATuto from tomo.tests.test_base_centralized_layer import TestBaseCentralizedLayer +TS_54 = 'TS_54' +TS_03 = 'TS_03' + class TestEmanBasePPPT(TestBaseCentralizedLayer): ds = None importedTs = None - particlesUnbinnedBoxSize = 256 particlesExtractedBoxSize = 64 + unbinnedSRate = DataSetRe4STATuto.unbinnedPixSize.value @classmethod def setUpClass(cls): @@ -56,12 +59,12 @@ def setUpClass(cls): cls.ds = DataSet.getDataSet(RE4_STA_TUTO) @classmethod - def _runImportTs(cls, exclusionWords=DataSetRe4STATuto.exclusionWordsTs03.value): + def _runImportTs(cls): print(magentaStr("\n==> Importing the tilt series:")) - protImportTs = cls.newProtocol(ProtImportTs, + protTsImport = cls.newProtocol(ProtImportTs, filesPath=cls.ds.getFile(DataSetRe4STATuto.tsPath.value), filesPattern=DataSetRe4STATuto.tsPattern.value, - exclusionWords=exclusionWords, + exclusionWords=DataSetRe4STATuto.exclusionWordsTs03ts54.value, anglesFrom=2, # From tlt file voltage=DataSetRe4STATuto.voltage.value, magnification=DataSetRe4STATuto.magnification.value, @@ -69,59 +72,109 @@ def _runImportTs(cls, exclusionWords=DataSetRe4STATuto.exclusionWordsTs03.value) amplitudeContrast=DataSetRe4STATuto.amplitudeContrast.value, samplingRate=DataSetRe4STATuto.unbinnedPixSize.value, doseInitial=DataSetRe4STATuto.initialDose.value, - dosePerFrame=DataSetRe4STATuto.dosePerTiltImg.value, + dosePerFrame=DataSetRe4STATuto.dosePerTiltImgWithTltFile.value, tiltAxisAngle=DataSetRe4STATuto.tiltAxisAngle.value) - cls.launchProtocol(protImportTs) - tsImported = getattr(protImportTs, 'outputTiltSeries', None) + cls.launchProtocol(protTsImport) + tsImported = getattr(protTsImport, protTsImport.OUTPUT_NAME, None) return tsImported class TestEmanTsAlignAndTomoRec(TestEmanBasePPPT): + nAnglesDict = DataSetRe4STATuto.nAnglesDict.value + nTiltSeries = len(nAnglesDict) + expectedDimsTs = DataSetRe4STATuto.dimsTsBin1Dict.value + tsAcqDict = DataSetRe4STATuto.tsAcqDict.value + expectedTomoThk = 250 + expectedTomoDimsDict = {TS_54: [1080, 1050, expectedTomoThk]} @classmethod def setUpClass(cls): super().setUpClass() cls.importedTs = cls._runImportTs() - def test_ts_align_tomo_rec(self): - print(magentaStr("\n==> Aligning the tilt series and reconstructing the tomograms:")) + def test_ts_align_tomo_rec_01(self): + print(magentaStr("\n==> Testing Emantomo Tilt-Series Alignment and Tomogram Reconstruction:" + "\n\t- Align the TS" + "\n\t- Generate the interpolated TS" + "\n\t- Reconstruct the tomograms")) protTsAlignTomoRec = self.newProtocol(EmanProtTsAlignTomoRec, inputTS=self.importedTs, - boxSizeTrk=100, - tltStep=3, - clipz=250, + boxSizeTrk=32, + clipz=self.expectedTomoThk, + genInterpolatedTs=True, + correctrot=True, + extrapad=True, numberOfThreads=8) + protTsAlignTomoRec.setObjLabel('Align & Interp & Rec') self.launchProtocol(protTsAlignTomoRec) + + # CHECK THE OUTPUTS + self.expectedTomoDimsDict[TS_03] = [1024, 1000, self.expectedTomoThk] tsAligned = getattr(protTsAlignTomoRec, protTsAlignTomoRec._possibleOutputs.tiltSeries.name, None) + tsInterp = getattr(protTsAlignTomoRec, protTsAlignTomoRec._possibleOutputs.tiltSeriesInterpolated.name, None) tomosRec = getattr(protTsAlignTomoRec, protTsAlignTomoRec._possibleOutputs.tomograms.name, None) - self.assertIsNotNone(tsAligned, "There was a problem aligning the tilt series") - self.assertIsNotNone(tomosRec, "There was a problem reconstructing the tomograms") + self.assertIsNotNone(tsAligned, msg="There was a problem aligning the tilt-series") + self.assertIsNotNone(tsInterp, msg="There was a problem with the interpolated the tilt-series") + self.assertIsNotNone(tomosRec, msg="There was a problem reconstructing the tomograms") # Check the tilt series - testAcq = TomoAcquisition(voltage=DataSetRe4STATuto.voltage.value, - sphericalAberration=DataSetRe4STATuto.sphericalAb.value, - amplitudeContrast=DataSetRe4STATuto.amplitudeContrast.value, - magnification=DataSetRe4STATuto.magnification.value, - tiltAxisAngle=DataSetRe4STATuto.tiltAxisAngle.value, - doseInitial=DataSetRe4STATuto.initialDose.value, - dosePerFrame=DataSetRe4STATuto.dosePerTiltImg.value, - accumDose=DataSetRe4STATuto.accumDose.value - ) self.checkTiltSeries(tsAligned, - expectedSetSize=1, - expectedSRate=DataSetRe4STATuto.unbinnedPixSize.value, - expectedDimensions=[3710, 3838, 40], - testAcqObj=testAcq, - alignment=ALIGN_2D, + expectedSetSize=self.nTiltSeries, + expectedSRate=self.unbinnedSRate, + expectedDimensions=self.expectedDimsTs, + testAcqObj=self.tsAcqDict, hasAlignment=True, - anglesCount=40, - ) + alignment=ALIGN_2D, + anglesCount=self.nAnglesDict) + # Check the interpolated tilt series + testAcqInterpDict = DataSetRe4STATuto.tsAcqInterpDict.value + self.checkTiltSeries(tsInterp, + expectedSetSize=self.nTiltSeries, + expectedSRate=self.unbinnedSRate, + expectedDimensions=DataSetRe4STATuto.dimsTsInterpBin1Dict.value, + testAcqObj=testAcqInterpDict, + anglesCount=self.nAnglesDict, + isInterpolated=True) + # Check the tomograms + self.checkTomograms(tomosRec, + expectedSetSize=self.nTiltSeries, + expectedSRate=DataSetRe4STATuto.unbinnedPixSize.value * 4, # Bin 4 + expectedDimensions=self.expectedTomoDimsDict) + + def test_ts_align_tomo_rec_02(self): + print(magentaStr("\n==> Testing Emantomo Tilt-Series Alignment and Tomogram Reconstruction:" + "\n\t- Only reconstruct the tomogram")) + # Import the transformation matrix + protImportTrMatrix = self.newProtocol(ProtImodImportTransformationMatrix, + filesPath=self.ds.getFile(DataSetRe4STATuto.tsPath.value), + filesPattern=DataSetRe4STATuto.transformPattern.value, + inputSetOfTiltSeries=self.importedTs) + self.launchProtocol(protImportTrMatrix) + outTsSet = getattr(protImportTrMatrix, OUTPUT_TILTSERIES_NAME, None) + # Reconstruct the tomograms + protTsAlignTomoRec = self.newProtocol(EmanProtTsAlignTomoRec, + inputTS=outTsSet, + doAlignment=False, + clipz=self.expectedTomoThk, + correctrot=True, + extrapad=True, + numberOfThreads=8) + protTsAlignTomoRec.setObjLabel('Tomo rec') + self.launchProtocol(protTsAlignTomoRec) + + # Check the results + self.expectedTomoDimsDict[TS_03] = [1050, 1000, self.expectedTomoThk] + tomosRec = getattr(protTsAlignTomoRec, protTsAlignTomoRec._possibleOutputs.tomograms.name, None) + tsAligned = getattr(protTsAlignTomoRec, protTsAlignTomoRec._possibleOutputs.tiltSeries.name, None) + tsInterp = getattr(protTsAlignTomoRec, protTsAlignTomoRec._possibleOutputs.tiltSeriesInterpolated.name, None) + self.assertIsNone(tsAligned, msg="There was a problem aligning the tilt-series") + self.assertIsNone(tsInterp, msg="There was a problem with the interpolated the tilt-series") + self.assertIsNotNone(tomosRec, msg="There was a problem reconstructing the tomograms") # Check the tomograms self.checkTomograms(tomosRec, - expectedSetSize=1, + expectedSetSize=self.nTiltSeries, expectedSRate=DataSetRe4STATuto.unbinnedPixSize.value * 4, # Bin 4 - expectedDimensions=[1050, 1000, 250] - ) + expectedDimensions=self.expectedTomoDimsDict) class TestEmanEstimateCtf(TestEmanBasePPPT): @@ -154,6 +207,7 @@ class TestBaseRefineCyclePPPT(TestEmanBasePPPT): importedTomos = None importedCoords = None extractedParticles = None + particlesUnbinnedBoxSize = 256 @classmethod def setUpClass(cls): @@ -162,7 +216,7 @@ def setUpClass(cls): @classmethod def _runPreviousProtocols(cls): - cls.importedTs = cls._runImportTs(exclusionWords=DataSetRe4STATuto.exclusionWordsTs03ts54.value) + cls.importedTs = cls._runImportTs() cls.tsWithAlignment = cls._runImportTrMatrix() cls.importedCtfs = cls._runImportCtfs() cls.importedTomos = cls._runImportTomograms() diff --git a/requirements.txt b/requirements.txt index ebf335e..c14d5a9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ scipion-em>=3.5.0 scipion-em-eman2 -scipion-em-tomo>=3.5.1 +scipion-em-tomo>=3.6.1 h5py==3.8.0 \ No newline at end of file