def parallelStart_splitVol(self,
job,
splitX,
splitY,
splitZ,
verbose=True):
"""
"""
self.parallelInit()
import pytom_mpi
mpi_myid = pytom_mpi.rank()
if mpi_myid == 0: # manager
# distribute the job to the workers and get the original size of the volume
self.splitVolumes(job, splitX, splitY, splitZ, verbose)
[vsizeX, vsizeY, vsizeZ] = self.jobInfo["originalSize"]
# gather results and post processing
from pytom_volume import vol
volume = vol(vsizeX, vsizeY, vsizeZ)
volume.setAll(0)
orient = vol(vsizeX, vsizeY, vsizeZ)
orient.setAll(0)
i = 0
while i < splitX * splitY * splitZ:
mpi_msgString = getMsgStr()
msg = self.getResMsg(mpi_msgString)
resFromWorker = self.resFromMsg(msg)
if verbose == True:
print("Manager: processing result from worker " +
msg.getSender())
resV = resFromWorker.result.getVolume()
resO = resFromWorker.orient.getVolume()
jobID = resFromWorker.jobID
[sizeX, sizeY, sizeZ] = self.jobInfo["splitSize"]
[sub_start, start] = self.jobInfo[jobID]
from pytom_volume import subvolume, putSubVolume
sub_resV = subvolume(resV, sub_start[0], sub_start[1],
sub_start[2], sizeX, sizeY, sizeZ)
sub_resO = subvolume(resO, sub_start[0], sub_start[1],
sub_start[2], sizeX, sizeY, sizeZ)
putSubVolume(sub_resV, volume, start[0], start[1], start[2])
putSubVolume(sub_resO, orient, start[0], start[1], start[2])
i = i + 1
# write the final result back to the disk
volume.write(self.name + '_res.em')
orient.write(self.name + '_orient.em')
# delete the temporary files on the disk
import os
files = os.listdir('.')
for name in files:
if 'job' in name and '.em' in name and not 'sum' in name and not 'sqr' in name:
os.remove(name)
# rename the result files name
os.rename('node_0_res.em', 'scores{}{}.em'.format(suffix))
os.rename('node_0_orient.em', 'angles{}{}.em'.format(suffix))
# finishing, stop all workers
self.parallelEnd(verbose)
if verbose == True:
print("Manager: end")
else: # worker
worker = PeakWorker()
worker.parallelRun(verbose)
pytom_mpi.finalise()
def maskOut(self, mask, center, size):
"""
maskOut: Set part of mask volume to all zero. The region is specified by center and size.
@param mask: volume that you handle with
@type mask: L{pytom_volume.vol}
@param center: center of the region
@type center: [x,y,z]
@param size: size of the region
@type size: [sizeX, sizeY, sizeZ] or radius
"""
from pytom_volume import vol, putSubVolume
if size.__class__ == list:
p_sizeX = size[0]
p_sizeY = size[1]
p_sizeZ = size[2]
elif size.__class__ == vol:
mm = size
p_sizeX = mm.sizeX()
p_sizeY = mm.sizeY()
p_sizeZ = mm.sizeZ()
else:
radius = size
p_sizeX = radius * 2
p_sizeY = radius * 2
p_sizeZ = radius * 2
maskSize = [mask.sizeX(), mask.sizeY(), mask.sizeZ()]
if maskSize < center:
raise RuntimeError('Center out of range!')
# [)
# mask out double size. CHANGED!!!
startX = int(center[0] - p_sizeX / 2)
endX = int(center[0] + p_sizeX / 2)
startY = int(center[1] - p_sizeY / 2)
endY = int(center[1] + p_sizeY / 2)
startZ = int(center[2] - p_sizeZ / 2)
endZ = int(center[2] + p_sizeZ / 2)
# only used for radius
sub_startX = 0
sub_startY = 0
sub_startZ = 0
if startX < 0:
sub_startX = -startX
startX = 0
if endX > maskSize[0]:
endX = maskSize[0]
if startY < 0:
sub_startY = -startY
startY = 0
if endY > maskSize[1]:
endY = maskSize[1]
if startZ < 0:
sub_startZ = -startZ
startZ = 0
if endZ > maskSize[2]:
endZ = maskSize[2]
sizeX = endX - startX
sizeY = endY - startY
sizeZ = endZ - startZ
if size.__class__ == list:
subV = vol(sizeX, sizeY, sizeZ)
subV.setAll(0)
elif size.__class__ == vol:
from pytom_volume import limit, subvolume
subV = (mm - 1) / -1
limit(subV, 0.999, 0, 0, 0, True, False)
subV = subvolume(subV, sub_startX, sub_startY, sub_startZ, sizeX,
sizeY, sizeZ)
tempV = subvolume(mask, startX, startY, startZ, sizeX, sizeY,
sizeZ)
subV = subV * tempV # AND operation
else:
from pytom_volume import initSphere, subvolume
subV = vol(radius * 2, radius * 2, radius * 2)
initSphere(subV, radius, 0, 0, radius, radius, radius)
tempV = vol(radius * 2, radius * 2, radius * 2)
tempV.setAll(1)
subV = tempV - subV
subV = subvolume(subV, sub_startX, sub_startY, sub_startZ, sizeX,
sizeY, sizeZ)
tempV = subvolume(mask, startX, startY, startZ, sizeX, sizeY,
sizeZ)
subV = subV * tempV # AND operation
putSubVolume(subV, mask, startX, startY, startZ)
def summarize(self, result, jobID, verbose=True):
"""
summarize: Get the result and do the summarization accordingly.\
(Send back, or update result, or write to the disk).
@param result: result
@type result: L{pytom.localization.peak_job.PeakResult}
@param jobID: ID of job
@type jobID: integer
"""
resV = result[0]
orientV = result[1]
if self.jobInfoPool == {}: # leaf node
assert self.backTo != None
#self.parallelEnd(verbose)
if not self.backTo is None:
result = self.writeRes(resV, orientV, jobID)
result.send(self.mpi_id, self.backTo)
else:
self.writeRes(resV, orientV, None)
self.parallelEnd(verbose)
return
# non leaf node, update the result
if self.jobInfoPool[jobID].splitType == "Ang":
self.jobInfoPool[
"numDoneJobsA"] = self.jobInfoPool["numDoneJobsA"] + 1
offset = self.jobInfoPool[jobID].angleOffset
# print self.name + ': JobID ' + str(jobID) + ' Offset ' + str(offset)
orientV = orientV + offset
if self.resVolA == None or self.resOrientA == None:
from pytom_volume import vol
self.resVolA = vol(resV.sizeX(), resV.sizeY(), resV.sizeZ())
self.resOrientA = vol(orientV.sizeX(), orientV.sizeY(),
orientV.sizeZ())
self.resVolA.copyVolume(resV)
self.resOrientA.copyVolume(orientV)
else:
from pytom_volume import updateResFromVol
updateResFromVol(self.resVolA, resV, self.resOrientA, orientV)
if self.jobInfoPool["numDoneJobsA"] == self.jobInfoPool[
"numJobsA"] and self.jobInfoPool["numJobsV"] > 0:
self.summarize([self.resVolA, self.resOrientA],
self.jobInfoPool[jobID].originalJobID, verbose)
return
elif self.jobInfoPool[jobID].splitType == "Vol":
self.jobInfoPool[
"numDoneJobsV"] = self.jobInfoPool["numDoneJobsV"] + 1
[originX, originY, originZ] = self.jobInfoPool[jobID].origin
if self.resVol == None or self.resOrient == None:
from pytom_volume import vol
[vsizeX, vsizeY, vsizeZ] = self.jobInfoPool[jobID].originalSize
self.resVol = vol(vsizeX, vsizeY, vsizeZ)
self.resVol.setAll(0)
self.resOrient = vol(vsizeX, vsizeY, vsizeZ)
self.resOrient.setAll(0)
[vsizeX, vsizeY, vsizeZ] = self.jobInfoPool[jobID].originalSize
[sizeX, sizeY, sizeZ] = self.jobInfoPool[jobID].splitSize
sub_start = self.jobInfoPool[jobID].sub_start
start = self.jobInfoPool[jobID].whole_start
stepSizeX = min(vsizeX - sub_start[0], sizeX)
stepSizeY = min(vsizeY - sub_start[1], sizeY)
stepSizeZ = min(vsizeZ - sub_start[2], sizeZ)
from pytom_volume import subvolume, putSubVolume
sub_resV = subvolume(resV, sub_start[0], sub_start[1],
sub_start[2], stepSizeX, stepSizeY, stepSizeZ)
sub_resO = subvolume(orientV, sub_start[0], sub_start[1],
sub_start[2], stepSizeX, stepSizeY, stepSizeZ)
putSubVolume(sub_resV, self.resVol, start[0] - originX,
start[1] - originY, start[2] - originZ)
putSubVolume(sub_resO, self.resOrient, start[0] - originX,
start[1] - originY, start[2] - originZ)
else:
raise RuntimeError("Unclear split type!")
# if all results are there, write back to disk and return to high level
if self.jobInfoPool["numDoneJobsA"] + self.jobInfoPool[
"numDoneJobsV"] == self.jobInfoPool[
"numJobsA"] + self.jobInfoPool["numJobsV"]:
if self.jobInfoPool["numJobsV"] == 0:
self.resVol = self.resVolA
self.resOrient = self.resOrientA
if self.backTo != None:
result = self.writeRes(self.resVol, self.resOrient,
self.jobInfoPool[jobID].originalJobID)
if verbose == True:
print(self.name + ': sending back result to ' +
str(self.backTo))
result.send(self.mpi_id, self.backTo)
else:
# write the final result to the disk
self.writeRes(self.resVol, self.resOrient)
self.parallelEnd(verbose)
def findParticles(self,
sizeParticle,
maxNumParticle=0,
minScore=-1,
write2disk=0,
margin=None,
offset=[0, 0, 0]):
"""
findParticles: Find particles in target volume according to the result volume.
@param sizeParticle: size or radius of searched particle
@type sizeParticle: [x,y,z] or integer
@param maxNumParticle: maximal number of particles you want to pick
@type maxNumParticle: integer
@param minScore: minimal score as threshold
@type minScore: float
@param write2disk: write the found particles to the disk or not (0: do not write, otherwise the length of each dimension)
@type write2disk: integer
@param margin: set the margin of the score volume
@param margin: [x,y,z] or integer
@return: list of found particles
@rtype: L{pytom.localization.structures.FoundParticle}
"""
from pytom_volume import vol, peak, putSubVolume, read
from pytom.localization.structures import FoundParticle
# prepare the mask
x = self.result.sizeX()
y = self.result.sizeY()
z = self.result.sizeZ()
if sizeParticle.__class__ == list:
xP = sizeParticle[0]
yP = sizeParticle[1]
zP = sizeParticle[2]
elif sizeParticle.__class__ == vol:
xP = sizeParticle.sizeX()
yP = sizeParticle.sizeY()
zP = sizeParticle.sizeZ()
else:
radius = sizeParticle
xP = 2 * sizeParticle
yP = 2 * sizeParticle
zP = 2 * sizeParticle
if margin:
if margin.__class__ == list:
marginX, marginY, marginZ = margin
else:
marginX = marginY = marginZ = margin
else: # no margin given, set automatically
marginX = int(xP / 2)
marginY = int(yP / 2)
marginZ = int(zP / 2)
mask = vol(x, y, z)
mask.setAll(0)
maskIn = vol(x - 2 * marginX, y - 2 * marginY, z - 2 * marginZ)
maskIn.setAll(1)
putSubVolume(maskIn, mask, marginX, marginY, marginZ)
# progress bar
from pytom.tools.ProgressBar import FixedProgBar
prog = FixedProgBar(0, maxNumParticle - 1, '')
# find the particles
resList = []
for i in range(maxNumParticle):
prog.update(i)
try:
posV = peak(self.result, mask)
except:
break # the mask is all zero
[scoreV, orientV] = self.get(posV[0], posV[1], posV[2])
# test if the peak score is bigger than minimal threshold
if scoreV > minScore:
particleFilename = 'particle_' + str(i) + '.em'
if write2disk:
# write the found particle to the disk
l = write2disk
v = read(self.volFilename, posV[0] - l / 2,
posV[1] - l / 2, posV[2] - l / 2, l, l, l, 0, 0,
0, 0, 0, 0)
v.write(particleFilename)
score = self.score()
score.setValue(scoreV)
from pytom.basic.structures import PickPosition, Rotation
pos = PickPosition(posV, originFilename=self.volFilename)
pos + offset
orientation = Rotation(orientV)
p = FoundParticle(pos, orientation, score, particleFilename)
resList.append(p)
if sizeParticle.__class__ == list:
self.maskOut(mask, posV, [xP, yP, zP])
elif sizeParticle.__class__ == vol:
self.maskOut(mask, posV, sizeParticle)
else:
self.maskOut(mask, posV, radius)
else:
break
return resList