Rescaling in BeagleLib

To add rescaling, we will need to modify several BeagleLib function calls inside our Likelihood class.

First step: tell BeagleLib to create scale buffers

In the Likelihood::initBeagleLib function, add BEAGLE_FLAG_SCALING_MANUAL to the requirement flags and tell BeagleLib to create num_internals+1 scaling buffers (i.e. arrays). Each scaling buffer has one element for each pattern in the data, and if scaling is done at each internal node, there needs to be at least the same number of scaling buffers as there are internal nodes. We allocate num_internals+1 rather than num_internals because one scaling buffer is needed to hold the cumulative sum of log scaling factors for each pattern. The items to add or change are shown below in bold, blue text.

template <class T>
inline void Likelihood<T>::initBeagleLib()
    long preferenceFlags = 0;
    if (_prefer_gpu)
        preferenceFlags |= BEAGLE_FLAG_PROCESSOR_GPU;
        preferenceFlags |= BEAGLE_FLAG_PROCESSOR_CPU;

    requirementFlags |= BEAGLE_FLAG_SCALING_MANUAL;
    BeagleInstanceDetails instance_details;
    _instance = beagleCreateInstance(
         _ntaxa,                    // tips
         num_internals,             // partials
         _ntaxa,                    // sequences
         _nstates,                  // states
         _npatterns,                // patterns
         1,                         // models
         num_transition_probs,      // transition matrices
         _model->_num_categ,        // rate categories
         num_internals + 1,         // scale buffers
         NULL,                      // resource restrictions
         0,                         // length of resource list
         preferenceFlags,           // preferred flags
         requirementFlags,          // required flags
         &instance_details);        // pointer for details

Second step: add scale buffer index to each operation

Modify the Likelihood::defineOperations function to tell BeagleLib which scaling buffer to use when computing the partials for a given internal node. We will reserve scaling buffer 0 to hold the cumulative sum of log scalers, and because internal nodes are numbered starting with _ntaxa, we need to subtract _ntaxa and add 1 from the internal node number to get the index of the scaler buffer to use.

template <class T>
inline void Likelihood<T>::defineOperations(typename Tree<T>::SharedPtr t) 
    // 1. destination partial to be calculated
    int partial = nd->_number;

    // 2. destination scaling buffer index to write to
    int scaler = nd->_number - _ntaxa + 1;

    // 3. destination scaling buffer index to read from

In our application, we need only specify scaling buffers for writing, not reading, so the 4th element of each operation should be left as BEAGLE_OP_NONE.

Third step: specifying the cumulative scale buffer index

We need to tell BeagleLib which scaling buffer to use for the cumulative sum of log scaling factors. We do that in the function Likelihood::calculatePartials by replacing BEAGLE_OP_NONE with 0 as the last argument to beagleUpdatePartials function. We also need to initialize the cumulative scaling factors to make the cumulative sum initially 0.0 for each pattern, which is done by adding a call to beagleResetScaleFactors, passing in the BeagleLib instance handle (stored in the _instance data member) as well as the index of the cumulative scaling factor buffer (0).

template <class T>
inline void Likelihood<T>::calculatePartials()
    int code = beagleResetScaleFactors(_instance, 0);
    if (code != 0)
        throw XStrom(boost::str(boost::format("failed to reset scale factors in calculatePartials. BeagleLib error code was %d (%s)") % code % _beagle_error[code]));

    // Calculate or queue for calculation partials using a list of operations
    int totalOperations = (int)(_operations.size()/7);
    int code = beagleUpdatePartials(
        _instance,                              // Instance number
        (BeagleOperation *) &_operations[0],    // BeagleOperation list specifying operations
        totalOperations,                        // Number of operations
        0);                                     // Index number of scaleBuffer to store accumulated factors

    if (code != 0)
        throw XStrom(boost::str(boost::format("failed to update partials. BeagleLib error code was %d (%s)") % code % _beagle_error[code]));

Fourth step: telling beagleCalculateEdgeLogLikelihoods the cumulative scale buffer index

Finally, we need to provide the index of the scalar array holding the cumulative sum of log scaling factors to the beagleCalculateEdgeLogLikelihoods function, which is called in our Likelihood::calcLogLikelihood function. The beagleCalculateEdgeLogLikelihoods function expects the memory address, not a value. In our case, there is only one cumulative array, and hence only one index to provide, but we must nevertheless create a variable and pass its address rather than the actual value (0). The variable has already been created for you, so you need only replace BEAGLE_OP_NONE with 0.

template <class T>
inline double Likelihood<T>::calcLogLikelihood(typename Tree<T>::SharedPtr t)
    if (t->_is_rooted)
        throw XStrom("can only compute likelihoods for unrooted trees currently");

    if (!_data)
        throw XStrom("must call setData before calcLogLikelihood");

    initBeagleLib(); // this is a no-op if a valid instance already exists
    // Assuming "root" is leaf 0
    assert(t->getRoot()->_number == 0 && t->getRoot()->_left_child == t->_preorder[0] && !t->_preorder[0]->_right_sib);

    // Assuming there are as many transition matrices as there are edge lengths
    assert(_pmatrix_index.size() == _edge_lengths.size());


    // The beagleCalculateEdgeLogLikelihoods function integrates a list of partials 
    // at a parent and child node with respect to a set of partials-weights and 
    // state frequencies to return the log likelihood and first and second derivative sums
    int stateFrequencyIndex  = 0;
    int categoryWeightsIndex = 0;   
    int cumulativeScalingIndex = 0;
    double log_likelihood = 0.0;

    // index_focal_child is the root node
    int index_focal_child  = t->getRoot()->_number;

    // index_focal_parent is the only child of root node  
    int index_focal_parent = t->_preorder[0]->_number;  

    int code = beagleCalculateEdgeLogLikelihoods(
        _instance,                  // instance number
        &index_focal_parent,        // indices of parent partialsBuffers
        &index_focal_child,         // indices of child partialsBuffers
        &index_focal_child,         // transition probability matrices for this edge
        NULL,                       // first derivative matrices
        NULL,                       // second derivative matrices
        &categoryWeightsIndex,      // weights to apply to each partialsBuffer
        &stateFrequencyIndex,       // state frequencies for each partialsBuffer
        &cumulativeScalingIndex,    // scaleBuffers containing accumulated factors
        1,                          // Number of partialsBuffer
        &log_likelihood,            // destination for log likelihood
        NULL,                       // destination for first derivative
        NULL);                      // destination for second derivative

    if (code != 0)
        throw XStrom(boost::str(boost::format("failed to calculate edge logLikelihoods in CalcLogLikelihood. BeagleLib error code was %d (%s)") % code % _beagle_error[code]));

    return log_likelihood;

Run again

If you now run your program, you should find that it computes the log-likelihood correctly.


Adding rescaling has added some computational expense to our likelihood calculation. You can lessen the burden by not rescaling at every internal node. All that is needed to avoid rescaling for a particular internal node is to modify the operation association with that internal node, specifying BEAGLE_OP_NONE instead of the node number for the 3rd element of the operation. Skimping too much, however, runs the risk of overflow. For now, we will rescale at every internal node just to be safe.