text: // sgmm2/am-sgmm2.cc
  
  // Copyright 2009-2011  Microsoft Corporation;  Lukas Burget;
  //                      Saarland University (Author: Arnab Ghoshal);
  //                      Ondrej Glembek;  Yanmin Qian;
  // Copyright 2012-2013  Johns Hopkins University (Author: Daniel Povey)
  //                      Liang Lu;  Arnab Ghoshal
  
  // 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 <functional>
  
  #include "sgmm2/am-sgmm2.h"
  #include "util/kaldi-thread.h"
  
  namespace kaldi {
  using std::vector;
  
  // This function needs to be added because std::generate is complaining
  // about RandGauss(), which takes an optional arguments.
  static inline float _RandGauss()
  {
    return RandGauss();
  }
  
  void Sgmm2LikelihoodCache::NextFrame() {
    t++;
    if (t == 0) {
      t++; // skip over zero; zero is used to invalidate frames.
      for (size_t i = 0; i < substate_cache.size(); i++)
        substate_cache[i].t = 0;
      for (size_t i = 0; i < pdf_cache.size(); i++)
        pdf_cache[i].t = 0;
    }
  }
  
  void AmSgmm2::ComputeGammaI(const Vector<BaseFloat> &state_occupancies,
                              Vector<BaseFloat> *gamma_i) const {
    KALDI_ASSERT(state_occupancies.Dim() == NumPdfs());
    Vector<BaseFloat> w_jm(NumGauss());
    gamma_i->Resize(NumGauss());
    for (int32 j1 = 0; j1 < NumGroups(); j1++) {
      int32 M = NumSubstatesForGroup(j1);
      const std::vector<int32> &pdfs = group2pdf_[j1];
      Vector<BaseFloat> substate_weight(M); // total weight for each substate.
      for (size_t i = 0; i < pdfs.size(); i++) {
        int32 j2 = pdfs[i];
        substate_weight.AddVec(state_occupancies(j2), c_[j2]);
      }
      for (int32 m = 0; m < M; m++) {
        w_jm.AddMatVec(1.0, w_, kNoTrans, v_[j1].Row(m), 0.0);
        w_jm.ApplySoftMax();
        gamma_i->AddVec(substate_weight(m), w_jm);
      }
    }
  }
  
  
  void AmSgmm2::ComputePdfMappings() {
    if (pdf2group_.empty()) {
      KALDI_WARN << "ComputePdfMappings(): no pdf2group_ map, assuming you "
          "are reading in old model.";
      KALDI_ASSERT(v_.size() != 0);
      pdf2group_.resize(v_.size());
      for (int32 j2 = 0; j2 < static_cast<int32>(pdf2group_.size()); j2++)
        pdf2group_[j2] = j2;
    }
    group2pdf_.clear();
    for (int32 j2 = 0; j2 < static_cast<int32>(pdf2group_.size()); j2++) {
      int32 j1 = pdf2group_[j2];
      if (group2pdf_.size() <= j1) group2pdf_.resize(j1+1);
      group2pdf_[j1].push_back(j2);
    }
  }
  
  void AmSgmm2::Read(std::istream &in_stream, bool binary) {
    { // We want this to work even if the object was previously
      // populated, so we clear the items that are more likely
      // to cause problems.
      pdf2group_.clear();
      group2pdf_.clear();
      u_.Resize(0,0);
      w_jmi_.clear();
      v_.clear();
    }
    // removing anything that was in the object before.
    int32 num_pdfs = -1, feat_dim, num_gauss;
    std::string token;
  
    ExpectToken(in_stream, binary, "<SGMM>");
    ExpectToken(in_stream, binary, "<NUMSTATES>");
    ReadBasicType(in_stream, binary, &num_pdfs);
    ExpectToken(in_stream, binary, "<DIMENSION>");
    ReadBasicType(in_stream, binary, &feat_dim);
    ExpectToken(in_stream, binary, "<NUMGAUSS>");
    ReadBasicType(in_stream, binary, &num_gauss);
  
    KALDI_ASSERT(num_pdfs > 0 && feat_dim > 0);
  
    ReadToken(in_stream, binary, &token);
  
    while (token != "</SGMM>") {
      if (token == "<PDF2GROUP>") {
        ReadIntegerVector(in_stream, binary, &pdf2group_);
        ComputePdfMappings();
      } else if (token == "<WEIGHTIDX2GAUSS>") {  // TEMP!   Will remove.
        std::vector<int32> garbage;
        ReadIntegerVector(in_stream, binary, &garbage);
      } else if (token == "<DIAG_UBM>") {
        diag_ubm_.Read(in_stream, binary);
      } else if (token == "<FULL_UBM>") {
        full_ubm_.Read(in_stream, binary);
      } else if (token == "<SigmaInv>") {
        SigmaInv_.resize(num_gauss);
        for (int32 i = 0; i < num_gauss; i++) {
          SigmaInv_[i].Read(in_stream, binary);
        }
      } else if (token == "<M>") {
        M_.resize(num_gauss);
        for (int32 i = 0; i < num_gauss; i++) {
          M_[i].Read(in_stream, binary);
        }
      } else if (token == "<N>") {
        N_.resize(num_gauss);
        for (int32 i = 0; i < num_gauss; i++) {
          N_[i].Read(in_stream, binary);
        }
      } else if (token == "<w>") {
        w_.Read(in_stream, binary);
      } else if (token == "<u>") {
        u_.Read(in_stream, binary);
      } else if (token == "<v>") {
        int32 num_groups = group2pdf_.size();
        if (num_groups == 0) {
          KALDI_WARN << "Reading old model with new code (should still work)";
          num_groups = num_pdfs;
        }
        v_.resize(num_groups);
        for (int32 j1 = 0; j1 < num_groups; j1++) {
          v_[j1].Read(in_stream, binary);
        }
      } else if (token == "<c>") {
        c_.resize(num_pdfs);
        for (int32 j2 = 0; j2 < num_pdfs; j2++) {
          c_[j2].Read(in_stream, binary);
        }
      } else if (token == "<n>") {
        int32 num_groups = group2pdf_.size();
        if (num_groups == 0) num_groups = num_pdfs;
        n_.resize(num_groups);
        for (int32 j1 = 0; j1 < num_groups; j1++) {
          n_[j1].Read(in_stream, binary);
        }
        // The following are the Gaussian prior parameters for MAP adaptation of M
        // They may be moved to somewhere else eventually.
      } else if (token == "<M_Prior>") {
        ExpectToken(in_stream, binary, "<NUMGaussians>");
        ReadBasicType(in_stream, binary, &num_gauss);
        M_prior_.resize(num_gauss);
        for (int32 i = 0; i < num_gauss; i++) {
          M_prior_[i].Read(in_stream, binary);
        }
      } else if (token == "<Row_Cov_Inv>") {
        row_cov_inv_.Read(in_stream, binary);
      } else if (token == "<Col_Cov_Inv>") {
        col_cov_inv_.Read(in_stream, binary);
      } else {
        KALDI_ERR << "Unexpected token '" << token << "' in model file ";
      }
      ReadToken(in_stream, binary, &token);
    }
  
    if (pdf2group_.empty())
      ComputePdfMappings(); // sets up group2pdf_, and pdf2group_ if reading
    // old model.
  
    if (n_.empty())
      ComputeNormalizers();
    if (HasSpeakerDependentWeights())
      ComputeWeights();
  }
  
  int32 AmSgmm2::Pdf2Group(int32 j2) const {
    KALDI_ASSERT(static_cast<size_t>(j2) < pdf2group_.size());
    int32 j1 = pdf2group_[j2];
    return j1;
  }
  
  
  void AmSgmm2::Write(std::ostream &out_stream,
                     bool binary,
                     SgmmWriteFlagsType write_params) const {
    int32 num_pdfs = NumPdfs(),
        feat_dim = FeatureDim(),
        num_gauss = NumGauss();
  
    WriteToken(out_stream, binary, "<SGMM>");
    if (!binary) out_stream << "\n";
    WriteToken(out_stream, binary, "<NUMSTATES>");
    WriteBasicType(out_stream, binary, num_pdfs);
    WriteToken(out_stream, binary, "<DIMENSION>");
    WriteBasicType(out_stream, binary, feat_dim);
    WriteToken(out_stream, binary, "<NUMGAUSS>");
    WriteBasicType(out_stream, binary, num_gauss);
    if (!binary) out_stream << "\n";
  
    if (write_params & kSgmmBackgroundGmms) {
      WriteToken(out_stream, binary, "<DIAG_UBM>");
      diag_ubm_.Write(out_stream, binary);
      WriteToken(out_stream, binary, "<FULL_UBM>");
      full_ubm_.Write(out_stream, binary);
    }
  
    if (write_params & kSgmmGlobalParams) {
      WriteToken(out_stream, binary, "<SigmaInv>");
      if (!binary) out_stream << "\n";
      for (int32 i = 0; i < num_gauss; i++) {
        SigmaInv_[i].Write(out_stream, binary);
      }
      WriteToken(out_stream, binary, "<M>");
      if (!binary) out_stream << "\n";
      for (int32 i = 0; i < num_gauss; i++) {
        M_[i].Write(out_stream, binary);
      }
      if (N_.size() != 0) {
        WriteToken(out_stream, binary, "<N>");
        if (!binary) out_stream << "\n";
        for (int32 i = 0; i < num_gauss; i++) {
          N_[i].Write(out_stream, binary);
        }
      }
      WriteToken(out_stream, binary, "<w>");
      w_.Write(out_stream, binary);
      WriteToken(out_stream, binary, "<u>");
      u_.Write(out_stream, binary);
    }
  
    if (write_params & kSgmmStateParams) {
      WriteToken(out_stream, binary, "<PDF2GROUP>");
      WriteIntegerVector(out_stream, binary, pdf2group_);
      WriteToken(out_stream, binary, "<v>");
      for (int32 j1 = 0; j1 < NumGroups(); j1++) {
        v_[j1].Write(out_stream, binary);
      }
      WriteToken(out_stream, binary, "<c>");
      for (int32 j2 = 0; j2 < num_pdfs; j2++) {
        c_[j2].Write(out_stream, binary);
      }
    }
  
    if (write_params & kSgmmNormalizers) {
      WriteToken(out_stream, binary, "<n>");
      if (n_.empty())
        KALDI_WARN << "Not writing normalizers since they are not present.";
      else
        for (int32 j1 = 0; j1 < NumGroups(); j1++)
          n_[j1].Write(out_stream, binary);
    }
    WriteToken(out_stream, binary, "</SGMM>");
  }
  
  
  void AmSgmm2::Check(bool show_properties) {
    int32 J1 = NumGroups(),
        J2 = NumPdfs(),
        num_gauss = NumGauss(),
        feat_dim = FeatureDim(),
        phn_dim = PhoneSpaceDim(),
        spk_dim = SpkSpaceDim();
  
    if (show_properties)
      KALDI_LOG << "AmSgmm2: #pdfs = " << J2 << ", #pdf-groups = "
                << J1 << ", #Gaussians = "
                << num_gauss << ", feature dim = " << feat_dim
                << ", phone-space dim =" << phn_dim
                << ", speaker-space dim =" << spk_dim;
    KALDI_ASSERT(J1 > 0 && num_gauss > 0 && feat_dim > 0 && phn_dim > 0
                 && J2 > 0 && J2 >= J1);
  
    std::ostringstream debug_str;
  
    // First check the diagonal-covariance UBM.
    KALDI_ASSERT(diag_ubm_.NumGauss() == num_gauss);
    KALDI_ASSERT(diag_ubm_.Dim() == feat_dim);
  
    // Check the full-covariance UBM.
    KALDI_ASSERT(full_ubm_.NumGauss() == num_gauss);
    KALDI_ASSERT(full_ubm_.Dim() == feat_dim);
  
    // Check the globally-shared covariance matrices.
    KALDI_ASSERT(SigmaInv_.size() == static_cast<size_t>(num_gauss));
    for (int32 i = 0; i < num_gauss; i++) {
      KALDI_ASSERT(SigmaInv_[i].NumRows() == feat_dim &&
                   SigmaInv_[i](0, 0) > 0.0);  // or it wouldn't be +ve definite.
    }
  
    if (spk_dim != 0) {
      KALDI_ASSERT(N_.size() == static_cast<size_t>(num_gauss));
      for (int32 i = 0; i < num_gauss; i++)
        KALDI_ASSERT(N_[i].NumRows() == feat_dim && N_[i].NumCols() == spk_dim);
      if (u_.NumRows() == 0) {
        debug_str << "Speaker-weight projections: no.";
      } else {
        KALDI_ASSERT(u_.NumRows() == num_gauss && u_.NumCols() == spk_dim);
        debug_str << "Speaker-weight projections: yes.";
      }
    } else {
      KALDI_ASSERT(N_.size() == 0 && u_.NumRows() == 0);
    }
  
    KALDI_ASSERT(M_.size() == static_cast<size_t>(num_gauss));
    for (int32 i = 0; i < num_gauss; i++) {
      KALDI_ASSERT(M_[i].NumRows() == feat_dim && M_[i].NumCols() == phn_dim);
    }
  
    KALDI_ASSERT(w_.NumRows() == num_gauss && w_.NumCols() == phn_dim);
  
    {  // check v, c.
      KALDI_ASSERT(v_.size() == static_cast<size_t>(J1) &&
                   c_.size() == static_cast<size_t>(J2));
      int32 nSubstatesTot = 0;
      for (int32 j1 = 0; j1 < J1; j1++) {
        int32 M_j = NumSubstatesForGroup(j1);
        nSubstatesTot += M_j;
        KALDI_ASSERT(M_j > 0 && v_[j1].NumRows() == M_j &&
                     v_[j1].NumCols() == phn_dim);
      }
      debug_str << "Substates: "<< (nSubstatesTot) << ".  ";
      int32 nSubstateWeights = 0;
      for (int32 j2 = 0; j2 < J2; j2++) {
        int32 j1 = Pdf2Group(j2);
        int32 M = NumSubstatesForPdf(j2);
        KALDI_ASSERT(M == NumSubstatesForGroup(j1));
        nSubstateWeights += M;
      }
      KALDI_ASSERT(nSubstateWeights >= nSubstatesTot);
      debug_str << "SubstateWeights: "<< (nSubstateWeights) << ".  ";
    }
  
    // check normalizers.
    if (n_.size() == 0) {
      debug_str << "Normalizers: no.  ";
    } else {
      debug_str << "Normalizers: yes.  ";
      KALDI_ASSERT(n_.size() == static_cast<size_t>(J1));
      for (int32 j1 = 0; j1 < J1; j1++) {
        KALDI_ASSERT(n_[j1].NumRows() == num_gauss &&
                     n_[j1].NumCols() == NumSubstatesForGroup(j1));
      }
    }
  
    // check w_jmi_.
    if (w_jmi_.size() == 0) {
      debug_str << "Computed weights: no.  ";
    } else {
      debug_str << "Computed weights: yes.  ";
      KALDI_ASSERT(w_jmi_.size() == static_cast<size_t>(J1));
      for (int32 j1 = 0; j1 < J1; j1++) {
        KALDI_ASSERT(w_jmi_[j1].NumRows() == NumSubstatesForGroup(j1) &&
                     w_jmi_[j1].NumCols() == num_gauss);
      }
    }
  
    if (show_properties)
      KALDI_LOG << "Subspace GMM model properties: " << debug_str.str();
  }
  
  void AmSgmm2::InitializeFromFullGmm(const FullGmm &full_gmm,
                                      const std::vector<int32> &pdf2group,
                                      int32 phn_subspace_dim,
                                      int32 spk_subspace_dim,
                                      bool speaker_dependent_weights,
                                      BaseFloat self_weight) {
    pdf2group_ = pdf2group;
    ComputePdfMappings();
    full_ubm_.CopyFromFullGmm(full_gmm);
    diag_ubm_.CopyFromFullGmm(full_gmm);
    if (phn_subspace_dim < 1 || phn_subspace_dim > full_gmm.Dim() + 1) {
      KALDI_WARN << "Initial phone-subspace dimension must be >= 1, value is "
                 << phn_subspace_dim << "; setting to " << full_gmm.Dim() + 1;
      phn_subspace_dim = full_gmm.Dim() + 1;
    }
    KALDI_ASSERT(spk_subspace_dim >= 0);
  
    w_.Resize(0, 0);
    N_.clear();
    c_.clear();
    v_.clear();
    SigmaInv_.clear();
  
    KALDI_LOG << "Initializing model";
    Matrix<BaseFloat> norm_xform;
    ComputeFeatureNormalizingTransform(full_gmm, &norm_xform);
    InitializeMw(phn_subspace_dim, norm_xform);
    if (spk_subspace_dim > 0)
      InitializeNu(spk_subspace_dim, norm_xform, speaker_dependent_weights);
    InitializeVecsAndSubstateWeights(self_weight);
    KALDI_LOG << "Initializing variances";
    InitializeCovars();
  }
  
  void AmSgmm2::CopyFromSgmm2(const AmSgmm2 &other,
                            bool copy_normalizers,
                            bool copy_weights) {
    KALDI_LOG << "Copying AmSgmm2";
    pdf2group_ = other.pdf2group_;
    group2pdf_ = other.group2pdf_;
  
    // Copy background GMMs
    diag_ubm_.CopyFromDiagGmm(other.diag_ubm_);
    full_ubm_.CopyFromFullGmm(other.full_ubm_);
  
    // Copy global params
    SigmaInv_ = other.SigmaInv_;
    M_ = other.M_;
    w_ = other.w_;
    N_ = other.N_;
    u_ = other.u_;
  
    // Copy state-specific params, but only copy normalizers if requested.
    v_ = other.v_;
    c_ = other.c_;
    if (copy_normalizers) n_ = other.n_;
    if (copy_weights) w_jmi_ = other.w_jmi_;
  
    KALDI_LOG << "Done.";
  }
  
  void AmSgmm2::ComputePerFrameVars(const VectorBase<BaseFloat> &data,
                                   const std::vector<int32> &gselect,
                                   const Sgmm2PerSpkDerivedVars &spk_vars,
                                   Sgmm2PerFrameDerivedVars *per_frame_vars) const {
    KALDI_ASSERT(!n_.empty() && "ComputeNormalizers() must be called.");
  
    per_frame_vars->Resize(gselect.size(), FeatureDim(), PhoneSpaceDim());
  
    per_frame_vars->gselect = gselect;
    per_frame_vars->xt.CopyFromVec(data);
  
    for (int32 ki = 0, last = gselect.size(); ki < last; ki++) {
      int32 i = gselect[ki];
      per_frame_vars->xti.Row(ki).CopyFromVec(per_frame_vars->xt);
      if (spk_vars.v_s.Dim() != 0)
        per_frame_vars->xti.Row(ki).AddVec(-1.0, spk_vars.o_s.Row(i));
    }
    Vector<BaseFloat> SigmaInv_xt(FeatureDim());
  
    bool speaker_dep_weights =
        (spk_vars.v_s.Dim() != 0 && HasSpeakerDependentWeights());
    for (int32 ki = 0, last = gselect.size(); ki < last; ki++) {
      int32 i = gselect[ki];
      BaseFloat ssgmm_term = (speaker_dep_weights ? spk_vars.log_b_is(i) : 0.0);
      SigmaInv_xt.AddSpVec(1.0, SigmaInv_[i], per_frame_vars->xti.Row(ki), 0.0);
      // Eq (35): z_{i}(t) = M_{i}^{T} \Sigma_{i}^{-1} x_{i}(t)
      per_frame_vars->zti.Row(ki).AddMatVec(1.0, M_[i], kTrans, SigmaInv_xt, 0.0);
      // Eq.(36): n_{i}(t) = -0.5 x_{i}^{T} \Sigma_{i}^{-1} x_{i}(t)
      per_frame_vars->nti(ki) = -0.5 * VecVec(per_frame_vars->xti.Row(ki),
                                              SigmaInv_xt) + ssgmm_term;
    }
  }
  
  // inline
  void AmSgmm2::ComponentLogLikes(const Sgmm2PerFrameDerivedVars &per_frame_vars,
                                 int32 j1,
                                 Sgmm2PerSpkDerivedVars *spk_vars,
                                 Matrix<BaseFloat> *loglikes) const {
    const vector<int32> &gselect = per_frame_vars.gselect;
    int32 num_gselect = gselect.size(), num_substates = v_[j1].NumRows();
  
    // Eq.(37): log p(x(t), m, i|j)  [indexed by j, ki]
    // Although the extra memory allocation of storing this as a
    // matrix might seem unnecessary, we save time in the LogSumExp()
    // via more effective pruning.
    loglikes->Resize(num_gselect, num_substates);
    bool speaker_dep_weights =
        (spk_vars->v_s.Dim() != 0 && HasSpeakerDependentWeights());
    if (speaker_dep_weights) {
      KALDI_ASSERT(static_cast<int32>(spk_vars->log_d_jms.size()) == NumGroups());
      KALDI_ASSERT(static_cast<int32>(w_jmi_.size()) == NumGroups() ||
                   "You need to call ComputeWeights().");
    }
    for (int32 ki = 0;  ki < num_gselect; ki++) {
      SubVector<BaseFloat> logp_xi(*loglikes, ki);
      int32 i = gselect[ki];
      // for all substates, compute z_{i}^T v_{jm}
      logp_xi.AddMatVec(1.0, v_[j1], kNoTrans, per_frame_vars.zti.Row(ki), 0.0);
      logp_xi.AddVec(1.0, n_[j1].Row(i));  // for all substates, add n_{jim}
      logp_xi.Add(per_frame_vars.nti(ki));  // for all substates, add n_{i}(t)
    }
    if (speaker_dep_weights) { // [SSGMM]
      Vector<BaseFloat> &log_d = spk_vars->log_d_jms[j1];
      if (log_d.Dim() == 0) { // have not yet cached this quantity.
        log_d.Resize(num_substates);
        log_d.AddMatVec(1.0, w_jmi_[j1], kNoTrans, spk_vars->b_is, 0.0);
        log_d.ApplyLog();
      }
      loglikes->AddVecToRows(-1.0, log_d); // [SSGMM] this is the term
      // - log d_{jm}^{(s)} in the likelihood function [eq. 25 in
      // the techreport]
    }
  }
  
  
  BaseFloat AmSgmm2::LogLikelihood(const Sgmm2PerFrameDerivedVars &per_frame_vars,
                                  int32 j2,
                                  Sgmm2LikelihoodCache *cache,
                                  Sgmm2PerSpkDerivedVars *spk_vars,
                                  BaseFloat log_prune) const {
    int32 t = cache->t; // not a real time; used to uniquely identify frames.
    // Forgo asserts here, as this is frequently called.
    // We'll probably get a segfault if an error is made.
    Sgmm2LikelihoodCache::PdfCacheElement &pdf_cache =
        cache->pdf_cache[j2];
  #ifdef KALDI_PARANOID
    bool random_test = (Rand() % 1000 == 1); // to check that the user is
    // calling Next() on the cache, as they should.
  #else
    bool random_test = false; // compiler will ignore test branches.
  #endif
    if (pdf_cache.t == t) {
      if (!random_test) return pdf_cache.log_like;
    } else {
      random_test = false;
    }
    // if random_test == true at this point, it was already cached, and we will
    // verify that we return the same value as the cached one.
    pdf_cache.t = t;
  
    int32 j1 = pdf2group_[j2];
    Sgmm2LikelihoodCache::SubstateCacheElement &substate_cache =
        cache->substate_cache[j1];
    if (substate_cache.t != t) { // Need to compute sub-state likelihoods.
      substate_cache.t = t;
      Matrix<BaseFloat> loglikes; // indexed [gselect-index][substate-index]
      ComponentLogLikes(per_frame_vars, j1, spk_vars, &loglikes);
      BaseFloat max = loglikes.Max(); // use this to keep things in good numerical range.
      loglikes.Add(-max);
      loglikes.ApplyExp();
      substate_cache.remaining_log_like = max;
      int32 num_substates = loglikes.NumCols();
      substate_cache.likes.Resize(num_substates); // zeroes it.
      substate_cache.likes.AddRowSumMat(1.0, loglikes); // add likelihoods [not in log!] for
      // each column [i.e. summing over the rows], so we get the sum for
      // each substate index.  You have to multiply by exp(remaining_log_like)
      // to get a real likelihood.
    }
  
    BaseFloat log_like = substate_cache.remaining_log_like
        + Log(VecVec(substate_cache.likes, c_[j2]));
  
    if (random_test)
      KALDI_ASSERT(ApproxEqual(pdf_cache.log_like, log_like));
  
    pdf_cache.log_like = log_like;
    KALDI_ASSERT(log_like == log_like && log_like - log_like == 0); // check
    // that it's not NaN or infinity.
    return log_like;
  }
  
  BaseFloat
  AmSgmm2::ComponentPosteriors(const Sgmm2PerFrameDerivedVars &per_frame_vars,
                              int32 j2,
                              Sgmm2PerSpkDerivedVars *spk_vars,
                              Matrix<BaseFloat> *post) const {
    KALDI_ASSERT(j2 < NumPdfs() && post != NULL);
    int32 j1 = pdf2group_[j2];
    ComponentLogLikes(per_frame_vars, j1, spk_vars, post); // now
    // post is a matrix of log-likelihoods indexed by [gaussian-selection index]
    // [sub-state index].  It doesn't include the sub-state weights,
    // though.
    BaseFloat loglike = post->Max();
    post->Add(-loglike); // get it to nicer numeric range.
    post->ApplyExp(); // so we're dealing with likelihoods (with an arbitrary offset
    // "loglike" removed to make it in a nice numeric range)
    post->MulColsVec(c_[j2]); // include the sub-state weights.
  
    BaseFloat tot_like = post->Sum();
    KALDI_ASSERT(tot_like != 0.0); // note: not valid to have zero weights.
    loglike += Log(tot_like);
    post->Scale(1.0 / tot_like); // so "post" now sums to one, and "loglike"
    // contains the correct log-likelihood of the data given the pdf.
  
    return loglike;
  }
  
  void AmSgmm2::SplitSubstatesInGroup(const Vector<BaseFloat> &pdf_occupancies,
                                      const Sgmm2SplitSubstatesConfig &opts,
                                      const SpMatrix<BaseFloat> &sqrt_H_sm,
                                      int32 j1,
                                      int32 tgt_M) {
    const std::vector<int32> &pdfs = group2pdf_[j1];
    int32 phn_dim = PhoneSpaceDim(), cur_M = NumSubstatesForGroup(j1),
        num_pdfs_for_group = pdfs.size();
    Vector<BaseFloat> rand_vec(phn_dim), v_shift(phn_dim);
  
    KALDI_ASSERT(tgt_M >= cur_M);
    if (cur_M == tgt_M) return;
    // Resize v[j1] to fit new substates
    {
      Matrix<BaseFloat> tmp_v_j(v_[j1]);
      v_[j1].Resize(tgt_M, phn_dim);
      v_[j1].Range(0, cur_M, 0, phn_dim).CopyFromMat(tmp_v_j);
    }
  
    // we'll use a temporary matrix for the c quantities.
    Matrix<BaseFloat> c_j(num_pdfs_for_group, tgt_M);
    for (int32 i = 0; i < num_pdfs_for_group; i++) {
      int32 j2 = pdfs[i];
      c_j.Row(i).Range(0, cur_M).CopyFromVec(c_[j2]);
    }
  
    // Keep splitting substates until obtaining the desired number
    for (; cur_M < tgt_M; cur_M++) {
      int32 split_m; // substate to split.
      {
        Vector<BaseFloat> substate_count(tgt_M);
        substate_count.AddRowSumMat(1.0, c_j);
        BaseFloat *data = substate_count.Data();
        split_m = std::max_element(data, data+cur_M) - data;
      }
      for (int32 i = 0; i < num_pdfs_for_group; i++) { // divide count of split
        // substate. [extended for SCTM]
        // c_{jkm} := c_{jmk}' := c_{jkm} / 2
        c_j(i, split_m) = c_j(i, cur_M) = c_j(i, split_m) / 2;
      }
      // v_{jkm} := +/- split_perturb * H_k^{(sm)}^{-0.5} * rand_vec
      std::generate(rand_vec.Data(), rand_vec.Data() + rand_vec.Dim(),
                    _RandGauss);
      v_shift.AddSpVec(opts.perturb_factor, sqrt_H_sm, rand_vec, 0.0);
      v_[j1].Row(cur_M).CopyFromVec(v_[j1].Row(split_m));
      v_[j1].Row(cur_M).AddVec(1.0, v_shift);
      v_[j1].Row(split_m).AddVec(-1.0, v_shift);
    }
    // copy the temporary matrix for the c_ (sub-state weight)
    // quantities back to the place it belongs.
    for (int32 i = 0; i < num_pdfs_for_group; i++) {
      int32 j2 = pdfs[i];
      c_[j2].Resize(tgt_M);
      c_[j2].CopyFromVec(c_j.Row(i));
    }
  }
  
  
  void AmSgmm2::SplitSubstates(const Vector<BaseFloat> &pdf_occupancies,
                               const Sgmm2SplitSubstatesConfig &opts) {
    KALDI_ASSERT(pdf_occupancies.Dim() == NumPdfs());
    int32 J1 = NumGroups(), J2 = NumPdfs();
    Vector<BaseFloat> group_occupancies(J1);
    for (int32 j2 = 0; j2 < J2; j2++)
      group_occupancies(Pdf2Group(j2)) += pdf_occupancies(j2);
  
    vector<int32> tgt_num_substates;
  
    GetSplitTargets(group_occupancies, opts.split_substates,
                    opts.power, opts.min_count, &tgt_num_substates);
  
    int32 tot_num_substates_old = 0, tot_num_substates_new = 0;
    vector< SpMatrix<BaseFloat> > H_i;
    SpMatrix<BaseFloat> sqrt_H_sm;
  
    ComputeH(&H_i);  // set up that array.
    ComputeHsmFromModel(H_i, pdf_occupancies, &sqrt_H_sm, opts.max_cond);
    H_i.clear();
    sqrt_H_sm.ApplyPow(-0.5);
  
    for (int32 j1 = 0; j1 < J1; j1++) {
      int32 cur_M = NumSubstatesForGroup(j1),
          tgt_M = tgt_num_substates[j1];
      tot_num_substates_old += cur_M;
      tot_num_substates_new += std::max(cur_M, tgt_M);
      if (cur_M < tgt_M)
        SplitSubstatesInGroup(pdf_occupancies, opts, sqrt_H_sm, j1, tgt_M);
    }
    if (tot_num_substates_old == tot_num_substates_new) {
      KALDI_LOG << "Not splitting substates; current #substates is "
                << tot_num_substates_old << " and target is "
                << opts.split_substates;
    } else {
      KALDI_LOG << "Getting rid of normalizers as they will no longer be valid";
      n_.clear();
      KALDI_LOG << "Split " << tot_num_substates_old << " substates to "
                << tot_num_substates_new;
    }
  }
  
  void AmSgmm2::IncreasePhoneSpaceDim(int32 target_dim,
                                     const Matrix<BaseFloat> &norm_xform) {
    KALDI_ASSERT(!M_.empty());
    int32 initial_dim = PhoneSpaceDim(),
        feat_dim = FeatureDim();
    KALDI_ASSERT(norm_xform.NumRows() == feat_dim);
  
    if (target_dim < initial_dim)
      KALDI_ERR << "You asked to increase phn dim to a value lower than the "
                << " current dimension, " << target_dim << " < " << initial_dim;
  
    if (target_dim > initial_dim + feat_dim) {
      KALDI_WARN << "Cannot increase phone subspace dimensionality from "
                 << initial_dim << " to " << target_dim << ", increasing to "
                 << initial_dim + feat_dim;
      target_dim = initial_dim + feat_dim;
    }
  
    if (initial_dim < target_dim) {
      Matrix<BaseFloat> tmp_M(feat_dim, initial_dim);
      for (int32 i = 0; i < NumGauss(); i++) {
        tmp_M.CopyFromMat(M_[i]);
        M_[i].Resize(feat_dim, target_dim);
        M_[i].Range(0, feat_dim, 0, tmp_M.NumCols()).CopyFromMat(tmp_M);
        M_[i].Range(0, feat_dim, tmp_M.NumCols(),
            target_dim - tmp_M.NumCols()).CopyFromMat(norm_xform.Range(0,
                feat_dim, 0, target_dim-tmp_M.NumCols()));
      }
      Matrix<BaseFloat> tmp_w = w_;
      w_.Resize(tmp_w.NumRows(), target_dim);
      w_.Range(0, tmp_w.NumRows(), 0, tmp_w.NumCols()).CopyFromMat(tmp_w);
  
      for (int32 j1 = 0; j1 < NumGroups(); j1++) {
        // Resize phonetic-subspce vectors.
        Matrix<BaseFloat> tmp_v_j = v_[j1];
        v_[j1].Resize(tmp_v_j.NumRows(), target_dim);
        v_[j1].Range(0, tmp_v_j.NumRows(), 0, tmp_v_j.NumCols()).CopyFromMat(
            tmp_v_j);
      }
      KALDI_LOG << "Phone subspace dimensionality increased from " <<
          initial_dim << " to " << target_dim;
    } else {
      KALDI_LOG << "Phone subspace dimensionality unchanged, since target " <<
          "dimension (" << target_dim << ") <= initial dimansion (" <<
          initial_dim << ")";
    }
  }
  
  void AmSgmm2::IncreaseSpkSpaceDim(int32 target_dim,
                                   const Matrix<BaseFloat> &norm_xform,
                                   bool speaker_dependent_weights) {
    int32 initial_dim = SpkSpaceDim(),
        feat_dim = FeatureDim();
    KALDI_ASSERT(norm_xform.NumRows() == feat_dim);
  
    if (N_.size() == 0)
      N_.resize(NumGauss());
  
    if (target_dim < initial_dim)
      KALDI_ERR << "You asked to increase spk dim to a value lower than the "
                << " current dimension, " << target_dim << " < " << initial_dim;
  
    if (target_dim > initial_dim + feat_dim) {
      KALDI_WARN << "Cannot increase speaker subspace dimensionality from "
                 << initial_dim << " to " << target_dim << ", increasing to "
                 << initial_dim + feat_dim;
      target_dim = initial_dim + feat_dim;
    }
  
    if (initial_dim < target_dim) {
      int32 dim_change = target_dim - initial_dim;
      Matrix<BaseFloat> tmp_N((initial_dim != 0) ? feat_dim : 0,
                              initial_dim);
      for (int32 i = 0; i < NumGauss(); i++) {
        if (initial_dim != 0) tmp_N.CopyFromMat(N_[i]);
        N_[i].Resize(feat_dim, target_dim);
        if (initial_dim != 0) {
          N_[i].Range(0, feat_dim, 0, tmp_N.NumCols()).CopyFromMat(tmp_N);
        }
        N_[i].Range(0, feat_dim, tmp_N.NumCols(), dim_change).CopyFromMat(
            norm_xform.Range(0, feat_dim, 0, dim_change));
      }
      // if we already have speaker-dependent weights or we are increasing
      // spk-dim from zero and are asked to add them...
      if (u_.NumRows() != 0 || (initial_dim == 0 && speaker_dependent_weights))
        u_.Resize(NumGauss(), target_dim, kCopyData); // extend dim of u_i's
      KALDI_LOG << "Speaker subspace dimensionality increased from " <<
          initial_dim << " to " << target_dim;
      if (initial_dim == 0 && speaker_dependent_weights)
        KALDI_LOG << "Added parameters u for speaker-dependent weights.";
    } else {
      KALDI_LOG << "Speaker subspace dimensionality unchanged, since target " <<
          "dimension (" << target_dim << ") <= initial dimansion (" <<
          initial_dim << ")";
    }
  }
  
  void AmSgmm2::ComputeWeights() {
    int32 J1 = NumGroups();
    w_jmi_.resize(J1);
    int32 i = NumGauss();
    for (int32 j1 = 0; j1 < J1; j1++) {
      int32 M = NumSubstatesForGroup(j1);
      w_jmi_[j1].Resize(M, i);
      w_jmi_[j1].AddMatMat(1.0, v_[j1], kNoTrans, w_, kTrans, 0.0);
      // now w_jmi_ contains un-normalized log weights.
      for (int32 m = 0; m < M; m++)
        w_jmi_[j1].Row(m).ApplySoftMax(); // get the actual weights.
    }
  }
  
  void AmSgmm2::ComputeDerivedVars() {
    if (n_.empty()) ComputeNormalizers();
    if (diag_ubm_.NumGauss() != full_ubm_.NumGauss()
        || diag_ubm_.Dim() != full_ubm_.Dim()) {
      diag_ubm_.CopyFromFullGmm(full_ubm_);
    }
    if (w_jmi_.empty() && HasSpeakerDependentWeights())
      ComputeWeights();
  }
  
  class ComputeNormalizersClass: public MultiThreadable { // For multi-threaded.
   public:
    ComputeNormalizersClass(AmSgmm2 *am_sgmm,
                            int32 *entropy_count_ptr,
                            double *entropy_sum_ptr):
        am_sgmm_(am_sgmm), entropy_count_ptr_(entropy_count_ptr),
        entropy_sum_ptr_(entropy_sum_ptr), entropy_count_(0),
        entropy_sum_(0.0) { }
  
    ComputeNormalizersClass(const ComputeNormalizersClass &other):
        MultiThreadable(other),
        am_sgmm_(other.am_sgmm_), entropy_count_ptr_(other.entropy_count_ptr_),
        entropy_sum_ptr_(other.entropy_sum_ptr_), entropy_count_(0),
        entropy_sum_(0.0) { }
  
    ~ComputeNormalizersClass() {
      *entropy_count_ptr_ += entropy_count_;
      *entropy_sum_ptr_ += entropy_sum_;
    }
  
    inline void operator() () {
      // Note: give them local copy of the sums we're computing,
      // which will be propagated to original pointer in the destructor.
      am_sgmm_->ComputeNormalizersInternal(num_threads_, thread_id_,
                                           &entropy_count_,
                                           &entropy_sum_);
    }
   private:
    ComputeNormalizersClass() { } // Disallow empty constructor.
    AmSgmm2 *am_sgmm_;
    int32 *entropy_count_ptr_;
    double *entropy_sum_ptr_;
    int32 entropy_count_;
    double entropy_sum_;
  
  };
  
  void AmSgmm2::ComputeNormalizers() {
    KALDI_LOG << "Computing normalizers";
    n_.resize(NumPdfs());
    int32 entropy_count = 0;
    double entropy_sum = 0.0;
    ComputeNormalizersClass c(this, &entropy_count, &entropy_sum);
    RunMultiThreaded(c);
  
    KALDI_LOG << "Entropy of weights in substates is "
              << (entropy_sum / entropy_count) << " over " << entropy_count
              << " substates, equivalent to perplexity of "
              << (Exp(entropy_sum /entropy_count));
    KALDI_LOG << "Done computing normalizers";
  }
  
  
  void AmSgmm2::ComputeNormalizersInternal(int32 num_threads, int32 thread,
                                           int32 *entropy_count,
                                           double *entropy_sum) {
  
    BaseFloat DLog2pi = FeatureDim() * Log(2 * M_PI);
    Vector<BaseFloat> log_det_Sigma(NumGauss());
  
    for (int32 i = 0; i < NumGauss(); i++) {
      try {
        log_det_Sigma(i) = - SigmaInv_[i].LogPosDefDet();
      } catch(...) {
        if (thread == 0) // just for one thread, print errors [else, duplicates]
          KALDI_WARN << "Covariance is not positive definite, setting to unit";
        SigmaInv_[i].SetUnit();
        log_det_Sigma(i) = 0.0;
      }
    }
  
    int32 J1 = NumGroups();
  
    int block_size = (NumPdfs() + num_threads-1) / num_threads;
    int j_start = thread * block_size, j_end = std::min(J1, j_start + block_size);
  
    int32 I = NumGauss();
    for (int32 j1 = j_start; j1 < j_end; j1++) {
      int32 M = NumSubstatesForGroup(j1);
      Matrix<BaseFloat> log_w_jm(M, I);
      n_[j1].Resize(I, M);
      Matrix<BaseFloat> mu_jmi(M, FeatureDim());
      Matrix<BaseFloat> SigmaInv_mu(M, FeatureDim());
  
      // (in logs): w_jm = softmax([w_{k1}^T ... w_{kD}^T] * v_{jkm}) eq.(7)
      log_w_jm.AddMatMat(1.0, v_[j1], kNoTrans, w_, kTrans, 0.0);
      for (int32 m = 0; m < M; m++) {
        log_w_jm.Row(m).Add(-1.0 * log_w_jm.Row(m).LogSumExp());
        {  // DIAGNOSTIC CODE
          (*entropy_count)++;
          for (int32 i = 0; i < NumGauss(); i++) {
            (*entropy_sum) -= log_w_jm(m, i) * Exp(log_w_jm(m, i));
          }
        }
      }
  
      for (int32 i = 0; i < I; i++) {
        // mu_jmi = M_{i} * v_{jm}
        mu_jmi.AddMatMat(1.0, v_[j1], kNoTrans, M_[i], kTrans, 0.0);
        SigmaInv_mu.AddMatSp(1.0, mu_jmi, kNoTrans, SigmaInv_[i], 0.0);
  
        for (int32 m = 0; m < M; m++) {
          // mu_{jmi} * \Sigma_{i}^{-1} * mu_{jmi}
          BaseFloat mu_SigmaInv_mu = VecVec(mu_jmi.Row(m), SigmaInv_mu.Row(m));
          // Previously had:
          // BaseFloat logc = log(c_[j](m));
          // but because of STCM aspect, we can't include the sub-state mixture weights
          // at this point [included later on.]
  
          // eq.(31)
          n_[j1](i, m) = log_w_jm(m, i) - 0.5 * (log_det_Sigma(i) + DLog2pi
              + mu_SigmaInv_mu);
          {  // Mainly diagnostic code.  Not necessary.
            BaseFloat tmp = n_[j1](i, m);
            if (!KALDI_ISFINITE(tmp)) {  // NaN or inf
              KALDI_LOG << "Warning: normalizer for j1 = " << j1 << ", m = " << m
                        << ", i = " << i << " is infinite or NaN " << tmp << "= "
                        << log_w_jm(m, i) << "+"
                        << (-0.5 * log_det_Sigma(i)) << "+" << (-0.5 * DLog2pi)
                        << "+" << (mu_SigmaInv_mu) << ", setting to finite.";
              n_[j1](i, m) = -1.0e+40;  // future work(arnab): get rid of magic number
            }
          }
        }
      }
    }
  }
  
  BaseFloat AmSgmm2::GetDjms(int32 j1, int32 m,
                            Sgmm2PerSpkDerivedVars *spk_vars) const {
    // This relates to SSGMMs (speaker-dependent weights).
    if (spk_vars->log_d_jms.empty()) return -1; // this would be
    // because we don't have speaker-dependent weights ("u" not set up).
  
    KALDI_ASSERT(!w_jmi_.empty() && "You need to call ComputeWeights() on SGMM.");
    Vector<BaseFloat> &log_d = spk_vars->log_d_jms[j1];
    if (log_d.Dim() == 0) {
      log_d.Resize(NumSubstatesForGroup(j1));
      log_d.AddMatVec(1.0, w_jmi_[j1], kNoTrans, spk_vars->b_is, 0.0);
      log_d.ApplyLog();
    }
    return Exp(log_d(m));
  }
  
  
  void AmSgmm2::ComputeFmllrPreXform(const Vector<BaseFloat> &state_occs,
                                    Matrix<BaseFloat> *xform,
                                     Matrix<BaseFloat> *inv_xform,
                                    Vector<BaseFloat> *diag_mean_scatter) const {
    int32 num_pdfs = NumPdfs(),
        num_gauss = NumGauss(),
        dim = FeatureDim();
    KALDI_ASSERT(state_occs.Dim() == num_pdfs);
  
    BaseFloat total_occ = state_occs.Sum();
  
    // Degenerate case: unlikely to ever happen.
    if (total_occ == 0) {
      KALDI_WARN << "Zero probability (computing transform). Using unit "
                 << "pre-transform";
      xform->Resize(dim, dim + 1, kUndefined);
      xform->SetUnit();
      inv_xform->Resize(dim, dim + 1, kUndefined);
      inv_xform->SetUnit();
      diag_mean_scatter->Resize(dim, kSetZero);
      return;
    }
  
    // Convert state occupancies to posteriors; Eq. (B.1)
    Vector<BaseFloat> state_posteriors(state_occs);
    state_posteriors.Scale(1/total_occ);
  
    Vector<BaseFloat> mu_jmi(dim), global_mean(dim);
    SpMatrix<BaseFloat> within_class_covar(dim), between_class_covar(dim);
    Vector<BaseFloat> gauss_weight(num_gauss);  // weights for within-class vars.
    Vector<BaseFloat> w_jm(num_gauss);
    for (int32 j1 = 0; j1 < NumGroups(); j1++) {
      const std::vector<int32> &pdfs = group2pdf_[j1];
      int32 M = NumSubstatesForGroup(j1);
      Vector<BaseFloat> substate_weight(M); // total weight for each substate.
      for (size_t i = 0; i < pdfs.size(); i++) {
        int32 j2 = pdfs[i];
        substate_weight.AddVec(state_posteriors(j2), c_[j2]);
      }
      for (int32 m = 0; m < M; m++) {
        BaseFloat this_substate_weight = substate_weight(m);
        // Eq. (7): w_jm = softmax([w_{1}^T ... w_{D}^T] * v_{jm})
        w_jm.AddMatVec(1.0, w_, kNoTrans, v_[j1].Row(m), 0.0);
        w_jm.ApplySoftMax();
  
        for (int32 i = 0; i < num_gauss; i++) {
          BaseFloat weight = this_substate_weight * w_jm(i);
          mu_jmi.AddMatVec(1.0, M_[i], kNoTrans, v_[j1].Row(m), 0.0);  // Eq. (6)
          // Eq. (B.3): \mu_avg = \sum_{jmi} p(j) c_{jm} w_{jmi} \mu_{jmi}
          global_mean.AddVec(weight, mu_jmi);
          // \Sigma_B = \sum_{jmi} p(j) c_{jm} w_{jmi} \mu_{jmi} \mu_{jmi}^T
          between_class_covar.AddVec2(weight, mu_jmi);  // Eq. (B.4)
          gauss_weight(i) += weight;
        }
      }
    }
    between_class_covar.AddVec2(-1.0, global_mean);  // Eq. (B.4)
  
    for (int32 i = 0; i < num_gauss; i++) {
      SpMatrix<BaseFloat> Sigma(SigmaInv_[i]);
      Sigma.InvertDouble();
      // Eq. (B.2): \Sigma_W = \sum_{jmi} p(j) c_{jm} w_{jmi} \Sigma_i
      within_class_covar.AddSp(gauss_weight(i), Sigma);
    }
  
    TpMatrix<BaseFloat> tmpL(dim);
    Matrix<BaseFloat> tmpLInvFull(dim, dim);
    tmpL.Cholesky(within_class_covar);  // \Sigma_W = L L^T
    tmpL.InvertDouble();  // L^{-1}
    tmpLInvFull.CopyFromTp(tmpL);  // get as full matrix.
  
    // B := L^{-1} * \Sigma_B * L^{-T}
    SpMatrix<BaseFloat> tmpB(dim);
    tmpB.AddMat2Sp(1.0, tmpLInvFull, kNoTrans, between_class_covar, 0.0);
  
    Matrix<BaseFloat> U(dim, dim);
    diag_mean_scatter->Resize(dim);
    xform->Resize(dim, dim + 1);
    inv_xform->Resize(dim, dim + 1);
  
    tmpB.Eig(diag_mean_scatter, &U);  // Eq. (B.5): B = U D V^T
  
    int32 n;
    diag_mean_scatter->ApplyFloor(1.0e-04, &n);
    if (n != 0)
      KALDI_WARN << "Floored " << n << " elements of the mean-scatter matrix.";
  
    // Eq. (B.6): A_{pre} = U^T * L^{-1}
    SubMatrix<BaseFloat> Apre(*xform, 0, dim, 0, dim);
    Apre.AddMatMat(1.0, U, kTrans, tmpLInvFull, kNoTrans, 0.0);
  
  #ifdef KALDI_PARANOID
    {
      SpMatrix<BaseFloat> tmp(dim);
      tmp.AddMat2Sp(1.0, Apre, kNoTrans, within_class_covar, 0.0);
      KALDI_ASSERT(tmp.IsUnit(0.01));
    }
    {
      SpMatrix<BaseFloat> tmp(dim);
      tmp.AddMat2Sp(1.0, Apre, kNoTrans, between_class_covar, 0.0);
      KALDI_ASSERT(tmp.IsDiagonal(0.01));
    }
  #endif
  
    // Eq. (B.7): b_{pre} = - A_{pre} \mu_{avg}
    Vector<BaseFloat> b_pre(dim);
    b_pre.AddMatVec(-1.0, Apre, kNoTrans, global_mean, 0.0);
    for (int32 r = 0; r < dim; r++) {
      xform->Row(r)(dim) = b_pre(r);  // W_{pre} = [ A_{pre}, b_{pre} ]
    }
  
    // Eq. (B.8) & (B.9): W_{inv} = [ A_{pre}^{-1}, \mu_{avg} ]
    inv_xform->CopyFromMat(*xform);
    inv_xform->Range(0, dim, 0, dim).InvertDouble();
    for (int32 r = 0; r < dim; r++)
      inv_xform->Row(r)(dim) = global_mean(r);
  }  // End of ComputePreXform()
  
  template<typename Real>
  void AmSgmm2::GetNtransSigmaInv(vector< Matrix<Real> > *out) const {
    KALDI_ASSERT(SpkSpaceDim() > 0 &&
        "Cannot compute N^{T} \\Sigma_{i}^{-1} without speaker projections.");
    out->resize(NumGauss());
    Matrix<Real> tmpcov(FeatureDim(), FeatureDim());
    Matrix<Real> tmp_n(FeatureDim(), SpkSpaceDim());
    for (int32 i = 0; i < NumGauss(); i++) {
      tmpcov.CopyFromSp(SigmaInv_[i]);
      tmp_n.CopyFromMat(N_[i]);
      (*out)[i].Resize(SpkSpaceDim(), FeatureDim());
      (*out)[i].AddMatMat(1.0, tmp_n, kTrans, tmpcov, kNoTrans, 0.0);
    }
  }
  
  // Instantiate the above template.
  template
  void AmSgmm2::GetNtransSigmaInv(vector< Matrix<float> > *out) const;
  template
  void AmSgmm2::GetNtransSigmaInv(vector< Matrix<double> > *out) const;
  
  ///////////////////////////////////////////////////////////////////////////////
  
  template<class Real>
  void AmSgmm2::ComputeH(std::vector< SpMatrix<Real> > *H_i) const {
    KALDI_ASSERT(NumGauss() != 0);
    (*H_i).resize(NumGauss());
    SpMatrix<BaseFloat> H_i_tmp(PhoneSpaceDim());
    for (int32 i = 0; i < NumGauss(); i++) {
      (*H_i)[i].Resize(PhoneSpaceDim());
      H_i_tmp.AddMat2Sp(1.0, M_[i], kTrans, SigmaInv_[i], 0.0);
      (*H_i)[i].CopyFromSp(H_i_tmp);
    }
  }
  
  // Instantiate the template.
  template
  void AmSgmm2::ComputeH(std::vector< SpMatrix<float> > *H_i) const;
  template
  void AmSgmm2::ComputeH(std::vector< SpMatrix<double> > *H_i) const;
  
  
  // Initializes the matrices M_{i} and w_i
  void AmSgmm2::InitializeMw(int32 phn_subspace_dim,
                             const Matrix<BaseFloat> &norm_xform) {
    int32 ddim = full_ubm_.Dim();
    KALDI_ASSERT(phn_subspace_dim <= ddim + 1);
    KALDI_ASSERT(phn_subspace_dim <= norm_xform.NumCols() + 1);
    KALDI_ASSERT(ddim <= norm_xform.NumRows());
  
    Vector<BaseFloat> mean(ddim);
    int32 num_gauss = full_ubm_.NumGauss();
    w_.Resize(num_gauss, phn_subspace_dim);
    M_.resize(num_gauss);
    for (int32 i = 0; i < num_gauss; i++) {
      full_ubm_.GetComponentMean(i, &mean);
      Matrix<BaseFloat> &thisM(M_[i]);
      thisM.Resize(ddim, phn_subspace_dim);
      // Eq. (27): M_{i} = [ \bar{\mu}_{i} (J)_{1:D, 1:(S-1)}]
      thisM.CopyColFromVec(mean, 0);
      int32 nonrandom_dim = std::min(phn_subspace_dim - 1, ddim),
          random_dim = phn_subspace_dim - 1 - nonrandom_dim;
      thisM.Range(0, ddim, 1, nonrandom_dim).CopyFromMat(
          norm_xform.Range(0, ddim, 0, nonrandom_dim), kNoTrans);
      // The following extension to the original paper allows us to
      // initialize the model with a larger dimension of phone-subspace vector.
      if (random_dim > 0)
        thisM.Range(0, ddim, nonrandom_dim + 1, random_dim).SetRandn();
    }
  }
  
  // Initializes the matrices N_i, and [if speaker_dependent_weights==true] u_i.
  void AmSgmm2::InitializeNu(int32 spk_subspace_dim,
                            const Matrix<BaseFloat> &norm_xform,
                            bool speaker_dependent_weights) {
    int32 ddim = full_ubm_.Dim();
  
    int32 num_gauss = full_ubm_.NumGauss();
    N_.resize(num_gauss);
    for (int32 i = 0; i < num_gauss; i++) {
      N_[i].Resize(ddim, spk_subspace_dim);
      // Eq. (28): N_{i} = [ (J)_{1:D, 1:T)}]
  
      int32 nonrandom_dim = std::min(spk_subspace_dim, ddim),
          random_dim = spk_subspace_dim - nonrandom_dim;
  
      N_[i].Range(0, ddim, 0, nonrandom_dim).
          CopyFromMat(norm_xform.Range(0, ddim, 0, nonrandom_dim), kNoTrans);
      // The following extension to the original paper allows us to
      // initialize the model with a larger dimension of speaker-subspace vector.
      if (random_dim > 0)
        N_[i].Range(0, ddim, nonrandom_dim, random_dim).SetRandn();
    }
    if (speaker_dependent_weights) {
      u_.Resize(num_gauss, spk_subspace_dim); // will set to zero.
    } else {
      u_.Resize(0, 0);
    }
  }
  
  void AmSgmm2::CopyGlobalsInitVecs(const AmSgmm2 &other,
                                    const std::vector<int32> &pdf2group,
                                    BaseFloat self_weight) {
    KALDI_LOG << "Initializing model";
    pdf2group_ = pdf2group;
    ComputePdfMappings();
  
    // Copy background GMMs
    diag_ubm_.CopyFromDiagGmm(other.diag_ubm_);
    full_ubm_.CopyFromFullGmm(other.full_ubm_);
  
    // Copy global params
    SigmaInv_ = other.SigmaInv_;
  
    M_ = other.M_;
    w_ = other.w_;
    u_ = other.u_;
    N_ = other.N_;
  
    InitializeVecsAndSubstateWeights(self_weight);
  }
  
  
  // Initializes the vectors v_{j1,m} and substate weights c_{j2,m}.
  void AmSgmm2::InitializeVecsAndSubstateWeights(BaseFloat self_weight) {
    int32 J1 = NumGroups(), J2 = NumPdfs();
    KALDI_ASSERT(J1 > 0 && J2 >= J1);
    int32 phn_subspace_dim = PhoneSpaceDim();
    KALDI_ASSERT(phn_subspace_dim > 0 && "Initialize M and w first.");
  
    v_.resize(J1);
    if (self_weight == 1.0) {
      for (int32 j1 = 0; j1 < J1; j1++) {
        v_[j1].Resize(1, phn_subspace_dim);
        v_[j1](0, 0) = 1.0;  // Eq. (26): v_{j1} = [1 0 0 ... 0]
      }
      c_.resize(J2);
      for (int32 j2 = 0; j2 < J2; j2++) {
        c_[j2].Resize(1);
        c_[j2](0) = 1.0;    // Eq. (25): c_{j1} = 1.0
      }
    } else {
      for (int32 j1 = 0; j1 < J1; j1++) {
        int32 npdfs = group2pdf_[j1].size();
        v_[j1].Resize(npdfs, phn_subspace_dim);
        for (int32 m = 0; m < npdfs; m++)
          v_[j1](m, 0) = 1.0;  // Eq. (26): v_{j1} = [1 0 0 ... 0]
      }
      c_.resize(J2);
      for (int32 j2 = 0; j2 < J2; j2++) {
        int32 j1 = pdf2group_[j2], npdfs = group2pdf_[j1].size();
        c_[j2].Resize(npdfs);
        if (npdfs == 1) c_[j2].Set(1.0);
        else {
          // note: just avoid NaNs if npdfs-1... value won't matter.
          double other_weight = (1.0 - self_weight) / std::max((1-npdfs), 1);
          c_[j2].Set(other_weight);
          for (int32 k = 0; k < npdfs; k++)
            if(group2pdf_[j1][k] == j2) c_[j2](k) = self_weight;
        }
      }
    }
  }
  
  // Initializes the within-class vars Sigma_{ki}
  void AmSgmm2::InitializeCovars() {
    std::vector< SpMatrix<BaseFloat> > &inv_covars(full_ubm_.inv_covars());
    int32 num_gauss = full_ubm_.NumGauss();
    int32 dim = full_ubm_.Dim();
    SigmaInv_.resize(num_gauss);
    for (int32 i = 0; i < num_gauss; i++) {
      SigmaInv_[i].Resize(dim);
      SigmaInv_[i].CopyFromSp(inv_covars[i]);
    }
  }
  
  // Compute the "smoothing" matrix H^{(sm)} from expected counts given the model.
  void AmSgmm2::ComputeHsmFromModel(
      const std::vector< SpMatrix<BaseFloat> > &H,
      const Vector<BaseFloat> &state_occupancies,
      SpMatrix<BaseFloat> *H_sm,
      BaseFloat max_cond) const {
    int32 num_gauss = NumGauss();
    BaseFloat tot_sum = 0.0;
    KALDI_ASSERT(state_occupancies.Dim() == NumPdfs());
    Vector<BaseFloat> w_jm(num_gauss);
    H_sm->Resize(PhoneSpaceDim());
    H_sm->SetZero();
    Vector<BaseFloat> gamma_i;
    ComputeGammaI(state_occupancies, &gamma_i);
  
    BaseFloat sum = 0.0;
    for (int32 i = 0; i < num_gauss; i++) {
      if (gamma_i(i) > 0) {
        H_sm->AddSp(gamma_i(i), H[i]);
        sum += gamma_i(i);
      }
    }
    if (sum == 0.0) {
      KALDI_WARN << "Sum of counts is zero. ";
      // set to unit matrix--arbitrary non-singular matrix.. won't ever matter.
      H_sm->SetUnit();
    } else {
      H_sm->Scale(1.0 / sum);
      int32 tmp = H_sm->LimitCondDouble(max_cond);
      if (tmp > 0) {
        KALDI_WARN << "Limited " << (tmp) << " eigenvalues of H_sm";
      }
    }
    tot_sum += sum;
  
    KALDI_LOG << "total count is " << tot_sum;
  }
  
  void ComputeFeatureNormalizingTransform(const FullGmm &gmm, Matrix<BaseFloat> *xform) {
    int32 dim = gmm.Dim();
    int32 num_gauss = gmm.NumGauss();
    SpMatrix<BaseFloat> within_class_covar(dim);
    SpMatrix<BaseFloat> between_class_covar(dim);
    Vector<BaseFloat> global_mean(dim);
  
    // Accumulate LDA statistics from the GMM parameters.
    {
      BaseFloat total_weight = 0.0;
      Vector<BaseFloat> tmp_weight(num_gauss);
      Matrix<BaseFloat> tmp_means;
      std::vector< SpMatrix<BaseFloat> > tmp_covars;
      tmp_weight.CopyFromVec(gmm.weights());
      gmm.GetCovarsAndMeans(&tmp_covars, &tmp_means);
      for (int32 i = 0; i < num_gauss; i++) {
        BaseFloat w_i = tmp_weight(i);
        total_weight += w_i;
        within_class_covar.AddSp(w_i, tmp_covars[i]);
        between_class_covar.AddVec2(w_i, tmp_means.Row(i));
        global_mean.AddVec(w_i, tmp_means.Row(i));
      }
      KALDI_ASSERT(total_weight > 0);
      if (fabs(total_weight - 1.0) > 0.001) {
        KALDI_WARN << "Total weight across the GMMs is " << (total_weight)
            << ", renormalizing.";
        global_mean.Scale(1.0 / total_weight);
        within_class_covar.Scale(1.0 / total_weight);
        between_class_covar.Scale(1.0 / total_weight);
      }
      between_class_covar.AddVec2(-1.0, global_mean);
    }
  
    TpMatrix<BaseFloat> chol(dim);
    chol.Cholesky(within_class_covar);  // Sigma_W = L L^T
    TpMatrix<BaseFloat> chol_inv(chol);
    chol_inv.InvertDouble();
    Matrix<BaseFloat> chol_full(dim, dim);
    chol_full.CopyFromTp(chol_inv);
    SpMatrix<BaseFloat> LBL(dim);
    // LBL = L^{-1} \Sigma_B L^{-T}
    LBL.AddMat2Sp(1.0, chol_full, kNoTrans, between_class_covar, 0.0);
    Vector<BaseFloat> Dvec(dim);
    Matrix<BaseFloat> U(dim, dim);
    LBL.Eig(&Dvec, &U);
    SortSvd(&Dvec, &U);
  
    xform->Resize(dim, dim);
    chol_full.CopyFromTp(chol);
    // T := L U, eq (23)
    xform->AddMatMat(1.0, chol_full, kNoTrans, U, kNoTrans, 0.0);
  
  #ifdef KALDI_PARANOID
    Matrix<BaseFloat> inv_xform(*xform);
    inv_xform.InvertDouble();
    {  // Check that T*within_class_covar*T' = I.
      Matrix<BaseFloat> wc_covar_full(dim, dim), tmp(dim, dim);
      wc_covar_full.CopyFromSp(within_class_covar);
      tmp.AddMatMat(1.0, inv_xform, kNoTrans, wc_covar_full, kNoTrans, 0.0);
      wc_covar_full.AddMatMat(1.0, tmp, kNoTrans, inv_xform, kTrans, 0.0);
      KALDI_ASSERT(wc_covar_full.IsUnit(0.01));
    }
    {  // Check that T*between_class_covar*T' = diagonal.
      Matrix<BaseFloat> bc_covar_full(dim, dim), tmp(dim, dim);
      bc_covar_full.CopyFromSp(between_class_covar);
      tmp.AddMatMat(1.0, inv_xform, kNoTrans, bc_covar_full, kNoTrans, 0.0);
      bc_covar_full.AddMatMat(1.0, tmp, kNoTrans, inv_xform, kTrans, 0.0);
      KALDI_ASSERT(bc_covar_full.IsDiagonal(0.01));
    }
  #endif
  }
  
  void AmSgmm2::ComputePerSpkDerivedVars(Sgmm2PerSpkDerivedVars *vars) const {
    KALDI_ASSERT(vars != NULL);
    if (vars->v_s.Dim() != 0) {
      KALDI_ASSERT(vars->v_s.Dim() == SpkSpaceDim());
      vars->o_s.Resize(NumGauss(), FeatureDim());
      int32 num_gauss = NumGauss();
      // first compute the o_i^{(s)} quantities.
      for (int32 i = 0; i < num_gauss; i++) {
         // Eqn. (32): o_i^{(s)} = N_i v^{(s)}
        vars->o_s.Row(i).AddMatVec(1.0, N_[i], kNoTrans, vars->v_s, 0.0);
      }
      // the rest relates to the SSGMM.  We only need to to this
      // if we're using speaker-dependent weights.
      if (HasSpeakerDependentWeights()) {
        vars->log_d_jms.clear();
        vars->log_d_jms.resize(NumGroups());
        vars->log_b_is.Resize(NumGauss());
        vars->log_b_is.AddMatVec(1.0, u_, kNoTrans, vars->v_s, 0.0);
        vars->b_is.Resize(NumGauss());
        vars->b_is.CopyFromVec(vars->log_b_is);
        vars->b_is.ApplyExp();
        for (int32 i = 0; i < vars->b_is.Dim(); i++) {
          if (vars->b_is(i) - vars->b_is(i) != 0.0) { // NaN.
            vars->b_is(i) = 1.0;
            KALDI_WARN << "Set NaN in b_is to 1.0";
          }
        }
      } else {
        vars->b_is.Resize(0);
        vars->log_b_is.Resize(0);
        vars->log_d_jms.resize(0);
      }
    } else {
      vars->Clear(); // make sure everything is cleared.
    }
  }
  
  BaseFloat AmSgmm2::GaussianSelection(const Sgmm2GselectConfig &config,
                                      const VectorBase<BaseFloat> &data,
                                      std::vector<int32> *gselect) const {
    KALDI_ASSERT(diag_ubm_.NumGauss() != 0 &&
                 diag_ubm_.NumGauss() == full_ubm_.NumGauss() &&
                 diag_ubm_.Dim() == data.Dim());
    KALDI_ASSERT(config.diag_gmm_nbest > 0 && config.full_gmm_nbest > 0 &&
                 config.full_gmm_nbest < config.diag_gmm_nbest);
    int32 num_gauss = diag_ubm_.NumGauss();
  
    std::vector< std::pair<BaseFloat, int32> > pruned_pairs;
    if (config.diag_gmm_nbest < num_gauss) {    Vector<BaseFloat> loglikes(num_gauss);
      diag_ubm_.LogLikelihoods(data, &loglikes);
      Vector<BaseFloat> loglikes_copy(loglikes);
      BaseFloat *ptr = loglikes_copy.Data();
      std::nth_element(ptr, ptr+num_gauss-config.diag_gmm_nbest, ptr+num_gauss);
      BaseFloat thresh = ptr[num_gauss-config.diag_gmm_nbest];
      for (int32 g = 0; g < num_gauss; g++)
        if (loglikes(g) >= thresh)  // met threshold for diagonal phase.
          pruned_pairs.push_back(
              std::make_pair(full_ubm_.ComponentLogLikelihood(data, g), g));
    } else {
      Vector<BaseFloat> loglikes(num_gauss);
      full_ubm_.LogLikelihoods(data, &loglikes);
      for (int32 g = 0; g < num_gauss; g++)
        pruned_pairs.push_back(std::make_pair(loglikes(g), g));
    }
    KALDI_ASSERT(!pruned_pairs.empty());
    if (pruned_pairs.size() > static_cast<size_t>(config.full_gmm_nbest)) {
      std::nth_element(pruned_pairs.begin(),
                       pruned_pairs.end() - config.full_gmm_nbest,
                       pruned_pairs.end());
      pruned_pairs.erase(pruned_pairs.begin(),
                         pruned_pairs.end() - config.full_gmm_nbest);
    }
    Vector<BaseFloat> loglikes_tmp(pruned_pairs.size());  // for return value.
    KALDI_ASSERT(gselect != NULL);
    gselect->resize(pruned_pairs.size());
    // Make sure pruned Gaussians appear from best to worst.
    std::sort(pruned_pairs.begin(), pruned_pairs.end(),
              std::greater< std::pair<BaseFloat, int32> >());
    for (size_t i = 0; i < pruned_pairs.size(); i++) {
      loglikes_tmp(i) = pruned_pairs[i].first;
      (*gselect)[i] = pruned_pairs[i].second;
    }
    return loglikes_tmp.LogSumExp();
  }
  
  void Sgmm2GauPost::Write(std::ostream &os, bool binary) const {
    WriteToken(os, binary, "<Sgmm2GauPost>");
    int32 T = this->size();
    WriteBasicType(os, binary, T);
    for (int32 t = 0; t < T; t++) {
      WriteToken(os, binary, "<gselect>");
      WriteIntegerVector(os, binary, (*this)[t].gselect);
      WriteToken(os, binary, "<tids>");
      WriteIntegerVector(os, binary, (*this)[t].tids);
      KALDI_ASSERT((*this)[t].tids.size() == (*this)[t].posteriors.size());
      for (size_t i = 0; i < (*this)[t].posteriors.size(); i++) {
        (*this)[t].posteriors[i].Write(os, binary);
      }
    }
    WriteToken(os, binary, "</Sgmm2GauPost>");
  }
  
  
  void Sgmm2GauPost::Read(std::istream &is, bool binary) {
    ExpectToken(is, binary, "<Sgmm2GauPost>");
    int32 T;
    ReadBasicType(is, binary, &T);
    KALDI_ASSERT(T >= 0);
    this->resize(T);
    for (int32 t = 0; t < T; t++) {
      ExpectToken(is, binary, "<gselect>");
      ReadIntegerVector(is, binary, &((*this)[t].gselect));
      ExpectToken(is, binary, "<tids>");
      ReadIntegerVector(is, binary, &((*this)[t].tids));
      size_t sz = (*this)[t].tids.size();
      (*this)[t].posteriors.resize(sz);
      for (size_t i = 0; i < sz; i++)
        (*this)[t].posteriors[i].Read(is, binary);
    }
    ExpectToken(is, binary, "</Sgmm2GauPost>");
  }
  
  
  
  }  // namespace kaldi
  