/***************************************************************************
 * PHAST: PHylogenetic Analysis with Space/Time models
 * Copyright (c) 2002-2005 University of California, 2006-2010 Cornell
 * University.  All rights reserved.
 *
 * This source code is distributed under a BSD-style license.  See the
 * file LICENSE.txt for details.
 ***************************************************************************/

/* $Id: fit_column.c,v 1.24 2009-02-19 17:23:17 acs Exp $ */

/* Functions to compute likelihoods for individual alignment columns,
   estimate column-by-column scale factors by maximum likelihood,
   perform single-base LRTs, score tests, phyloP, etc. */

#include <stdlib.h>
#include <phast/fit_column.h>
#include <phast/sufficient_stats.h>
#include <phast/tree_likelihoods.h>
#include <time.h>

#define DERIV_EPSILON 1e-6
/* for numerical computation of derivatives */

#define NSAMPLES_FIM 50
/* number of samples to use in estimating FIM */

#define SIGFIGS 4
/* number of significant figures to which to estimate column scale
   parameters (currently affects 1d parameter estimation only) */

/* Compute and return the log likelihood of a tree model with respect
   to a single column tuple in an alignment.  This is a pared-down
   version of tl_compute_log_likelihood for use in estimation of
   base-by-base scale factors.  It assumes a 0th order model,
   leaf-to-sequence mapping already available, prob matrices computed,
   sufficient stats already available.  Note that this function uses
   natural log rather than log2.  This function does allow for rate
   variation. */
double col_compute_likelihood(TreeModel *mod, MSA *msa, int tupleidx,
                                  double **scratch) {

  int i, j, k, nodeidx, rcat;
  int nstates = mod->rate_matrix->size;
  TreeNode *n;
  double total_prob = 0;
  List *traversal = tr_postorder(mod->tree);
  double **pL = NULL;

  if (msa->ss->tuple_size != 1)
    die("ERROR col_compute_likelihood: need tuple size 1, got %i\n",
	msa->ss->tuple_size);
  if (mod->order != 0)
    die("ERROR col_compute_likelihood: got mod->order of %i, expected 0\n",
	mod->order);
  if (!mod->allow_gaps)
    die("ERROR col_compute_likelihood: need mod->allow_gaps to be TRUE\n");

  /* allocate memory or use scratch if avail */
  if (scratch != NULL)
    pL = scratch;
  else {
    pL = smalloc(nstates * sizeof(double*));
    for (j = 0; j < nstates; j++)
      pL[j] = smalloc((mod->tree->nnodes+1) * sizeof(double));
  }

  for (rcat = 0; rcat < mod->nratecats; rcat++) {
    for (nodeidx = 0; nodeidx < lst_size(traversal); nodeidx++) {
      n = lst_get_ptr(traversal, nodeidx);
      if (n->lchild == NULL) {
        /* leaf: base case of recursion */
        int state = mod->rate_matrix->
          inv_states[(int)ss_get_char_tuple(msa, tupleidx,
                                            mod->msa_seq_idx[n->id], 0)];
        for (i = 0; i < nstates; i++) {
          if (state < 0 || i == state)
            pL[i][n->id] = 1;
          else
            pL[i][n->id] = 0;
        }
      }
      else {
        /* general recursive case */
        MarkovMatrix *lsubst_mat = mod->P[n->lchild->id][rcat];
        MarkovMatrix *rsubst_mat = mod->P[n->rchild->id][rcat];
        for (i = 0; i < nstates; i++) {
          double totl = 0, totr = 0;
          for (j = 0; j < nstates; j++)
            totl += pL[j][n->lchild->id] *
              mm_get(lsubst_mat, i, j);

          for (k = 0; k < nstates; k++)
            totr += pL[k][n->rchild->id] *
              mm_get(rsubst_mat, i, k);

          pL[i][n->id] = totl * totr;
        }
      }
    }

    /* termination (for each rate cat) */
    for (i = 0; i < nstates; i++)
      total_prob += vec_get(mod->backgd_freqs, i) *
        pL[i][mod->tree->id] * mod->freqK[rcat];
  }

  if (scratch == NULL) {
    for (j = 0; j < nstates; j++) sfree(pL[j]);
    sfree(pL);
  }

  return(total_prob);
}



/* See col_compute_likelihood above for notes.
   Note that this function uses natural log rather than log2
 */
double col_compute_log_likelihood(TreeModel *mod, MSA *msa, int tupleidx,
                                  double **scratch) {
  return log(col_compute_likelihood(mod, msa, tupleidx, scratch));
}


/* version of col_scale_derivs_subst that allows for the general case
   of complex eigenvalues and eigenvectors */
void col_scale_derivs_subst_complex(ColFitData *d) {
  Zmatrix *S = d->mod->rate_matrix->evec_matrix_z,
    *Sinv = d->mod->rate_matrix->evec_matrix_inv_z;
  MarkovMatrix *Q = d->mod->rate_matrix;
  int size = Q->size;
  int rcat, nid, i;

  if (S==NULL)
    die("ERROR col_scale_derivs_subst_complex: got S==NULL\n");
  if (Sinv==NULL)
    die("ERROR col_scale_derivs_subst_complex: got Sinv==NULL\n");
  if (d->mod->alt_subst_mods != NULL)
    die("ERROR col_scale_derivs_subst_complex cannot handle lineage-specific models");

  for (rcat = 0; rcat < d->mod->nratecats; rcat++) {
    for (nid = 1; nid < d->mod->tree->nnodes; nid++) { /* skip root */

      double t = ((TreeNode*)lst_get_ptr(d->mod->tree->nodes, nid))->dparent;
      double l1 = d->mod->scale;
      double l2 = (d->stype == SUBTREE && d->mod->in_subtree[nid] ?
                   d->mod->scale_sub : 1);

      /* set up exponentiated diagonal matrix */
      for (i = 0; i < size; i++) {
        double r = t * l1 * l2 * d->mod->rK[rcat];
        zvec_set(d->expdiag_z, i, z_exp(z_mul_real(zvec_get(Q->evals_z, i), r)));
      }

      /* PP */
      zvec_copy(d->vec_scratch1_z, Q->evals_z);
      zvec_scale(d->vec_scratch1_z, t * l2);
      zvec_had_prod(d->vec_scratch2_z, d->vec_scratch1_z, d->expdiag_z);
      zmat_mult_real_diag(d->PP[nid][rcat], S, d->vec_scratch2_z, Sinv,
                          d->mat_scratch_z);

      if (d->second_derivs) {
        /* PPP */
        zvec_had_prod(d->vec_scratch2_z, d->vec_scratch1_z, d->vec_scratch1_z);
        zvec_had_prod(d->vec_scratch1_z, d->vec_scratch2_z, d->expdiag_z);
        zmat_mult_real_diag(d->PPP[nid][rcat], S, d->vec_scratch1_z, Sinv,
                            d->mat_scratch_z);
      }

      if (d->stype == SUBTREE && d->mod->in_subtree[nid]) {
        /* if not in subtree, leave all of these equal to 0 (as
           initialized) */

        /* QQ */
        zvec_copy(d->vec_scratch1_z, Q->evals_z);
        zvec_scale(d->vec_scratch1_z, t * l1);
        zvec_had_prod(d->vec_scratch2_z, d->vec_scratch1_z, d->expdiag_z);
        zmat_mult_real_diag(d->QQ[nid][rcat], S, d->vec_scratch2_z, Sinv,
                            d->mat_scratch_z);

        if (d->second_derivs) {
          /* QQQ */
          zvec_had_prod(d->vec_scratch2_z, d->vec_scratch1_z, d->vec_scratch1_z);
          zvec_had_prod(d->vec_scratch1_z, d->vec_scratch2_z, d->expdiag_z);
          zmat_mult_real_diag(d->QQQ[nid][rcat], S, d->vec_scratch1_z, Sinv,
                              d->mat_scratch_z);

          /* RRR */
          zvec_copy(d->vec_scratch1_z, Q->evals_z);
          zvec_scale(d->vec_scratch1_z, t);
          zvec_had_prod(d->vec_scratch2_z, Q->evals_z, Q->evals_z);
          zvec_scale(d->vec_scratch2_z, t * t * l1 * l2);
          zvec_plus_eq(d->vec_scratch2_z, d->vec_scratch1_z);
          zvec_had_prod(d->vec_scratch1_z, d->vec_scratch2_z, d->expdiag_z);
          zmat_mult_real_diag(d->RRR[nid][rcat], S, d->vec_scratch1_z, Sinv,
                              d->mat_scratch_z);
        }
      }
    }
  }
}

/* version of col_scale_derivs_subst that is optimized for the case in
   which eigenvalues and eigenvectors can be assumed to be real */
void col_scale_derivs_subst_real(ColFitData *d) {
  Matrix *S = d->mod->rate_matrix->evec_matrix_r,
    *Sinv = d->mod->rate_matrix->evec_matrix_inv_r;
  MarkovMatrix *Q = d->mod->rate_matrix;
  int size = Q->size;
  int rcat, nid, i;

  if (S==NULL)
    die("ERROR col_scale_derivs_subst_real: got S==NULL\n");
  if (Sinv==NULL)
    die("ERROR col_scale_derivs_subst_real: got Sinv==NULL\n");
  if (d->mod->alt_subst_mods != NULL)
    die("ERROR col_scale_derivs_subst_real: cannot handle lineage-specific models");

  for (rcat = 0; rcat < d->mod->nratecats; rcat++) {
    for (nid = 1; nid < d->mod->tree->nnodes; nid++) { /* skip root */

      double t = ((TreeNode*)lst_get_ptr(d->mod->tree->nodes, nid))->dparent;
      double l1 = d->mod->scale;
      double l2 = (d->stype == SUBTREE && d->mod->in_subtree[nid] ?
                   d->mod->scale_sub : 1);

      /* set up exponentiated diagonal matrix */
      for (i = 0; i < size; i++)
        d->expdiag_r->data[i] = exp(Q->evals_r->data[i] * t *
                                    l1 * l2 * d->mod->rK[rcat]);

      /* PP */
      for (i = 0; i < size; i++) {
        d->vec_scratch1_r->data[i] = Q->evals_r->data[i] * t * l2;
        d->vec_scratch2_r->data[i] = d->vec_scratch1_r->data[i] *
          d->expdiag_r->data[i];
      }
      mat_mult_diag(d->PP[nid][rcat], S, d->vec_scratch2_r, Sinv);

      if (d->second_derivs) {
        /* PPP */
        for (i = 0; i < size; i++) {
          d->vec_scratch2_r->data[i] = d->vec_scratch1_r->data[i] *
            d->vec_scratch1_r->data[i];
          d->vec_scratch1_r->data[i] = d->vec_scratch2_r->data[i] *
            d->expdiag_r->data[i];
        }
        mat_mult_diag(d->PPP[nid][rcat], S, d->vec_scratch1_r, Sinv);
      }

      if (d->stype == SUBTREE && d->mod->in_subtree[nid]) {
        /* if not in subtree, leave all of these equal to 0 (as
           initialized) */

        /* QQ */
        for (i = 0; i < size; i++) {
          d->vec_scratch1_r->data[i] = Q->evals_r->data[i] * t * l1;
          d->vec_scratch2_r->data[i] = d->vec_scratch1_r->data[i] *
            d->expdiag_r->data[i];
        }
        mat_mult_diag(d->QQ[nid][rcat], S, d->vec_scratch2_r, Sinv);

        if (d->second_derivs) {
          /* QQQ */
          for (i = 0; i < size; i++) {
            d->vec_scratch2_r->data[i] = d->vec_scratch1_r->data[i] *
              d->vec_scratch1_r->data[i];
            d->vec_scratch1_r->data[i] = d->vec_scratch2_r->data[i] *
              d->expdiag_r->data[i];
          }
          mat_mult_diag(d->QQQ[nid][rcat], S, d->vec_scratch1_r, Sinv);

          /* RRR */
          for (i = 0; i < size; i++)
            d->vec_scratch1_r->data[i] = d->expdiag_r->data[i] *
              (Q->evals_r->data[i] * Q->evals_r->data[i] * t * t * l1 * l2 +
               Q->evals_r->data[i] * t);
          mat_mult_diag(d->RRR[nid][rcat], S, d->vec_scratch1_r, Sinv);
        }
      }
    }
  }
}

/* Compute 1st and 2nd derivs wrt scale params of substitution
   matrices for each branch of the tree (and each rate category).
   These are used in the recursive computation of derivatives of the
   column likelihoods */
void col_scale_derivs_subst(ColFitData *d) {
  /* now broken into real and complex cases for efficiency */
  if (d->mod->rate_matrix->eigentype == REAL_NUM)
    col_scale_derivs_subst_real(d);
  else
    col_scale_derivs_subst_complex(d);
}

/* Compute the first and (optionally) second derivatives with respect
   to the scale parameter for the single-column log likelihood
   function (col_compute_log_likelihood).  This version assumes a
   single scale parameter; see below for the subtree case.  Return
   value is log likelihood, which is computed as a by-product.  Derivs
   will be stored in *first_deriv and *second_deriv.  If second_deriv
   == NULL, it will not be computed (saves some time).  */
double col_scale_derivs(ColFitData *d, double *first_deriv,
                        double *second_deriv, double ***scratch) {

  int i, j, k, nodeidx, rcat;
  int nstates = d->mod->rate_matrix->size;
  TreeNode *n;
  double total_prob = 0;
  List *traversal = tr_postorder(d->mod->tree);
  double **L=NULL;                   /* partial likelihoods */
  double **LL=NULL;                  /* 1st deriv of partial likelihoods wrt
                                   scale param */
  double **LLL=NULL;                 /* 2nd deriv of partial likelihoods
                                   wrt scale param */
  if (d->msa->ss->tuple_size != 1)
    die("ERROR col_scale_derivs: need tuple size 1, got %i\n",
	d->msa->ss->tuple_size);
  if (d->mod->order != 0)
    die("ERROR col_scale_derivs: got mod->order of %i, expected 0\n",
	d->mod->order);
  if (!d->mod->allow_gaps)
    die("ERROR col_scale_derivs: need mod->allow_gaps to be TRUE\n");

  *first_deriv = 0;
  if (second_deriv != NULL) *second_deriv = 0;

  /* allocate memory or use scratch if available */
  if (scratch == NULL) {
    L = smalloc(nstates * sizeof(double*));
    LL = smalloc(nstates * sizeof(double*));
    if (second_deriv != NULL)
      LLL = smalloc(nstates * sizeof(double*));
    for (j = 0; j < nstates; j++) {
      L[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
      LL[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
      if (second_deriv != NULL)
        LLL[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
    }
  }
  else {
    L = scratch[0];
    LL = scratch[1];
    LLL = scratch[2];
  }

  col_scale_derivs_subst(d);

  for (rcat = 0; rcat < d->mod->nratecats; rcat++) {
    for (nodeidx = 0; nodeidx < lst_size(traversal); nodeidx++) {
      n = lst_get_ptr(traversal, nodeidx);
      if (n->lchild == NULL) {
        /* leaf: base case of recursion */
        int state = d->mod->rate_matrix->
          inv_states[(int)ss_get_char_tuple(d->msa, d->tupleidx,
                                            d->mod->msa_seq_idx[n->id], 0)];
        for (i = 0; i < nstates; i++) {
          if (state < 0 || i == state)
            L[i][n->id] = 1;
          else
            L[i][n->id] = 0;

          LL[i][n->id] = 0;
          if (second_deriv != NULL) LLL[i][n->id] = 0;
        }
      }
      else {
        /* general recursive case */
        MarkovMatrix *lsubst_mat = d->mod->P[n->lchild->id][rcat];
        MarkovMatrix *rsubst_mat = d->mod->P[n->rchild->id][rcat];
        for (i = 0; i < nstates; i++) {
          double totl = 0, totr = 0, A = 0, B = 0, E = 0, F = 0;
          for (j = 0; j < nstates; j++) {
            totl += L[j][n->lchild->id] * mm_get(lsubst_mat, i, j);

            A += (L[j][n->lchild->id] * d->PP[n->lchild->id][rcat]->data[i][j]) +
              (LL[j][n->lchild->id] * mm_get(lsubst_mat, i, j));
          }

          for (k = 0; k < nstates; k++) {
            totr += L[k][n->rchild->id] * mm_get(rsubst_mat, i, k);

            B += (L[k][n->rchild->id] * d->PP[n->rchild->id][rcat]->data[i][k]) +
              (LL[k][n->rchild->id] * mm_get(rsubst_mat, i, k));

          }

          L[i][n->id] = totl * totr;
          LL[i][n->id] = totr*A + totl*B;

          if (second_deriv != NULL) {
            for (j = 0; j < nstates; j++)
              E += L[j][n->lchild->id] * d->PPP[n->lchild->id][rcat]->data[i][j] +
                2 * LL[j][n->lchild->id] * d->PP[n->lchild->id][rcat]->data[i][j] +
                LLL[j][n->lchild->id] * mm_get(lsubst_mat, i, j);

            for (k = 0; k < nstates; k++)
              F += L[k][n->rchild->id] * d->PPP[n->rchild->id][rcat]->data[i][k] +
                2 * LL[k][n->rchild->id] * d->PP[n->rchild->id][rcat]->data[i][k] +
                LLL[k][n->rchild->id] * mm_get(rsubst_mat, i, k);

            LLL[i][n->id] = totr*E + 2*A*B + totl*F;
          }
        }
      }
    }

    /* termination (for each rate cat) */
    for (i = 0; i < nstates; i++) {
      total_prob += L[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
        d->mod->freqK[rcat];

      *first_deriv += LL[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
        d->mod->freqK[rcat];

      if (second_deriv != NULL)
        *second_deriv += LLL[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
          d->mod->freqK[rcat];
    }
  }

  /* convert to log space */
  if (second_deriv != NULL)
    *second_deriv = (*second_deriv)/total_prob -
      ((*first_deriv)/total_prob) * ((*first_deriv)/total_prob);
                                /* deriv of log followed by quotient
                                   rule; rearrange terms to avoid
                                   underflow */

  *first_deriv = *first_deriv / total_prob; /* deriv of log */
  total_prob = log(total_prob);

  if (scratch == NULL) {
    for (j = 0; j < nstates; j++) {
      sfree(L[j]);
      sfree(LL[j]);
      if (second_deriv != NULL) sfree(LLL[j]);
    }
    sfree(L);
    sfree(LL);
    if (second_deriv != NULL) sfree(LLL);
  }

  return(total_prob);
}

/* Compute the first and (optionally) second derivatives with respect
   to the scale parameters for the single-column log likelihood
   function (col_compute_log_likelihood).  This version assumes scale
   parameters for the whole tree and for the subtree.  Return value is
   log likelihood, which is computed as a by-product.  Derivs will be
   stored in *gradient and *hessian.  If hessian == NULL,
   it will not be computed (saves some time).  */
double col_scale_derivs_subtree(ColFitData *d, Vector *gradient,
                                Matrix *hessian, double ***scratch) {
  int i, j, k, nodeidx, rcat;
  int nstates = d->mod->rate_matrix->size;
  TreeNode *n;
  double total_prob = 0;
  List *traversal = tr_postorder(d->mod->tree);
  double **L=NULL;                   /* partial likelihoods */
  double **LL=NULL;                  /* 1st deriv of partial likelihoods wrt
                                   1st scale param */
  double **LLL=NULL;                 /* 2nd deriv of partial likelihoods
                                   wrt 1st scale param */
  double **MM=NULL;                  /* 1st deriv of partial likelihoods
                                   wrt 2nd scale param */
  double **MMM=NULL;                 /* 2nd deriv of partial likelihoods
                                   wrt 2nd scale param */
  double **NNN=NULL;                 /* 2nd cross deriv (off diagonal in
                                   Hessian) of partial likelihoods */

  double *pd = gradient->data;  /* 1st partial derivatives */
  double **pd2 = (hessian == NULL ? NULL : hessian->data);
                                /* 2nd partial derivatives; because of
                                   symmetry, only pd2[0][0],
                                   pd2[1][1], and pd2[1][0] need to be
                                   considered during computation */
  if (d->msa->ss->tuple_size != 1)
    die("ERROR col_scale_derivs_subtree: need tuple size 1, got %i\n",
	d->msa->ss->tuple_size);
  if (d->mod->order != 0)
    die("ERROR col_scale_derivs_subtree: got mod->order of %i, expected 0\n",
	d->mod->order);
  if (!d->mod->allow_gaps)
    die("ERROR col_scale_derivs_subtree: need mod->allow_gaps to be TRUE\n");

  pd[0] = pd[1] = 0;
  if (pd2 != NULL)
    pd2[0][0] = pd2[1][1] = pd2[0][1] = pd2[1][0] = 0;

  /* allocate memory or use scratch if available */
  if (scratch == NULL) {
    L = smalloc(nstates * sizeof(double*));
    LL = smalloc(nstates * sizeof(double*));
    MM = smalloc(nstates * sizeof(double*));
    if (pd2 != NULL) {
      LLL = smalloc(nstates * sizeof(double*));
      MMM = smalloc(nstates * sizeof(double*));
      NNN = smalloc(nstates * sizeof(double*));
    }
    for (j = 0; j < nstates; j++) {
      L[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
      LL[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
      MM[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
      if (pd2 != NULL) {
        LLL[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
        MMM[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
        NNN[j] = smalloc((d->mod->tree->nnodes+1) * sizeof(double));
      }
    }
  }
  else {
    L = scratch[0];
    LL = scratch[1];
    LLL = scratch[2];
    MM = scratch[3];
    MMM = scratch[4];
    NNN = scratch[5];
  }

  col_scale_derivs_subst(d);

  for (rcat = 0; rcat < d->mod->nratecats; rcat++) {
    for (nodeidx = 0; nodeidx < lst_size(traversal); nodeidx++) {
      n = lst_get_ptr(traversal, nodeidx);
      if (n->lchild == NULL) {
        /* leaf: base case of recursion */
        int state = d->mod->rate_matrix->
          inv_states[(int)ss_get_char_tuple(d->msa, d->tupleidx,
                                            d->mod->msa_seq_idx[n->id], 0)];
        for (i = 0; i < nstates; i++) {
          if (state < 0 || i == state)
            L[i][n->id] = 1;
          else
            L[i][n->id] = 0;

          LL[i][n->id] = MM[i][n->id] = 0;
          if (pd2 != NULL)
            LLL[i][n->id] = MMM[i][n->id] = NNN[i][n->id] = 0;
        }
      }
      else {
        /* general recursive case */
        MarkovMatrix *lsubst_mat = d->mod->P[n->lchild->id][rcat];
        MarkovMatrix *rsubst_mat = d->mod->P[n->rchild->id][rcat];
        for (i = 0; i < nstates; i++) {
          double totl = 0, totr = 0, A = 0, B = 0, C = 0, D = 0, E = 0,
            F = 0, G = 0, H = 0, I = 0, J = 0;
          for (j = 0; j < nstates; j++) {
            totl += L[j][n->lchild->id] * mm_get(lsubst_mat, i, j);

            A += (L[j][n->lchild->id] * d->PP[n->lchild->id][rcat]->data[i][j]) +
              (LL[j][n->lchild->id] * mm_get(lsubst_mat, i, j));

            C += (L[j][n->lchild->id] * d->QQ[n->lchild->id][rcat]->data[i][j]) +
              (MM[j][n->lchild->id] * mm_get(lsubst_mat, i, j));
          }

          for (k = 0; k < nstates; k++) {
            totr += L[k][n->rchild->id] * mm_get(rsubst_mat, i, k);

            B += (L[k][n->rchild->id] * d->PP[n->rchild->id][rcat]->data[i][k]) +
              (LL[k][n->rchild->id] * mm_get(rsubst_mat, i, k));

            D += (L[k][n->rchild->id] * d->QQ[n->rchild->id][rcat]->data[i][k]) +
              (MM[k][n->rchild->id] * mm_get(rsubst_mat, i, k));

          }

          L[i][n->id] = totl * totr;
          LL[i][n->id] = totr*A + totl*B;
          MM[i][n->id] = totr*C + totl*D;

          if (pd2 != NULL) {
            for (j = 0; j < nstates; j++) {
              E += L[j][n->lchild->id] * d->PPP[n->lchild->id][rcat]->data[i][j] +
                2 * LL[j][n->lchild->id] * d->PP[n->lchild->id][rcat]->data[i][j] +
                LLL[j][n->lchild->id] * mm_get(lsubst_mat, i, j);
              G += L[j][n->lchild->id] * d->QQQ[n->lchild->id][rcat]->data[i][j] +
                2 * MM[j][n->lchild->id] * d->QQ[n->lchild->id][rcat]->data[i][j] +
                MMM[j][n->lchild->id] * mm_get(lsubst_mat, i, j);
              I += L[j][n->lchild->id] * d->RRR[n->lchild->id][rcat]->data[i][j] +
                MM[j][n->lchild->id] * d->PP[n->lchild->id][rcat]->data[i][j] +
                LL[j][n->lchild->id] * d->QQ[n->lchild->id][rcat]->data[i][j] +
                NNN[j][n->lchild->id] * mm_get(lsubst_mat, i, j);
            }

            for (k = 0; k < nstates; k++) {
              F += L[k][n->rchild->id] * d->PPP[n->rchild->id][rcat]->data[i][k] +
                2 * LL[k][n->rchild->id] * d->PP[n->rchild->id][rcat]->data[i][k] +
                LLL[k][n->rchild->id] * mm_get(rsubst_mat, i, k);
              H += L[k][n->rchild->id] * d->QQQ[n->rchild->id][rcat]->data[i][k] +
                2 * MM[k][n->rchild->id] * d->QQ[n->rchild->id][rcat]->data[i][k] +
                MMM[k][n->rchild->id] * mm_get(rsubst_mat, i, k);
              J += L[k][n->rchild->id] * d->RRR[n->rchild->id][rcat]->data[i][k] +
                MM[k][n->rchild->id] * d->PP[n->rchild->id][rcat]->data[i][k] +
                LL[k][n->rchild->id] * d->QQ[n->rchild->id][rcat]->data[i][k] +
                NNN[k][n->rchild->id] * mm_get(rsubst_mat, i, k);
            }

            LLL[i][n->id] = totr*E + 2*A*B + totl*F;
            MMM[i][n->id] = totr*G + 2*C*D + totl*H;
            NNN[i][n->id] = totr*I + A*D + B*C + totl*J;
          }
        }
      }
    }

    /* termination (for each rate cat) */
    for (i = 0; i < nstates; i++) {
      total_prob += L[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
        d->mod->freqK[rcat];

      pd[0] += LL[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
        d->mod->freqK[rcat];
      pd[1] += MM[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
        d->mod->freqK[rcat];

      if (pd2 != NULL) {
        pd2[0][0] += LLL[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
          d->mod->freqK[rcat];
        pd2[1][1] += MMM[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
          d->mod->freqK[rcat];
        pd2[1][0] += NNN[i][d->mod->tree->id] * vec_get(d->mod->backgd_freqs, i) *
          d->mod->freqK[rcat];
      }
    }
 }

  /* convert to log space */
  if (pd2 != NULL) {
    /* deriv of log and quotient rule */
    pd2[0][0] =  pd2[0][0]/total_prob - (pd[0]/total_prob)*(pd[0]/total_prob);
    pd2[1][1] =  pd2[1][1]/total_prob - (pd[1]/total_prob)*(pd[1]/total_prob);
    pd2[1][0] =  pd2[1][0]/total_prob - (pd[1]/total_prob)*(pd[0]/total_prob);
    pd2[0][1] =  pd2[1][0];
  }
  pd[0] = pd[0] / total_prob; /* deriv of log */
  pd[1] = pd[1] / total_prob;
  total_prob = log(total_prob);

  if (scratch == NULL) {
    for (j = 0; j < nstates; j++) {
      sfree(L[j]);
      sfree(LL[j]);
      sfree(MM[j]);
      if (pd2 != NULL) {
        sfree(LLL[j]);
        sfree(MMM[j]);
        sfree(NNN[j]);
      }
    }
    sfree(L);
    sfree(LL);
    sfree(MM);
    if (pd2 != NULL) {
      sfree(LLL);
      sfree(MMM);
      sfree(NNN);
    }
  }

  return(total_prob);
}


/* Wrapper for likelihood function for use in parameter estimation */
double col_likelihood_wrapper(Vector *params, void *data) {
  ColFitData *d = (ColFitData*)data;

  d->mod->scale = vec_get(params, 0);
  if (d->stype == SUBTREE)
    d->mod->scale_sub = vec_get(params, 1);

  /* reestimate subst models on edges */
  tm_set_subst_matrices(d->mod);

  return -1 * col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                         d->fels_scratch[0]);
}

/* Wrapper for likelihood function for use in parameter estimation;
   version for use with opt_newton_1d */
double col_likelihood_wrapper_1d(double x, void *data) {
  ColFitData *d = (ColFitData*)data;
  if (d->stype == SUBTREE)
    die("ERROR col_likelihood_wrapper_1d: d->stype cannot be SUBTREE\n");

  d->mod->scale = x;

  /* reestimate subst models on edges */
  tm_set_subst_matrices(d->mod);

  return -1 * col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                         d->fels_scratch[0]);
}

/* Wrapper for gradient function for use in parameter estimation */
void col_grad_wrapper(Vector *grad, Vector *params, void *data,
                      Vector *lb, Vector *ub) {
  ColFitData *d = (ColFitData*)data;
  double deriv;

  if (d->stype == ALL) {
    col_scale_derivs(d, &deriv, NULL, d->fels_scratch);
    vec_set(grad, 0, -deriv);   /* because working with neg lnl */
  }
  else {
    col_scale_derivs_subtree(d, grad, NULL, d->fels_scratch);
    vec_scale(grad, -1);
  }
}

/* Wrapper for gradient function for use in parameter estimation;
   version for use with opt_newton_1d */
double col_grad_wrapper_1d(double x, void *data, double lb, double ub) {
  double deriv, deriv2;
  ColFitData *d = (ColFitData*)data;
  if (d->stype != ALL)
    die("ERROR col_grad_wrapper_1d: d->stype must be ALL\n");
  col_scale_derivs(d, &deriv, &deriv2, d->fels_scratch);
  d->deriv2 = -deriv2;           /* store for use by wrapper below */
  return -deriv; /* because working with neg lnl */
}

/* Wrapper for second derivative function for use in parameter
   estimation; version for use with opt_newton_1d.  Simply returns
   value computed in col_grad_wrapper_1d (1st and 2nd derivs are
   computed simultaneously) */
double col_deriv2_wrapper_1d(double x, void *data, double lb, double ub) {
  ColFitData *d = (ColFitData*)data;
  return d->deriv2;
}

/* Perform a likelihood ratio test for each column tuple in an
   alignment, comparing the given null model with an alternative model
   that has a free scaling parameter for all branches.  Assumes a 0th
   order model, leaf-to-sequence mapping already available, prob
   matrices computed, sufficient stats available.  Computes p-values
   based using the chi-sq distribution and stores them in tuple_pvals.
   Will optionally store the individual scale factors in tuple_scales
   and raw log likelihood ratios in tuple_llrs if these variables are
   non-NULL.  Must define mode as CON (for 0 <= scale <= 1), ACC
   (for 1 <= scale), NNEUT (0 <= scale), or CONACC (0 <= scale) */
void col_lrts(TreeModel *mod, MSA *msa, mode_type mode, double *tuple_pvals,
              double *tuple_scales, double *tuple_llrs, FILE *logf) {
  int i;
  ColFitData *d;
  double null_lnl, alt_lnl, delta_lnl, this_scale = 1;

  /* init ColFitData */
  d = col_init_fit_data(mod, msa, ALL, mode, FALSE);

  /* iterate through column tuples */
  for (i = 0; i < msa->ss->ntuples; i++) {
    checkInterruptN(i, 100);

    /* first check for actual substitution data in column; if none,
       don't waste time computing likelihoods */
    if (!col_has_data(mod, msa, i)) {
      delta_lnl = 0;
      this_scale = 1;
    }

    else {                      /* compute null and alt lnl */
      mod->scale = 1;
      tm_set_subst_matrices(mod);

      /* compute log likelihoods under null and alt hypotheses */
      null_lnl = col_compute_log_likelihood(mod, msa, i, d->fels_scratch[0]);

      vec_set(d->params, 0, d->init_scale);
      d->tupleidx = i;

      opt_newton_1d(col_likelihood_wrapper_1d, &d->params->data[0], d,
                    &alt_lnl, SIGFIGS, d->lb->data[0], d->ub->data[0],
                    logf, NULL, NULL);
      /* turns out to be faster (roughly 15% in limited experiments)
         to use numerical rather than exact derivatives */

      alt_lnl *= -1;
      this_scale = d->params->data[0];

      delta_lnl = alt_lnl - null_lnl;
      if (delta_lnl <= -0.01)
	die("ERROR col_lrts: delta_lnl = %e < -0.01\n", delta_lnl);
      if (delta_lnl < 0) delta_lnl = 0;
    } /* end estimation of delta_lnl */

    /* compute p-vals via chi-sq */
    if (tuple_pvals != NULL) {
      if (mode == NNEUT || mode == CONACC)
        tuple_pvals[i] = chisq_cdf(2*delta_lnl, 1, FALSE);
      else
        tuple_pvals[i] = half_chisq_cdf(2*delta_lnl, 1, FALSE);
        /* assumes 50:50 mix of chisq and point mass at zero, due to
           bounding of param */

      if (tuple_pvals[i] < 1e-20)
        tuple_pvals[i] = 1e-20;
      /* approx limit of eval of tail prob; pvals of 0 cause problems */

      if (mode == CONACC && this_scale > 1)
          tuple_pvals[i] *= -1; /* mark as acceleration */
    }

    /* store scales and log likelihood ratios if necessary */
    if (tuple_scales != NULL) tuple_scales[i] = this_scale;
    if (tuple_llrs != NULL) tuple_llrs[i] = delta_lnl;
  }

  col_free_fit_data(d);
}

/* Subtree version of LRT */
void col_lrts_sub(TreeModel *mod, MSA *msa, mode_type mode,
                  double *tuple_pvals, double *tuple_null_scales,
                  double *tuple_scales, double *tuple_sub_scales,
                  double *tuple_llrs, FILE *logf) {
  int i;
  ColFitData *d, *d2;
  double null_lnl, alt_lnl, delta_lnl;
  TreeModel *modcpy;
  List *inside=NULL, *outside=NULL;

  modcpy = tm_create_copy(mod);   /* need separate copy of tree model
                                     with different internal scaling
                                     data for supertree/subtree case */
  modcpy->subtree_root = NULL;

  /* init ColFitData -- one for null model, one for alt */
  d = col_init_fit_data(modcpy, msa, ALL, NNEUT, FALSE);
  d2 = col_init_fit_data(mod, msa, SUBTREE, mode, FALSE);
                                /* mod has the subtree info, modcpy
                                   does not */

  /* prepare lists of leaves inside and outside root, for use in
     checking for informative substitutions */
  if (mod->subtree_root != NULL) {
    inside = lst_new_ptr(mod->tree->nnodes);
    outside = lst_new_ptr(mod->tree->nnodes);
    tr_partition_leaves(mod->tree, mod->subtree_root, inside, outside);
  }

  /* iterate through column tuples */
  for (i = 0; i < msa->ss->ntuples; i++) {
    checkInterruptN(i, 100);

    /* first check for informative substitution data in column; if none,
       don't waste time computing likeihoods */
    if (!col_has_data_sub(mod, msa, i, inside, outside)) {
      delta_lnl = 0;
      d->params->data[0] = d2->params->data[0] = d2->params->data[1] = 1;
    }

    else {
      /* compute log likelihoods under null and alt hypotheses */
      d->tupleidx = i;
      vec_set(d->params, 0, d->init_scale);
      opt_newton_1d(col_likelihood_wrapper_1d, &d->params->data[0], d,
                    &null_lnl, SIGFIGS, d->lb->data[0], d->ub->data[0],
                    logf, NULL, NULL);

      //      opt_bfgs(col_likelihood_wrapper, d->params, d, &null_lnl, d->lb,
      //	       d->ub, logf, NULL, OPT_HIGH_PREC, NULL, NULL);

      /* turns out to be faster (roughly 15% in limited experiments)
         to use numerical rather than exact derivatives */
      null_lnl *= -1;

      d2->tupleidx = i;
      vec_set(d2->params, 0, max(0.05, d->params->data[0]));
      /* init to previous estimate to save time, but don't init to
         value at boundary */
      vec_set(d2->params, 1, d2->init_scale_sub);

      if (opt_bfgs(col_likelihood_wrapper, d2->params, d2, &alt_lnl, d2->lb,
                   d2->ub, logf, NULL, OPT_HIGH_PREC, NULL, NULL) != 0)
        ;                         /* do nothing; nonzero exit typically
                                     occurs when max iterations is
                                     reached; a warning is printed to
                                     the log */
      alt_lnl *= -1;

      delta_lnl = alt_lnl - null_lnl;
      if (delta_lnl <= -0.1)
	die("ERROR col_lrts_sub: delta_lnl = %e <= -0.1\n", delta_lnl);
      if (delta_lnl < 0) delta_lnl = 0;
    }

    /* compute p-vals via chi-sq */
    if (tuple_pvals != NULL) {
      if (mode == NNEUT || mode == CONACC)
        tuple_pvals[i] = chisq_cdf(2*delta_lnl, 1, FALSE);
      else
        tuple_pvals[i] = half_chisq_cdf(2*delta_lnl, 1, FALSE);
        /* assumes 50:50 mix of chisq and point mass at zero, due to
           bounding of param */

      if (tuple_pvals[i] < 1e-20)
        tuple_pvals[i] = 1e-20;
      /* approx limit of eval of tail prob; pvals of 0 cause problems */

      if (mode == CONACC && d2->params->data[1] > 1)
        tuple_pvals[i] *= -1;    /* mark as acceleration */
    }

    /* store scales and log likelihood ratios if necessary */
    if (tuple_null_scales != NULL)
      tuple_null_scales[i] = d->params->data[0];
    if (tuple_scales != NULL)
      tuple_scales[i] = d2->params->data[0];
    if (tuple_sub_scales != NULL)
      tuple_sub_scales[i] = d2->params->data[1];
    if (tuple_llrs != NULL)
      tuple_llrs[i] = delta_lnl;
  }

  col_free_fit_data(d);
  col_free_fit_data(d2);
  modcpy->estimate_branchlens = TM_BRANCHLENS_ALL;
                                /* have to revert for tm_free to work
                                   correctly */
  tm_free(modcpy);
  if (inside != NULL) lst_free(inside);
  if (outside != NULL) lst_free(outside);
}

/* Score test */
void col_score_tests(TreeModel *mod, MSA *msa, mode_type mode,
                     double *tuple_pvals, double *tuple_derivs,
                     double *tuple_teststats) {
  int i;
  ColFitData *d;
  double first_deriv, teststat, fim;

  /* init ColFitData */
  d = col_init_fit_data(mod, msa, ALL, NNEUT, FALSE);

  /* precompute FIM */
  fim = col_estimate_fim(mod);

  if (fim < 0)
    die("ERROR: negative fisher information in col_score_tests\n");

  /* iterate through column tuples */
  for (i = 0; i < msa->ss->ntuples; i++) {
    checkInterruptN(i, 1000);

    /* first check for actual substitution data in column; if none,
       don't waste time computing score */
    if (!col_has_data(mod, msa, i)) {
      first_deriv = 0;
      teststat = 0;
    }

    else {
      d->tupleidx = i;

      col_scale_derivs(d, &first_deriv, NULL, d->fels_scratch);

      teststat = first_deriv*first_deriv / fim;

      if ((mode == ACC && first_deriv < 0) ||
          (mode == CON && first_deriv > 0))
        teststat = 0;             /* derivative points toward boundary;
                                     truncate at 0 */
    }

    if (tuple_pvals != NULL) {
      if (mode == NNEUT || mode == CONACC)
        tuple_pvals[i] = chisq_cdf(teststat, 1, FALSE);
      else
        tuple_pvals[i] = half_chisq_cdf(teststat, 1, FALSE);
        /* assumes 50:50 mix of chisq and point mass at zero */

      if (tuple_pvals[i] < 1e-20)
        tuple_pvals[i] = 1e-20;
      /* approx limit of eval of tail prob; pvals of 0 cause problems */

      if (mode == CONACC && first_deriv > 0)
        tuple_pvals[i] *= -1; /* mark as acceleration */
    }

    /* store scales and log likelihood ratios if necessary */
    if (tuple_derivs != NULL) tuple_derivs[i] = first_deriv;
    if (tuple_teststats != NULL) tuple_teststats[i] = teststat;
  }

  col_free_fit_data(d);
}

/* Subtree version of score test */
void col_score_tests_sub(TreeModel *mod, MSA *msa, mode_type mode,
                         double *tuple_pvals, double *tuple_null_scales,
                         double *tuple_derivs, double *tuple_sub_derivs,
                         double *tuple_teststats, FILE *logf) {
  int i;
  ColFitData *d, *d2;
  Vector *grad = vec_new(2);
  Matrix *fim;
  double lnl, teststat;
  FimGrid *grid;
  List *inside=NULL, *outside=NULL;
  TreeModel *modcpy = tm_create_copy(mod); /* need separate copy of tree model
                                              with different internal scaling
                                              data for supertree/subtree case */
  modcpy->subtree_root = NULL;

  /* init ColFitData -- one for null model, one for alt */
  d = col_init_fit_data(modcpy, msa, ALL, NNEUT, FALSE);
  d2 = col_init_fit_data(mod, msa, SUBTREE, NNEUT, FALSE);
                                /* mod has the subtree info, modcpy
                                   does not */

  /* precompute Fisher information matrices for a grid of scale values */
  grid = col_fim_grid_sub(mod);

  /* prepare lists of leaves inside and outside root, for use in
     checking for informative substitutions */
  if (mod->subtree_root != NULL) {
    inside = lst_new_ptr(mod->tree->nnodes);
    outside = lst_new_ptr(mod->tree->nnodes);
    tr_partition_leaves(mod->tree, mod->subtree_root, inside, outside);
  }

  /* iterate through column tuples */
  for (i = 0; i < msa->ss->ntuples; i++) {
    checkInterruptN(i, 100);

    /* first check for informative substitution data in column; if none,
       don't waste time computing score */
    if (!col_has_data_sub(mod, msa, i, inside, outside)) {
      teststat = 0;
      vec_zero(grad);
      d->params->data[0] = 1.0;
    }

    else {
      d->tupleidx = i;
      vec_set(d->params, 0, d->init_scale);

      opt_newton_1d(col_likelihood_wrapper_1d, &d->params->data[0], d,
                    &lnl, SIGFIGS, d->lb->data[0], d->ub->data[0],
                    logf, NULL, NULL);
      /* turns out to be faster (roughly 15% in limited experiments)
         to use numerical rather than exact derivatives */

      d2->tupleidx = i;
      d2->mod->scale = d->params->data[0];
      d2->mod->scale_sub = 1;
      tm_set_subst_matrices(d2->mod);
      col_scale_derivs_subtree(d2, grad, NULL, d2->fels_scratch);

      fim = col_get_fim_sub(grid, d2->mod->scale);

      teststat = grad->data[1]*grad->data[1] /
        (fim->data[1][1] - fim->data[0][1]*fim->data[1][0]/fim->data[0][0]);

      if (teststat < 0) {
        fprintf(stderr, "WARNING: teststat < 0 (%f\t%f\t%f\t%f\t%f\t%f)\n",
                teststat, fim->data[0][0], fim->data[0][1],
                fim->data[1][0], fim->data[1][1],
                fim->data[0][1]*fim->data[1][0]/fim->data[0][0]);
        teststat = 0;
      }
      mat_free(fim);

      if ((mode == ACC && grad->data[1] < 0) ||
          (mode == CON && grad->data[1] > 0))
        teststat = 0;             /* derivative points toward boundary;
                                     truncate at 0 */
    }

    if (tuple_pvals != NULL) {
      if (mode == NNEUT || mode == CONACC)
        tuple_pvals[i] = chisq_cdf(teststat, 1, FALSE);
      else
        tuple_pvals[i] = half_chisq_cdf(teststat, 1, FALSE);
      /* assumes 50:50 mix of chisq and point mass at zero */

      if (tuple_pvals[i] < 1e-20)
        tuple_pvals[i] = 1e-20;
      /* approx limit of eval of tail prob; pvals of 0 cause problems */

      if (mode == CONACC && grad->data[1] > 0)
        tuple_pvals[i] *= -1; /* mark as acceleration */
    }

    /* store scales and log likelihood ratios if necessary */
    if (tuple_null_scales != NULL) tuple_null_scales[i] = d->params->data[0];
    if (tuple_derivs != NULL) tuple_derivs[i] = grad->data[0];
    if (tuple_sub_derivs != NULL) tuple_sub_derivs[i] = grad->data[1];
    if (tuple_teststats != NULL) tuple_teststats[i] = teststat;
  }

  col_free_fit_data(d);
  col_free_fit_data(d2);
  vec_free(grad);
  modcpy->estimate_branchlens = TM_BRANCHLENS_ALL;
                                /* have to revert for tm_free to work
                                   correctly */
  tm_free(modcpy);
  if (inside != NULL) lst_free(inside);
  if (outside != NULL) lst_free(outside);
  col_free_fim_grid(grid);
}

/* Create object with metadata and scratch memory for fitting scale
   factors */
ColFitData *col_init_fit_data(TreeModel *mod, MSA *msa, scale_type stype,
                              mode_type mode, int second_derivs) {
  ColFitData *d = smalloc(sizeof(ColFitData));
  int size = mod->rate_matrix->size, nrcats = mod->nratecats,
    nnodes = mod->tree->nnodes;
  int nid, rcat, i, j, dim;

  d->mod = mod;
  d->msa = msa;
  d->stype = stype;
  d->mode = mode;
  d->second_derivs = second_derivs;
  d->tupleidx = -1;         /* will be set as needed */

  d->mod->estimate_branchlens = TM_SCALE_ONLY;
  if (stype == SUBTREE) {
    if (!(mod->subtree_root != NULL || mod->in_subtree!=NULL))
      die("ERROR col_init_fit_data: mod->subtree_root or mod->in_subtree must not be NULL in SUBTREE mode\n");
  }
  else {
    if (mod->subtree_root != NULL)
      die("ERROR col_init_fit_data: mod->subtree_root must be NULL if not in subtree mode\n");
    mod->scale_sub = 1;
  }
  if (mod->msa_seq_idx == NULL)
    tm_build_seq_idx(mod, msa);
  tm_set_subst_matrices(mod);

  dim = (stype == ALL ? 1 : 2);
  d->params = vec_new(dim);
  d->lb = vec_new(dim);
  d->ub = vec_new(dim);
  vec_set(d->lb, 0, 0);
  vec_set(d->ub, 0, INFTY);
  d->init_scale = d->init_scale_sub = 1;

  if (stype == ALL) {
    if (mode == CON) {
      vec_set(d->ub, 0, 1);
      d->init_scale = 0.9;      /* don't start on boundary but avoid
                                   strong initialization bias */
    }
    else if (mode == ACC) {
      vec_set(d->lb, 0, 1);
      d->init_scale = 1.1;      /* don't start on boundary but avoid
                                   strong initialization bias */
    }
  }
  else {                        /* stype == SUBTREE */
    vec_set(d->lb, 1, 0);
    vec_set(d->ub, 1, INFTY);
    vec_set(d->lb, 0, 1e-6);    /* can't let this param go quite to
                                   zero, because other param becomes
                                   undefined */

    if (mode == CON) {
      vec_set(d->ub, 1, 1);
      d->init_scale_sub = 0.9;  /* don't start on boundary but avoid
                                   strong initialization bias */
    }
    else if (mode == ACC) {
      vec_set(d->lb, 1, 1);
      d->init_scale_sub = 1.1;  /* don't start on boundary but avoid
                                   strong initialization bias */
    }
  }


  d->PP = smalloc(nnodes * sizeof(void*));
  d->PPP = smalloc(nnodes * sizeof(void*));
  d->QQ = smalloc(nnodes * sizeof(void*));
  d->QQQ = smalloc(nnodes * sizeof(void*));
  d->RRR = smalloc(nnodes * sizeof(void*));

  for (nid = 0; nid < nnodes; nid++) {
    d->PP[nid] = smalloc(nrcats * sizeof(void*));
    d->PPP[nid] = smalloc(nrcats * sizeof(void*));
    d->QQ[nid] = smalloc(nrcats * sizeof(void*));
    d->QQQ[nid] = smalloc(nrcats * sizeof(void*));
    d->RRR[nid] = smalloc(nrcats * sizeof(void*));

    for (rcat = 0; rcat < nrcats; rcat++) {
      d->PP[nid][rcat] = mat_new(size, size);
      d->PPP[nid][rcat] = mat_new(size, size);
      d->QQ[nid][rcat] = mat_new(size, size);
      d->QQQ[nid][rcat] = mat_new(size, size);
      d->RRR[nid][rcat] = mat_new(size, size);

      mat_zero(d->PP[nid][rcat]);
      mat_zero(d->PPP[nid][rcat]);
      mat_zero(d->QQ[nid][rcat]);
      mat_zero(d->QQQ[nid][rcat]);
      mat_zero(d->RRR[nid][rcat]);
    }
  }

  d->expdiag_z = zvec_new(size);
  d->expdiag_r = vec_new(size);
  d->nfels_scratch = stype == SUBTREE ? 6 : 3;
  d->fels_scratch = smalloc(d->nfels_scratch * sizeof(void*));
  for (i = 0; i < d->nfels_scratch; i++) {
    d->fels_scratch[i] = smalloc(size * sizeof(void*));
    for (j = 0; j < size; j++)
      d->fels_scratch[i][j] = smalloc((nnodes+1) * sizeof(double));
  }
  d->mat_scratch_z = zmat_new(size, size);
  d->vec_scratch1_z = zvec_new(size);
  d->vec_scratch2_z = zvec_new(size);
  d->vec_scratch1_r = vec_new(size);
  d->vec_scratch2_r = vec_new(size);
  return d;
}

/* Free metadata and memory for fitting scale factors */
void col_free_fit_data(ColFitData *d) {
  int nid, rcat, i, j;

  vec_free(d->params);
  vec_free(d->lb);
  vec_free(d->ub);

  for (nid = 0; nid < d->mod->tree->nnodes; nid++) {
    for (rcat = 0; rcat < d->mod->nratecats; rcat++) {
      mat_free(d->PP[nid][rcat]);
      mat_free(d->PPP[nid][rcat]);
      mat_free(d->QQ[nid][rcat]);
      mat_free(d->QQQ[nid][rcat]);
      mat_free(d->RRR[nid][rcat]);
    }
    sfree(d->PP[nid]);
    sfree(d->PPP[nid]);
    sfree(d->QQ[nid]);
    sfree(d->QQQ[nid]);
    sfree(d->RRR[nid]);
  }
  sfree(d->PP);
  sfree(d->PPP);
  sfree(d->QQ);
  sfree(d->QQQ);
  sfree(d->RRR);

  zvec_free(d->expdiag_z);
  vec_free(d->expdiag_r);
  for (i = 0; i < d->nfels_scratch; i++) {
    for (j = 0; j < d->mod->rate_matrix->size; j++)
      sfree(d->fels_scratch[i][j]);
    sfree(d->fels_scratch[i]);
  }
  sfree(d->fels_scratch);
  zmat_free(d->mat_scratch_z);
  zvec_free(d->vec_scratch1_z);
  zvec_free(d->vec_scratch2_z);
  vec_free(d->vec_scratch1_r);
  vec_free(d->vec_scratch2_r);

  sfree(d);
}

/* Perform a GERP-like computation for each tuple.  Computes expected
   number of subst. under neutrality (tuple_nneut), expected number
   after rescaling by ML (tuple_nobs), expected number of rejected
   substitutions (tuple_nrejected), and number of species with data
   (tuple_nspecies).  If any arrays are NULL, values will not be
   retained.  Gaps and missing data are handled by working with
   induced subtree.  */
void col_gerp(TreeModel *mod, MSA *msa, mode_type mode, double *tuple_nneut,
              double *tuple_nobs, double *tuple_nrejected,
              double *tuple_nspec, FILE *logf) {
  int i, j, nspec = 0;
  double nneut, scale, lnl;
  int *has_data = smalloc(mod->tree->nnodes * sizeof(int));
  ColFitData *d;

  /* init ColFitData */
  d = col_init_fit_data(mod, msa, ALL, NNEUT, FALSE);

  /* iterate through column tuples */
  for (i = 0; i < msa->ss->ntuples;i++) {
    checkInterruptN(i, 1000);
    col_find_missing_branches(mod, msa, i, has_data, &nspec);

    if (nspec < 3)
      nneut = scale = 0;
    else {
      vec_set(d->params, 0, d->init_scale);
      d->tupleidx = i;

      opt_newton_1d(col_likelihood_wrapper_1d, &d->params->data[0], d,
                    &lnl, SIGFIGS, d->lb->data[0], d->ub->data[0],
                    logf, NULL, NULL);
      /* turns out to be faster (roughly 15% in limited experiments)
         to use numerical rather than exact derivatives */

      scale = d->params->data[0];
      for (j = 1, nneut = 0; j < mod->tree->nnodes; j++)  /* node 0 is root */
        if (has_data[j])
          nneut += ((TreeNode*)lst_get_ptr(mod->tree->nodes, j))->dparent;
    }

    if (tuple_nspec != NULL) tuple_nspec[i] = (double)nspec;
    if (tuple_nneut != NULL) tuple_nneut[i] = nneut;
    if (tuple_nobs != NULL) tuple_nobs[i] = scale * nneut;
    if (tuple_nrejected != NULL) {
      tuple_nrejected[i] = nneut * (1 - scale);
      if (mode == ACC) tuple_nrejected[i] *= -1;
      else if (mode == NNEUT) tuple_nrejected[i] = fabs(tuple_nrejected[i]);
    }
  }
  col_free_fit_data(d);
  sfree(has_data);
}

/* Identify branches wrt which a given column tuple is uninformative,
   in the sense that all leaves beneath these branches having missing
   data.  Will set (preallocated) array has_data[i] = I(branch above
   node i is informative).  Will also set *nspec equal to number of
   leaves that have data. */
void col_find_missing_branches(TreeModel *mod, MSA *msa, int tupleidx,
                               int *has_data, int *nspec) {
  int i;
  List *traversal = tr_postorder(mod->tree);
  *nspec = 0;
  for (i = 0; i < lst_size(traversal); i++) {
    TreeNode *n = lst_get_ptr(traversal, i);
    if (!((n->lchild == NULL && n->rchild == NULL) ||
	  (n->lchild != NULL && n->rchild != NULL)))
      die("ERROR: col_find_missing_branches: lchild and rchild must be either both NULL or both non-NULL\n");
    if (n->parent == NULL)      /* root */
      has_data[n->id] = FALSE;
    else if (n->lchild == NULL) {    /* leaf */
      if (mod->rate_matrix->
          inv_states[(int)ss_get_char_tuple(msa, tupleidx,
                                            mod->msa_seq_idx[n->id], 0)] >= 0) {
        has_data[n->id] = TRUE;
        (*nspec)++;
      }
      else
        has_data[n->id] = FALSE;
    }
    else {                      /* non-root ancestral node */
      if (has_data[n->lchild->id] || has_data[n->rchild->id])
        has_data[n->id] = TRUE;
      else
        has_data[n->id] = FALSE;
    }
  }
}

/* Numerically compute first and second derivatives of single-column
   log likelihood function.  For debugging */
void col_scale_derivs_num(ColFitData *d, double *first_deriv,
                          double *second_deriv) {
  double lnl1, lnl2, lnl3;
  double orig_scale = d->mod->scale;
  d->mod->scale += (2*DERIV_EPSILON);
  tm_set_subst_matrices(d->mod);
  lnl1 = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);
  d->mod->scale = orig_scale + DERIV_EPSILON;
  tm_set_subst_matrices(d->mod);
  lnl2 = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);
  d->mod->scale = orig_scale;
  tm_set_subst_matrices(d->mod);
  lnl3 = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);
  *first_deriv = (lnl2 - lnl3) / DERIV_EPSILON;
  *second_deriv = (lnl1 - 2*lnl2 + lnl3) / (DERIV_EPSILON * DERIV_EPSILON);
}

/* Numerically compute first and second derivatives of single-column
   log likelihood function.  For debugging */
void col_scale_derivs_subtree_num(ColFitData *d, Vector *gradient,
                                  Matrix *hessian) {
  double lnl00, lnl11, lnl0, lnl1, lnl01, lnl;
  double orig_scale = d->mod->scale, orig_scale_sub = d->mod->scale_sub;

  d->mod->scale += (2*DERIV_EPSILON);
  tm_set_subst_matrices(d->mod);
  lnl00 = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);

  d->mod->scale = orig_scale + DERIV_EPSILON;
  tm_set_subst_matrices(d->mod);
  lnl0 = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);

  d->mod->scale = orig_scale + DERIV_EPSILON;
  d->mod->scale_sub = orig_scale_sub + DERIV_EPSILON;
  tm_set_subst_matrices(d->mod);
  lnl01 = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);

  d->mod->scale = orig_scale;
  d->mod->scale_sub = orig_scale_sub + DERIV_EPSILON;
  tm_set_subst_matrices(d->mod);
  lnl1 = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);

  d->mod->scale = orig_scale;
  d->mod->scale_sub = orig_scale_sub + (2*DERIV_EPSILON);
  tm_set_subst_matrices(d->mod);
  lnl11 = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);

  d->mod->scale = orig_scale;
  d->mod->scale_sub = orig_scale_sub;
  tm_set_subst_matrices(d->mod);
  lnl = col_compute_log_likelihood(d->mod, d->msa, d->tupleidx,
                                    d->fels_scratch[0]);

  gradient->data[0] = (lnl0 - lnl) / DERIV_EPSILON;
  gradient->data[1] = (lnl1 - lnl) / DERIV_EPSILON;

  hessian->data[0][0] = (lnl00 - 2*lnl0 + lnl) /
    (DERIV_EPSILON * DERIV_EPSILON);
  hessian->data[1][1] = (lnl11 - 2*lnl1 + lnl) /
    (DERIV_EPSILON * DERIV_EPSILON);
  hessian->data[0][1] = hessian->data[1][0] =
    (lnl01 - lnl0 - lnl1 + lnl) /
    (DERIV_EPSILON * DERIV_EPSILON);
}

/* Estimate 2x2 Fisher Information Matrix (expected value of the
   negative Hessian) for the subtree case, based on a particular value
   of the scale parameter (set in calling code).  Estimation is done
   by sampling. */
Matrix *col_estimate_fim_sub(TreeModel *mod) {
  Vector *grad = vec_new(2);
  Matrix *hessian = mat_new(2, 2), *fim = mat_new(2, 2);
  MSA *msa = tm_generate_msa(NSAMPLES_FIM, NULL, &mod, NULL);
  ColFitData *d = col_init_fit_data(mod, msa, SUBTREE, NNEUT, TRUE);
  int i;

  ss_from_msas(msa, 1, TRUE, NULL, NULL, NULL, -1, 0);
  mat_zero(fim);
  for (i = 0; i < msa->ss->ntuples; i++) {
    d->tupleidx = i;
    col_scale_derivs_subtree(d, grad, hessian, d->fels_scratch);
    mat_scale(hessian, -1 * msa->ss->counts[i]);
                                /* now (observed) Fisher matrix
                                   weighted by count */
    mat_plus_eq(fim, hessian);   /* add to running total */
  }
  mat_scale(fim, 1.0/NSAMPLES_FIM);   /* convert total to sample mean */

  msa_free(msa);
  col_free_fit_data(d);
  vec_free(grad);
  mat_free(hessian);
  return (fim);
}

/* Precompute estimates of FIM for a grid of possible scale params
   (subtree case) */
FimGrid *col_fim_grid_sub(TreeModel *mod) {
  int i;
  FimGrid *g = smalloc(sizeof(FimGrid));

  g->ngrid1 = (int)(1.0/GRIDSIZE1);
  g->ngrid2 = (int)((1.0 * GRIDMAXLOG / GRIDSIZE2) + 1);
  g->ngrid = g->ngrid1 + g->ngrid2;
  g->scales = smalloc(g->ngrid * sizeof(double));

  mod->scale_sub = 1;

  for (i = 0; i < g->ngrid1; i++)
    g->scales[i] = i * GRIDSIZE1;

  for (i = 0; i < g->ngrid2; i++)
    g->scales[g->ngrid1 + i] = exp(i * GRIDSIZE2);

  g->fim = smalloc(g->ngrid * sizeof(void*));
  for (i = 0; i < g->ngrid; i++) {
    mod->scale = g->scales[i];
    tm_set_subst_matrices(mod);
    g->fim[i] = col_estimate_fim_sub(mod);
  }

  return g;
}

/* free FimGrid object */
void col_free_fim_grid(FimGrid *g) {
  int i;
  for (i = 0; i < g->ngrid; i++)
    mat_free(g->fim[i]);
  sfree(g->fim);
  sfree(g->scales);
}

/* Estimate scale Fisher Information Matrix for the non-subtree case.
   This version does not depend on any free parameters, so no grid is
   required.  Estimation is done by sampling, as above */
double col_estimate_fim(TreeModel *mod) {
  double deriv1, deriv2, retval = 0;
  MSA *msa = tm_generate_msa(NSAMPLES_FIM, NULL, &mod, NULL);
  ColFitData *d = col_init_fit_data(mod, msa, ALL, NNEUT, FALSE);
  int i;

  ss_from_msas(msa, 1, TRUE, NULL, NULL, NULL, -1, 0);

  for (i = 0; i < msa->ss->ntuples; i++) {
    d->tupleidx = i;
    col_scale_derivs(d, &deriv1, &deriv2, d->fels_scratch);
    retval += (-deriv2 * msa->ss->counts[i]); /* add (observed) Fisher matrix
                                                 weighted by count */
  }
  retval /= NSAMPLES_FIM;   /* convert total to sample mean */

  msa_free(msa);
  col_free_fit_data(d);
  return (retval);
}

/* Retrieve estimated FIM for given scale function; uses linear
   interpolation from precomputed grid */
Matrix *col_get_fim_sub(FimGrid *g, double scale) {
  int idx;
  double frac;
  Matrix *retval;

  if (scale < 0)
    die("ERROR col_get_fix_sub: scale should be >= 0 but is %e\n", scale);

  if (scale < 1)
    idx = (int)floor(scale / GRIDSIZE1);
  else
    idx = g->ngrid1 + (int)floor(log(scale) / GRIDSIZE2);

  if (idx >= g->ngrid - 1)
    retval = mat_create_copy(g->fim[g->ngrid - 1]);
                                /* just use last one in this case */

  else {
    if (!(g->scales[idx] <= scale && g->scales[idx+1] > scale))
      die("ERROR col_get_fim_sub: g->scales[%i]=%e should be <= %e and g->scales[%i+1]=%e should be > %e\n", idx, g->scales[idx], scale, idx+1, g->scales[idx+1], scale);

    frac = (scale - g->scales[idx]) / (g->scales[idx+1] - g->scales[idx]);

    if (frac < 0.1)
      retval = mat_create_copy(g->fim[idx]);
    else if (frac > 0.9)
      retval = mat_create_copy(g->fim[idx+1]);
    else {  /* interpolate */
      retval = mat_new(2,2);
      mat_linear_comb(retval, g->fim[idx], frac, g->fim[idx+1], 1-frac);
    }
  }
  return retval;
}

/* returns TRUE if column has two or more actual bases (not gaps or
   missing data), otherwise returns FALSE */
int col_has_data(TreeModel *mod, MSA *msa, int tupleidx) {
  int i, nbases = 0;
  for (i = 0; i < msa->nseqs && nbases < 2; i++) {
    int state = mod->rate_matrix->
      inv_states[(int)ss_get_char_tuple(msa, tupleidx, i, 0)];
    if (state >= 0)
      nbases++;
  }
  return(nbases >= 2);
}

/* returns TRUE if column has at least one base in the subtree of
   interest, at least one in the supertree of interest, and at least
   three bases total (the minimum required for a meaningful subtree
   test), otherwise returns FALSE.
   If inside and outside are both NULL, then returns TRUE if there
   are at least three bases total.
 */
int col_has_data_sub(TreeModel *mod, MSA *msa, int tupleidx, List *inside,
                     List *outside) {
  int i, nbases = 0, state;
  TreeNode *n;

  if (inside == NULL && outside == NULL) {
    for (i=0; i<mod->tree->nnodes; i++) {
      n = lst_get_ptr(mod->tree->nodes, i);
      if (n->lchild == NULL) {
	state = mod->rate_matrix->
	  inv_states[(int)ss_get_char_tuple(msa, tupleidx,
					    mod->msa_seq_idx[n->id], 0)];
	if (state >=0) nbases++;
	if (nbases==2) return TRUE;
      }
    }
    return FALSE;
  }
  for (i = 0; i < lst_size(inside) && nbases < 2; i++) {
    n = lst_get_ptr(inside, i);
    state = mod->rate_matrix->
      inv_states[(int)ss_get_char_tuple(msa, tupleidx,
                                        mod->msa_seq_idx[n->id], 0)];
    if (state >= 0) nbases++;
  }

  if (nbases == 0)              /* no leaves in subtree */
    return FALSE;
  /* NOTE: nbases at most two here */

  for (i = 0; i < lst_size(outside) && nbases < 3; i++) {
    n = lst_get_ptr(outside, i);
    state = mod->rate_matrix->
      inv_states[(int)ss_get_char_tuple(msa, tupleidx,
                                        mod->msa_seq_idx[n->id], 0)];
    if (state >= 0) nbases++;
  }

  if (nbases == 3)              /* has to be at least one in each
                                   partition and a third from one of
                                   the two partitions */
    return TRUE;

  return FALSE;
}
