text: // transform/fmllr-diag-gmm.cc
  
  // Copyright 2009-2011  Microsoft Corporation;  Saarland University;
  //                      Georg Stemmer
  //                2013  Johns Hopkins University (author: Daniel Povey)
  
  // See ../../COPYING for clarification regarding multiple authors
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //  http://www.apache.org/licenses/LICENSE-2.0
  //
  // THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  // WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  // MERCHANTABLITY OR NON-INFRINGEMENT.
  // See the Apache 2 License for the specific language governing permissions and
  // limitations under the License.
  
  #include <utility>
  #include <vector>
  using std::vector;
  
  #include "transform/fmllr-diag-gmm.h"
  
  namespace kaldi {
  
  void FmllrDiagGmmAccs:: AccumulateFromPosteriors(
      const DiagGmm &pdf,
      const VectorBase<BaseFloat> &data,
      const VectorBase<BaseFloat> &posterior) {
    
    if (this->DataHasChanged(data)) {
      CommitSingleFrameStats();
      InitSingleFrameStats(data);
    }
    SingleFrameStats &stats = this->single_frame_stats_;
    stats.count += posterior.Sum();
    stats.a.AddMatVec(1.0, pdf.means_invvars(), kTrans, posterior, 1.0);
    stats.b.AddMatVec(1.0, pdf.inv_vars(), kTrans, posterior, 1.0);
  }
  
  void FmllrDiagGmmAccs:: AccumulateFromPosteriorsPreselect(
      const DiagGmm &pdf,
      const std::vector<int32> &gselect,
      const VectorBase<BaseFloat> &data,
      const VectorBase<BaseFloat> &posterior) {
    
    if (this->DataHasChanged(data)) {
      CommitSingleFrameStats();
      InitSingleFrameStats(data);
    }
    SingleFrameStats &stats = this->single_frame_stats_;
    stats.count += posterior.Sum();
    
    const Matrix<BaseFloat> &means_invvars = pdf.means_invvars(),
        &inv_vars = pdf.inv_vars();
    KALDI_ASSERT(static_cast<int32>(gselect.size()) == posterior.Dim());
    for (size_t i = 0; i < gselect.size(); i++) {
      stats.a.AddVec(posterior(i), means_invvars.Row(gselect[i]));
      stats.b.AddVec(posterior(i), inv_vars.Row(gselect[i]));
    }
  }
  
  FmllrDiagGmmAccs::FmllrDiagGmmAccs(const DiagGmm &gmm,
                                     const AccumFullGmm &fgmm_accs):
      single_frame_stats_(gmm.Dim()), opts_(FmllrOptions()) {
    KALDI_ASSERT(gmm.NumGauss() == fgmm_accs.NumGauss()
                 && gmm.Dim() == fgmm_accs.Dim());
    Init(gmm.Dim());
    int32 dim = gmm.Dim(), num_gauss = gmm.NumGauss();
    for (int32 g = 0; g < num_gauss; g++) {
      double this_occ = fgmm_accs.occupancy()(g);
      if (this_occ == 0) continue;
      SubVector<BaseFloat> this_mean_invvar(gmm.means_invvars(), g);
      SubVector<BaseFloat> this_invvar(gmm.inv_vars(), g);
      SubVector<double> this_mean_acc(fgmm_accs.mean_accumulator(), g);
      Vector<double> this_mean_invvar_dbl(this_mean_invvar);
      Vector<double> this_extended_mean_acc(dim+1);
      this_extended_mean_acc.Range(0, dim).CopyFromVec(this_mean_acc);
      this_extended_mean_acc(dim) = this_occ; // acc of x^+
      Matrix<double> this_cov_acc(fgmm_accs.covariance_accumulator()[g]); // copy to
      // regular Matrix.
      Matrix<double> this_extended_cov_acc(dim+1, dim+1); // make as if accumulated
      // using x^+, not x.
      this_extended_cov_acc.Range(0, dim, 0, dim).CopyFromMat(this_cov_acc);
      this_extended_cov_acc.Row(dim).CopyFromVec(this_extended_mean_acc);
      this_extended_cov_acc.CopyColFromVec(this_extended_mean_acc, dim); // since
      // there is no Col() function, use a member-function of the matrix class.
      SpMatrix<double> this_extended_cov_acc_sp(this_extended_cov_acc);
      beta_ += this_occ;
      K_.AddVecVec(1.0, this_mean_invvar_dbl, this_extended_mean_acc);
      for (int32 d = 0; d < dim; d++)
        G_[d].AddSp(this_invvar(d), this_extended_cov_acc_sp);
    }
  }
  
  
  BaseFloat FmllrDiagGmmAccs::AccumulateForGmm(const DiagGmm &pdf,
                                               const VectorBase<BaseFloat> &data,
                                               BaseFloat weight) {
    int32 num_comp = pdf.NumGauss();
    Vector<BaseFloat> posterior(num_comp);
    BaseFloat loglike;
    
    loglike = pdf.ComponentPosteriors(data, &posterior);
    posterior.Scale(weight);
    AccumulateFromPosteriors(pdf, data, posterior);
    return loglike;
  }
  
  BaseFloat FmllrDiagGmmAccs::AccumulateForGmmPreselect(
      const DiagGmm &pdf,
      const std::vector<int32> &gselect,
      const VectorBase<BaseFloat> &data,
      BaseFloat weight) {
    KALDI_ASSERT(!gselect.empty() && "Empty gselect information");
    Vector<BaseFloat> loglikes;
    pdf.LogLikelihoodsPreselect(data, gselect, &loglikes);
  
    BaseFloat loglike = loglikes.ApplySoftMax(); // they are now posteriors.
    loglikes.Scale(weight);
    AccumulateFromPosteriorsPreselect(pdf, gselect, data, loglikes);
    return loglike;
  }
  
  
  
  void FmllrDiagGmmAccs::Update(const FmllrOptions &opts,
                                MatrixBase<BaseFloat> *fmllr_mat,
                                BaseFloat *objf_impr,
                                BaseFloat *count) {
    KALDI_ASSERT(fmllr_mat != NULL);
    CommitSingleFrameStats();
    if (fmllr_mat->IsZero())
      KALDI_ERR << "You must initialize the fMLLR matrix to a non-singular value "
          "(so we can report objective function changes); e.g. call SetUnit()";
    if (opts.update_type == "full" && this->opts_.update_type != "full") {
      KALDI_ERR << "You are requesting a full-fMLLR update but you accumulated "
                << "stats for more limited update type.";
    }
    if (beta_ > opts.min_count) {
      Matrix<BaseFloat> tmp_old(*fmllr_mat), tmp_new(*fmllr_mat);
      BaseFloat objf_change;
      if (opts.update_type == "full")
        objf_change = ComputeFmllrMatrixDiagGmmFull(tmp_old, *this, opts.num_iters, &tmp_new);
      else if (opts.update_type == "diag")
        objf_change = ComputeFmllrMatrixDiagGmmDiagonal(tmp_old, *this, &tmp_new);
      else if (opts.update_type == "offset")
        objf_change = ComputeFmllrMatrixDiagGmmOffset(tmp_old, *this, &tmp_new);
      else if (opts.update_type == "none")
        objf_change = 0.0;
      else
        KALDI_ERR << "Unknown fMLLR update type " << opts.update_type
                  << ", fmllr-update-type must be one of \"full\"|\"diag\"|\"offset\"|\"none\"";
      fmllr_mat->CopyFromMat(tmp_new);
      if (objf_impr) *objf_impr = objf_change;
      if (count) *count = beta_;
    } else {  // Not changing matrix.
      KALDI_WARN << "Not updating fMLLR since below min-count: count is " << beta_;
      if (objf_impr) *objf_impr = 0.0;
      if (count) *count = beta_;
    }
  }
  
  
  BaseFloat ComputeFmllrMatrixDiagGmm(const MatrixBase<BaseFloat> &in_xform,
                                      const AffineXformStats &stats,
                                      std::string fmllr_type,  // "none", "offset", "diag", "full"
                                      int32 num_iters,
                                      MatrixBase<BaseFloat> *out_xform) {
    if (fmllr_type == "full") {
      return ComputeFmllrMatrixDiagGmmFull(in_xform, stats, num_iters, out_xform);
    } else if (fmllr_type == "diag") {
      return ComputeFmllrMatrixDiagGmmDiagonal(in_xform, stats, out_xform);
    } else if (fmllr_type == "offset") {
      return ComputeFmllrMatrixDiagGmmOffset(in_xform, stats, out_xform);
    } else if (fmllr_type == "none") {
      if (!in_xform.IsUnit())
        KALDI_WARN << "You set fMLLR type to \"none\" but your starting transform "
            "is not unit [this is strange, and diagnostics will be wrong].";
      out_xform->SetUnit();
      return 0.0;
    } else
      KALDI_ERR << "Unknown fMLLR update type " << fmllr_type
                << ", must be one of \"full\"|\"diag\"|\"offset\"|\"none\"";
    return 0.0;
  }
  
  
  void FmllrInnerUpdate(SpMatrix<double> &inv_G,
                        VectorBase<double> &k,
                        double beta,
                        int32 row,
                        MatrixBase<double> *transform) {
    int32 dim = transform->NumRows();
    KALDI_ASSERT(transform->NumCols() == dim + 1);
    KALDI_ASSERT(row >= 0 && row < dim);
  
    double logdet;
    // Calculating the matrix of cofactors (transpose of adjugate)
    Matrix<double> cofact_mat(dim, dim);
    cofact_mat.CopyFromMat(transform->Range(0, dim, 0, dim), kTrans);
    cofact_mat.Invert(&logdet);
    // Removed this step because it's not necessary and could lead to
    // under/overflow [Dan]
    // cofact_mat.Scale(exp(logdet));
    
    // The extended cofactor vector for the current row
    Vector<double> cofact_row(dim + 1);
    cofact_row.Range(0, dim).CopyRowFromMat(cofact_mat, row);
    cofact_row(dim) = 0;
    Vector<double> cofact_row_invg(dim + 1);
    cofact_row_invg.AddSpVec(1.0, inv_G, cofact_row, 0.0);
  
    // Solve the quadratic equation for step size
    double e1 = VecVec(cofact_row_invg, cofact_row);
    double e2 = VecVec(cofact_row_invg, k);
    double discr = std::sqrt(e2 * e2 + 4 * e1 * beta);
    double alpha1 = (-e2 + discr) / (2 * e1);
    double alpha2 = (-e2 - discr) / (2 * e1);
    double auxf1 = beta * Log(std::abs(alpha1 * e1 + e2)) -
        0.5 * alpha1 * alpha1 * e1;
    double auxf2 = beta * Log(std::abs(alpha2 * e1 + e2)) -
        0.5 * alpha2 * alpha2 * e1;
    double alpha = (auxf1 > auxf2) ? alpha1 : alpha2;
  
    // Update transform row: w_d = (\alpha cofact_d + k_d) G_d^{-1}
    cofact_row.Scale(alpha);
    cofact_row.AddVec(1.0, k);
    transform->Row(row).AddSpVec(1.0, inv_G, cofact_row, 0.0);
  }
  
  BaseFloat ComputeFmllrMatrixDiagGmmFull(const MatrixBase<BaseFloat> &in_xform,
                                          const AffineXformStats &stats,
                                          int32 num_iters,
                                          MatrixBase<BaseFloat> *out_xform) {
    int32 dim = static_cast<int32>(stats.G_.size());
  
    // Compute the inverse matrices of second-order statistics
    vector< SpMatrix<double> > inv_g(dim);
    for (int32 d = 0; d < dim; d++) {
      inv_g[d].Resize(dim + 1);
      inv_g[d].CopyFromSp(stats.G_[d]);
      inv_g[d].Invert();
    }
  
    Matrix<double> old_xform(in_xform), new_xform(in_xform);
    BaseFloat old_objf = FmllrAuxFuncDiagGmm(old_xform, stats);
    
    for (int32 iter = 0; iter < num_iters; ++iter) {
      for (int32 d = 0; d < dim; d++) {
        SubVector<double> k_d(stats.K_, d);
        FmllrInnerUpdate(inv_g[d], k_d, stats.beta_, d, &new_xform);
      }  // end of looping over rows
    }  // end of iterations
  
    BaseFloat new_objf = FmllrAuxFuncDiagGmm(new_xform, stats),
        objf_improvement = new_objf - old_objf;
    KALDI_LOG << "fMLLR objf improvement is "
              << (objf_improvement / (stats.beta_ + 1.0e-10))
              << " per frame over " << stats.beta_ << " frames.";
    if (objf_improvement < 0.0 && !ApproxEqual(new_objf, old_objf)) {
      KALDI_WARN << "No applying fMLLR transform change because objective "
                 << "function did not increase.";
      return 0.0;
    } else {
      out_xform->CopyFromMat(new_xform, kNoTrans);
      return objf_improvement;
    }
  }
  
  BaseFloat ComputeFmllrMatrixDiagGmmDiagonal(const MatrixBase<BaseFloat> &in_xform,
                                              const AffineXformStats &stats,
                                              MatrixBase<BaseFloat> *out_xform) {
    // The "Diagonal" here means a diagonal fMLLR matrix, i.e. like W = [ A;  b] where
    // A is diagonal.
    // We re-derived the math (see exponential transform paper) to get a simpler
    // update rule.
  
    /*
    Write out_xform as D, which is a d x d+1 matrix (where d is the feature dimension).
    We are solving for s == d_{i,i}, and o == d_{i,d}  [assuming zero-based indexing];
        s is a scale, o is an offset.
    The stats are K (dimension d x d+1) and G_i for i=0..d-1 (dimension: d+1 x d+1),
      and the count beta.
  
   The auxf for the i'th row of the transform is (assuming zero-based indexing):
  
    s k_{i,i}  +  o k_{i,d}
    - \frac{1}{2} s^2 g_{i,i,i} - \frac{1}{2} o^2 g_{i,d,d} - s o g_{i,d,i}
     + \beta \log |s|
  
     Suppose we know s, we can solve for o:
        o = (k_{i,d} - s g_{i,d,i}) / g_{i,d,d}
     Substituting this expression for o into the auxf (and ignoring
     terms that don't vary with s), we have the auxf:
  
   \frac{1}{2} s^2 ( g_{i,d,i}^2 / g_{i,d,d}  -  g_{i,i,i} )
      +  s ( k_{i,i} - g_{i,d,i} k_{i,d} / g_{i,d,d} )
      + \beta \log |s|
  
    Differentiating w.r.t. s and assuming s is positive, we have
      a s + b + c/s = 0
   where
     a = (  g_{i,d,i}^2 / g_{i,d,d}  -  g_{i,i,i} ),
     b = ( k_{i,i} - g_{i,d,i} k_{i,d} / g_{i,d,d} )
     c = beta
   Multiplying by s, we have the equation
     a s^2 + b s + c = 0, where we assume s > 0.
   We solve it with:
    s = (-b - \sqrt{b^2 - 4ac}) / 2a
   [take the negative root because we know a is negative, and this gives
    the more positive solution for s; the other one would be negative].
   We then solve for o with the equation above, i.e.:
       o = (k_{i,d} - s g_{i,d,i}) / g_{i,d,d})
    */
  
    int32 dim = stats.G_.size();
    double beta = stats.beta_;
    out_xform->CopyFromMat(in_xform);
    if (beta == 0.0) {
      KALDI_WARN << "Computing diagonal fMLLR matrix: no stats [using original transform]";
      return 0.0;
    }
    BaseFloat old_obj = FmllrAuxFuncDiagGmm(*out_xform, stats);
    KALDI_ASSERT(out_xform->Range(0, dim, 0, dim).IsDiagonal()); // orig transform
    // must be diagonal.
    for(int32 i = 0; i < dim; i++) {
      double k_ii = stats.K_(i, i), k_id = stats.K_(i, dim),
          g_iii = stats.G_[i](i, i), g_idd = stats.G_[i](dim, dim),
          g_idi = stats.G_[i](dim, i);
      double a = g_idi*g_idi/g_idd - g_iii,
          b = k_ii - g_idi*k_id/g_idd,
          c = beta;
      double s = (-b - std::sqrt(b*b - 4*a*c)) / (2*a);
      KALDI_ASSERT(s > 0.0);
      double o = (k_id - s*g_idi) / g_idd;
      (*out_xform)(i, i) = s;
      (*out_xform)(i, dim) = o;
    }
    BaseFloat new_obj = FmllrAuxFuncDiagGmm(*out_xform, stats);
    KALDI_VLOG(2) << "fMLLR objective function improvement = "
                  << (new_obj - old_obj);
    return new_obj - old_obj;
  }
  
  BaseFloat ComputeFmllrMatrixDiagGmmOffset(const MatrixBase<BaseFloat> &in_xform,
                                            const AffineXformStats &stats,
                                            MatrixBase<BaseFloat> *out_xform) {
    int32 dim = stats.G_.size();
    KALDI_ASSERT(in_xform.NumRows() == dim && in_xform.NumCols() == dim+1);
    SubMatrix<BaseFloat> square_part(in_xform, 0, dim, 0, dim);
    KALDI_ASSERT(square_part.IsUnit());
    BaseFloat objf_impr = 0.0;
    out_xform->CopyFromMat(in_xform);
    for (int32 i = 0; i < dim; i++) {
      // auxf in this offset b_i is:
      //  -0.5 b_i^2 G_i(dim, dim) - b_i G_i(i, dim)*1.0 + b_i K(i, dim)  (1)
      // answer is:
      // b_i = [K(i, dim) - G_i(i, dim)] / G_i(dim, dim)
      // objf change is given by (1)
      BaseFloat b_i = (*out_xform)(i, dim);
      BaseFloat objf_before = -0.5 * b_i * b_i * stats.G_[i](dim, dim)
          - b_i * stats.G_[i](i, dim) + b_i * stats.K_(i, dim);
      b_i = (stats.K_(i, dim) - stats.G_[i](i, dim)) / stats.G_[i](dim, dim);
      (*out_xform)(i, dim) = b_i;
      BaseFloat objf_after = -0.5 * b_i * b_i * stats.G_[i](dim, dim)
          - b_i * stats.G_[i](i, dim) + b_i * stats.K_(i, dim);
      if (objf_after < objf_before)
        KALDI_WARN << "Objf decrease in offset estimation:"
                   << objf_after << " < " << objf_before;
      objf_impr += objf_after - objf_before;
    }
    return objf_impr;
  }
  
  
  void ApplyFeatureTransformToStats(const MatrixBase<BaseFloat> &xform,
                                    AffineXformStats *stats) {
    KALDI_ASSERT(stats != NULL && stats->Dim() != 0);
    int32 dim = stats->Dim();
    // make sure the stats are of the standard diagonal kind.
    KALDI_ASSERT(stats->G_.size() == static_cast<size_t>(dim));
    KALDI_ASSERT( (xform.NumRows() == dim && xform.NumCols() == dim) // linear
                  || (xform.NumRows() == dim && xform.NumCols() == dim+1) // affine
                  || (xform.NumRows() == dim+1 && xform.NumCols() == dim+1));  // affine w/ extra row.
    if (xform.NumRows() == dim+1) {  // check last row of input
      // has correct value. 0 0 0 ..  0 1.
      for (int32 i = 0; i < dim; i++)
        KALDI_ASSERT(xform(dim, i) == 0.0);
      KALDI_ASSERT(xform(dim, dim) == 1.0);
    }
  
    // Get the transform into the (dim+1 x dim+1) format, with
    // 0 0 0 .. 0 1 as the last row.
    SubMatrix<BaseFloat> xform_square(xform, 0, dim, 0, dim);
    Matrix<double> xform_full(dim+1, dim+1);
    SubMatrix<double> xform_full_square(xform_full, 0, dim, 0, dim);
    xform_full_square.CopyFromMat(xform_square);
    if (xform.NumCols() == dim+1)  // copy offset.
      for (int32 i = 0; i < dim; i++)
        xform_full(i, dim) = xform(i, dim);
  
    xform_full(dim, dim) = 1.0;
  
    SpMatrix<double> Gtmp(dim+1);
    for (int32 i = 0; i < dim; i++) {
      // Gtmp <-- xform_full * stats->G_[i] * xform_full^T
      Gtmp.AddMat2Sp(1.0, xform_full, kNoTrans, stats->G_[i], 0.0);
      stats->G_[i].CopyFromSp(Gtmp);
    }
    Matrix<double> Ktmp(dim, dim+1);
    // Ktmp <-- stats->K_ * xform_full^T
    Ktmp.AddMatMat(1.0, stats->K_, kNoTrans, xform_full, kTrans, 0.0);
    stats->K_.CopyFromMat(Ktmp);
  }
  
  void ApplyModelTransformToStats(const MatrixBase<BaseFloat> &xform,
                                  AffineXformStats *stats) {
    KALDI_ASSERT(stats != NULL && stats->Dim() != 0.0);
    int32 dim = stats->Dim();
    KALDI_ASSERT(xform.NumRows() == dim && xform.NumCols() == dim+1);
    {
      SubMatrix<BaseFloat> xform_square(xform, 0, dim, 0, dim);
      // Only works with diagonal transforms.
      KALDI_ASSERT(xform_square.IsDiagonal());
    }
  
    // Working out rules for transforming fMLLR statistics under diagonal
    // model-space transformations.
    //
    // We work out what the stats would be if we had accumulated
    // with offset/scaled means and vars. Let T be the transform
    // T = [ D; b ],
    // where D is diagonal, d_i is the i'th diagonal of D, and b_i
    // is the i'th offset element.  This is equivalent to the transform
    //   x_i -> y_i = d_i x_i + b_i,
    // so d_i is the diagonal and b_i is the offset term.  We work out the
    // reverse feature transform (from general to speaker-specific space),
    // which is
    //  y_i -> x_i = (y_i - b_i) / d_i
    // the corresponding mean transform to speaker-space is the same:
    //  mu_i -> (mu_i - b_i) / d_i
    // and the transfrom on the variances is:
    //  sigma_i^2 -> sigma_i^2 / d_i^2,
    // so on inverse variance this becomes:
    //  (1/sigma_i^2) -> (1/sigma_i^2) * d_i^2.
    //
    // Now, we work out the change in K and G_i from these effects on the
    // means and variances.
    //
    // Now, k_{ij} is \sum_{m, t} \gamma_m (1/\sigma^2_{m, i}) \mu_{m, i} x^+_j .
    //
    // If we are transforming to K', we want:
    //
    // k'_{ij} = \sum_{m, t} \gamma_m (d_i^2/\sigma^2_{m, i}) ((\mu_{m, i} - b_i)/d_i)  x^+_j .
    //         = d_i k_{i, j} - \sum_{m, t} \gamma_m (1/\sigma^2_{m, i}) d_i b_i x^+_j .
    //         = d_i k_{i, j} - d_i b_i g_{i, d, j},
    // where g_{i, d, j} is the {d, j}'th element of G_i. (in zero-based indexing).
    //
    //
    // G_i only depends on the variances and features, so the only change
    // in G_i is G_i -> d_i^2 G_i (this comes from the change in 1/sigma_i^2).
    // This is done after the change in K.
  
    for (int32 i = 0; i < dim; i++) {
      BaseFloat d_i = xform(i, i), b_i = xform(i, dim);
      for (int32 j = 0; j <= dim; j++) {
        stats->K_(i, j) = d_i * stats->K_(i, j) - d_i * b_i * stats->G_[i](dim, j);
      }
    }
    for (int32 i = 0; i < dim; i++) {
      BaseFloat d_i = xform(i, i);
      stats->G_[i].Scale(d_i * d_i);
    }
  }
  
  float FmllrAuxFuncDiagGmm(const MatrixBase<float> &xform,
                                const AffineXformStats &stats) {
    int32 dim = static_cast<int32>(stats.G_.size());
    Matrix<double> xform_d(xform);
    Vector<double> xform_row_g(dim + 1);
    SubMatrix<double> A(xform_d, 0, dim, 0, dim);
    double obj = stats.beta_ * A.LogDet() +
        TraceMatMat(xform_d, stats.K_, kTrans);
    for (int32 d = 0; d < dim; d++) {
      xform_row_g.AddSpVec(1.0, stats.G_[d], xform_d.Row(d), 0.0);
      obj -= 0.5 * VecVec(xform_row_g, xform_d.Row(d));
    }
    return obj;
  }
  
  double FmllrAuxFuncDiagGmm(const MatrixBase<double> &xform,
                             const AffineXformStats &stats) {
    int32 dim = static_cast<int32>(stats.G_.size());
    Vector<double> xform_row_g(dim + 1);
    SubMatrix<double> A(xform, 0, dim, 0, dim);
    double obj = stats.beta_ * A.LogDet() +
        TraceMatMat(xform, stats.K_, kTrans);
    for (int32 d = 0; d < dim; d++) {
      xform_row_g.AddSpVec(1.0, stats.G_[d], xform.Row(d), 0.0);
      obj -= 0.5 * VecVec(xform_row_g, xform.Row(d));
    }
    return obj;
  }
  
  BaseFloat FmllrAuxfGradient(const MatrixBase<BaseFloat> &xform,
                              // if this is changed back to Matrix<double>
                             // un-comment the Resize() below.
                              const AffineXformStats &stats,
                              MatrixBase<BaseFloat> *grad_out) {
    int32 dim = static_cast<int32>(stats.G_.size());
    Matrix<double> xform_d(xform);
    Vector<double> xform_row_g(dim + 1);
    SubMatrix<double> A(xform_d, 0, dim, 0, dim);
    double obj = stats.beta_ * A.LogDet() +
        TraceMatMat(xform_d, stats.K_, kTrans);
    Matrix<double> S(dim, dim + 1);
    for (int32 d = 0; d < dim; d++) {
      xform_row_g.AddSpVec(1.0, stats.G_[d], xform_d.Row(d), 0.0);
      obj -= 0.5 * VecVec(xform_row_g, xform_d.Row(d));
      S.CopyRowFromVec(xform_row_g, d);
    }
  
    // Compute the gradient: P = \beta [(A^{-1})^{T} , 0] + K - S
    // grad_out->Resize(dim, dim + 1);
    Matrix<double> tmp_grad(dim, dim + 1);
    tmp_grad.Range(0, dim, 0, dim).CopyFromMat(A);
    tmp_grad.Range(0, dim, 0, dim).Invert();
    tmp_grad.Range(0, dim, 0, dim).Transpose();
    tmp_grad.Scale(stats.beta_);
    tmp_grad.AddMat(-1.0, S, kNoTrans);
    tmp_grad.AddMat(1.0, stats.K_, kNoTrans);
    grad_out->CopyFromMat(tmp_grad, kNoTrans);
  
    return obj;
  }
  
  bool FmllrDiagGmmAccs::DataHasChanged(const VectorBase<BaseFloat> &data) const {
    KALDI_ASSERT(data.Dim() == this->Dim());
    return !data.ApproxEqual(single_frame_stats_.x, 0.0);
  }
  
  void FmllrDiagGmmAccs::SingleFrameStats::Init(int32 dim) {
    x.Resize(dim);
    a.Resize(dim);
    b.Resize(dim);
    count = 0.0;
  }  
  
  void FmllrDiagGmmAccs::InitSingleFrameStats(const VectorBase<BaseFloat> &data) {
    SingleFrameStats &stats = single_frame_stats_;
    stats.x.CopyFromVec(data);
    stats.count = 0.0;
    stats.a.SetZero();
    stats.b.SetZero();
  }
  
  void FmllrDiagGmmAccs::CommitSingleFrameStats() {
    // Commit the stats for this from (in SingleFrameStats).
    int32 dim = Dim();
    SingleFrameStats &stats = single_frame_stats_;
    if (stats.count == 0.0) return;
  
    Vector<double> xplus(dim+1);
    xplus.Range(0, dim).CopyFromVec(stats.x);
    xplus(dim) = 1.0;
    
    this->beta_ += stats.count;
    this->K_.AddVecVec(1.0, Vector<double>(stats.a), xplus);
  
  
    if (opts_.update_type == "full") {
      SpMatrix<double> scatter(dim+1);
      scatter.AddVec2(1.0, xplus);
      
      KALDI_ASSERT(static_cast<size_t>(dim) == this->G_.size());
      for (int32 i = 0; i < dim; i++)
        this->G_[i].AddSp(stats.b(i), scatter);
    } else {
      // We only need some elements of these stats, so just update those elements.
      for (int32 i = 0; i < dim; i++) {
        BaseFloat scale = stats.b(i), x_i = xplus(i);
        this->G_[i](i, i) += scale * x_i * x_i;
        this->G_[i](dim, i) += scale * 1.0 * x_i;
        this->G_[i](dim, dim) += scale * 1.0 * 1.0;
      }
    }
  
    stats.count = 0.0;
    stats.a.SetZero();
    stats.b.SetZero();
  }
      
  
  
  
  } // namespace kaldi
  