private void convolveAndRelease(
Beagle beagle,
int[] firstConvolutionBuffers,
int[] secondConvolutionBuffers,
int[] resultConvolutionBuffers,
int operationsCount) {
if (RUN_IN_SERIES) {
if (operationsCount > 1) {
throw new RuntimeException("Unable to convolve matrices in series");
}
}
beagle.convolveTransitionMatrices(
firstConvolutionBuffers, // A
secondConvolutionBuffers, // B
resultConvolutionBuffers, // C
operationsCount // count
);
for (int i = 0; i < operationsCount; i++) {
if (firstConvolutionBuffers[i] >= matrixBufferHelper.getBufferCount()
&& firstConvolutionBuffers[i] != reserveBufferIndex) {
pushAvailableBuffer(firstConvolutionBuffers[i]);
}
if (secondConvolutionBuffers[i] >= matrixBufferHelper.getBufferCount()
&& secondConvolutionBuffers[i] != reserveBufferIndex) {
pushAvailableBuffer(secondConvolutionBuffers[i]);
}
}
} // END: convolveAndRelease
public void updateSubstitutionModels(Beagle beagle) {
for (int i = 0; i < eigenCount; i++) {
eigenBufferHelper.flipOffset(i);
EigenDecomposition ed = substitutionModelList.get(i).getEigenDecomposition();
beagle.setEigenDecomposition(
eigenBufferHelper.getOffsetIndex(i),
ed.getEigenVectors(),
ed.getInverseEigenVectors(),
ed.getEigenValues());
}
}
public SubstitutionModelDelegate(Tree tree, BranchModel branchModel, int bufferPoolSize) {
if (MEASURE_RUN_TIME) {
updateTime = 0;
convolveTime = 0;
}
this.tree = tree;
this.substitutionModelList = branchModel.getSubstitutionModels();
this.branchModel = branchModel;
eigenCount = substitutionModelList.size();
nodeCount = tree.getNodeCount();
// two eigen buffers for each decomposition for store and restore.
eigenBufferHelper = new BufferIndexHelper(eigenCount, 0);
// two matrices for each node less the root
matrixBufferHelper = new BufferIndexHelper(nodeCount, 0);
this.extraBufferCount =
branchModel.requiresMatrixConvolution()
? (bufferPoolSize > 0 ? bufferPoolSize : BUFFER_POOL_SIZE_DEFAULT)
: 0;
if (branchModel.requiresMatrixConvolution() && this.extraBufferCount < eigenCount) {
throw new RuntimeException(
"SubstitutionModelDelegate requires at least "
+ eigenCount
+ " extra buffers to convolve matrices");
}
for (int i = 0; i < extraBufferCount; i++) {
pushAvailableBuffer(i + matrixBufferHelper.getBufferCount());
}
// one extra created as a reserve
// which is used to free up buffers when the avail stack is empty.
reserveBufferIndex = matrixBufferHelper.getBufferCount() + extraBufferCount;
if (DEBUG) {
System.out.println("Creating reserve buffer with index: " + reserveBufferIndex);
}
} // END: Constructor
private void computeTransitionMatrices(
Beagle beagle, int[][] probabilityIndices, double[][] edgeLengths, int[] counts) {
Timer timer;
if (MEASURE_RUN_TIME) {
timer = new Timer();
timer.start();
}
if (DEBUG) {
System.out.print("Computing matrices:");
}
for (int i = 0; i < eigenCount; i++) {
if (DEBUG) {
for (int j = 0; j < counts[i]; j++) {
// System.out.print(" " + probabilityIndices[i][j]);
System.out.print(" " + probabilityIndices[i][j] + " (" + edgeLengths[i][j] + ")");
}
}
if (counts[i] > 0) {
beagle.updateTransitionMatrices(
eigenBufferHelper.getOffsetIndex(i),
probabilityIndices[i],
null, // firstDerivativeIndices
null, // secondDerivativeIndices
edgeLengths[i],
counts[i]);
}
}
if (DEBUG) {
System.out.println();
}
if (MEASURE_RUN_TIME) {
timer.stop();
double timeInSeconds = timer.toSeconds();
updateTime += timeInSeconds;
}
} // END: computeTransitionMatrices
public void restoreState() {
eigenBufferHelper.restoreState();
matrixBufferHelper.restoreState();
}
public int getMatrixIndex(int branchIndex) {
return matrixBufferHelper.getOffsetIndex(branchIndex);
}
public void flipMatrixBuffer(int branchIndex) {
matrixBufferHelper.flipOffset(branchIndex);
}
public void updateTransitionMatrices(
Beagle beagle, int[] branchIndices, double[] edgeLength, int updateCount) {
int[][] probabilityIndices = new int[eigenCount][updateCount];
double[][] edgeLengths = new double[eigenCount][updateCount];
int[] counts = new int[eigenCount];
List<Deque<Integer>> convolutionList = new ArrayList<Deque<Integer>>();
for (int i = 0; i < updateCount; i++) {
BranchModel.Mapping mapping =
branchModel.getBranchModelMapping(tree.getNode(branchIndices[i]));
int[] order = mapping.getOrder();
double[] weights = mapping.getWeights();
if (order.length == 1) {
int k = order[0];
probabilityIndices[k][counts[k]] = matrixBufferHelper.getOffsetIndex(branchIndices[i]);
edgeLengths[k][counts[k]] = edgeLength[i];
counts[k]++;
} else {
double sum = 0.0;
for (double w : weights) {
sum += w;
}
if (getAvailableBufferCount() < order.length) {
// too few buffers available, process what we have and continue...
computeTransitionMatrices(beagle, probabilityIndices, edgeLengths, counts);
convolveMatrices(beagle, convolutionList);
// reset the counts
for (int k = 0; k < eigenCount; k++) {
counts[k] = 0;
}
}
Deque<Integer> bufferIndices = new ArrayDeque<Integer>();
for (int j = 0; j < order.length; j++) {
int buffer = popAvailableBuffer();
if (buffer < 0) {
// no buffers available
throw new RuntimeException(
"Ran out of buffers for transition matrices - computing current list.");
}
int k = order[j];
probabilityIndices[k][counts[k]] = buffer;
edgeLengths[k][counts[k]] = weights[j] * edgeLength[i] / sum;
// edgeLengths[k][counts[k]] = weights[j] ;
counts[k]++;
bufferIndices.add(buffer);
}
bufferIndices.add(matrixBufferHelper.getOffsetIndex(branchIndices[i]));
convolutionList.add(bufferIndices);
} // END: if convolution needed
} // END: i loop
computeTransitionMatrices(beagle, probabilityIndices, edgeLengths, counts);
convolveMatrices(beagle, convolutionList);
} // END: updateTransitionMatrices
public int getMatrixBufferCount() {
// plus one for the reserve buffer
return matrixBufferHelper.getBufferCount() + extraBufferCount + 1;
}
public int getEigenBufferCount() {
return eigenBufferHelper.getBufferCount();
}