// ivector/plda.cc
  
  // Copyright 2013     Daniel Povey
  //           2015     David Snyder
  
  // 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 <vector>
  #include "ivector/plda.h"
  
  namespace kaldi {
  
  void Plda::Write(std::ostream &os, bool binary) const {
    WriteToken(os, binary, "<Plda>");
    mean_.Write(os, binary);
    transform_.Write(os, binary);
    psi_.Write(os, binary);
    WriteToken(os, binary, "</Plda>");
  }
  
  void Plda::Read(std::istream &is, bool binary) {
    ExpectToken(is, binary, "<Plda>");
    mean_.Read(is, binary);
    transform_.Read(is, binary);
    psi_.Read(is, binary);
    ExpectToken(is, binary, "</Plda>");
    ComputeDerivedVars();
  }
  
  template<class Real>
  /// This function computes a projection matrix that when applied makes the
  /// covariance unit (i.e. all 1).
  static void ComputeNormalizingTransform(const SpMatrix<Real> &covar,
                                          MatrixBase<Real> *proj) {
    int32 dim = covar.NumRows();
    TpMatrix<Real> C(dim);  // Cholesky of covar, covar = C C^T
    C.Cholesky(covar);
    C.Invert();  // The matrix that makes covar unit is C^{-1}, because
                 // C^{-1} covar C^{-T} = C^{-1} C C^T C^{-T} = I.
    proj->CopyFromTp(C, kNoTrans);  // set "proj" to C^{-1}.
  }
  
  
  void Plda::ComputeDerivedVars() {
    KALDI_ASSERT(Dim() > 0);
    offset_.Resize(Dim());
    offset_.AddMatVec(-1.0, transform_, kNoTrans, mean_, 0.0);
  }
  
  
  /**
     This comment explains the thinking behind the function LogLikelihoodRatio.
     The reference is "Probabilistic Linear Discriminant Analysis" by
     Sergey Ioffe, ECCV 2006.
  
     I'm looking at the un-numbered equation between eqs. (4) and (5),
     that says
       P(u^p | u^g_{1...n}) =  N (u^p | \frac{n \Psi}{n \Psi + I} \bar{u}^g, I + \frac{\Psi}{n\Psi + I})
  
     Here, the superscript ^p refers to the "probe" example (e.g. the example
     to be classified), and u^g_1 is the first "gallery" example, i.e. the first
     training example of that class.  \psi is the between-class covariance
     matrix, assumed to be diagonalized, and I can be interpreted as the within-class
     covariance matrix which we have made unit.
  
     We want the likelihood ratio P(u^p | u^g_{1..n}) / P(u^p), where the
     numerator is the probability of u^p given that it's in that class, and the
     denominator is the probability of u^p with no class assumption at all
     (e.g. in its own class).
  
     The expression above even works for n = 0 (e.g. the denominator of the likelihood
     ratio), where it gives us
       P(u^p) = N(u^p | 0, I + \Psi)
     i.e. it's distributed with zero mean and covarance (within + between).
     The likelihood ratio we want is:
        N(u^p | \frac{n \Psi}{n \Psi + I} \bar{u}^g, I + \frac{\Psi}{n \Psi + I}) /
        N(u^p | 0, I + \Psi)
     where \bar{u}^g is the mean of the "gallery examples"; and we can expand the
     log likelihood ratio as
       - 0.5 [ (u^p - m) (I + \Psi/(n \Psi + I))^{-1} (u^p - m)  +  logdet(I + \Psi/(n \Psi + I)) ]
       + 0.5 [u^p (I + \Psi) u^p  +  logdet(I + \Psi) ]
     where m = (n \Psi)/(n \Psi + I) \bar{u}^g.
  
   */
  
  double Plda::GetNormalizationFactor(
      const VectorBase<double> &transformed_ivector,
      int32 num_examples) const {
    KALDI_ASSERT(num_examples > 0);
    // Work out the normalization factor.  The covariance for an average over
    // "num_examples" training iVectors equals \Psi + I/num_examples.
    Vector<double> transformed_ivector_sq(transformed_ivector);
    transformed_ivector_sq.ApplyPow(2.0);
    // inv_covar will equal 1.0 / (\Psi + I/num_examples).
    Vector<double> inv_covar(psi_);
    inv_covar.Add(1.0 / num_examples);
    inv_covar.InvertElements();
    // "transformed_ivector" should have covariance (\Psi + I/num_examples), i.e.
    // within-class/num_examples plus between-class covariance.  So
    // transformed_ivector_sq . (I/num_examples + \Psi)^{-1} should be equal to
    //  the dimension.
    double dot_prod = VecVec(inv_covar, transformed_ivector_sq);
    return sqrt(Dim() / dot_prod);
  }
  
  
  double Plda::TransformIvector(const PldaConfig &config,
                                const VectorBase<double> &ivector,
                                int32 num_examples,
                                VectorBase<double> *transformed_ivector) const {
    KALDI_ASSERT(ivector.Dim() == Dim() && transformed_ivector->Dim() == Dim());
    double normalization_factor;
    transformed_ivector->CopyFromVec(offset_);
    transformed_ivector->AddMatVec(1.0, transform_, kNoTrans, ivector, 1.0);
    if (config.simple_length_norm)
      normalization_factor = sqrt(transformed_ivector->Dim())
        / transformed_ivector->Norm(2.0);
    else
      normalization_factor = GetNormalizationFactor(*transformed_ivector,
                                                    num_examples);
    if (config.normalize_length)
      transformed_ivector->Scale(normalization_factor);
    return normalization_factor;
  }
  
  // "float" version of TransformIvector.
  float Plda::TransformIvector(const PldaConfig &config,
                               const VectorBase<float> &ivector,
                               int32 num_examples,
                               VectorBase<float> *transformed_ivector) const {
    Vector<double> tmp(ivector), tmp_out(ivector.Dim());
    float ans = TransformIvector(config, tmp, num_examples, &tmp_out);
    transformed_ivector->CopyFromVec(tmp_out);
    return ans;
  }
  
  
  // There is an extended comment within this file, referencing a paper by
  // Ioffe, that may clarify what this function is doing.
  double Plda::LogLikelihoodRatio(
      const VectorBase<double> &transformed_train_ivector,
      int32 n, // number of training utterances.
      const VectorBase<double> &transformed_test_ivector) const {
    int32 dim = Dim();
    double loglike_given_class, loglike_without_class;
    { // work out loglike_given_class.
      // "mean" will be the mean of the distribution if it comes from the
      // training example.  The mean is \frac{n \Psi}{n \Psi + I} \bar{u}^g
      // "variance" will be the variance of that distribution, equal to
      // I + \frac{\Psi}{n\Psi + I}.
      Vector<double> mean(dim, kUndefined);
      Vector<double> variance(dim, kUndefined);
      for (int32 i = 0; i < dim; i++) {
        mean(i) = n * psi_(i) / (n * psi_(i) + 1.0)
          * transformed_train_ivector(i);
        variance(i) = 1.0 + psi_(i) / (n * psi_(i) + 1.0);
      }
      double logdet = variance.SumLog();
      Vector<double> sqdiff(transformed_test_ivector);
      sqdiff.AddVec(-1.0, mean);
      sqdiff.ApplyPow(2.0);
      variance.InvertElements();
      loglike_given_class = -0.5 * (logdet + M_LOG_2PI * dim +
                                    VecVec(sqdiff, variance));
    }
    { // work out loglike_without_class.  Here the mean is zero and the variance
      // is I + \Psi.
      Vector<double> sqdiff(transformed_test_ivector); // there is no offset.
      sqdiff.ApplyPow(2.0);
      Vector<double> variance(psi_);
      variance.Add(1.0); // I + \Psi.
      double logdet = variance.SumLog();
      variance.InvertElements();
      loglike_without_class = -0.5 * (logdet + M_LOG_2PI * dim +
                                      VecVec(sqdiff, variance));
    }
    double loglike_ratio = loglike_given_class - loglike_without_class;
    return loglike_ratio;
  }
  
  
  void Plda::SmoothWithinClassCovariance(double smoothing_factor) {
    KALDI_ASSERT(smoothing_factor >= 0.0 && smoothing_factor <= 1.0);
    // smoothing_factor > 1.0 is possible but wouldn't really make sense.
  
    KALDI_LOG << "Smoothing within-class covariance by " << smoothing_factor
              << ", Psi is initially: " << psi_;
    Vector<double> within_class_covar(Dim());
    within_class_covar.Set(1.0); // It's now the current within-class covariance
                                 // (a diagonal matrix) in the space transformed
                                 // by transform_.
    within_class_covar.AddVec(smoothing_factor, psi_);
    /// We now revise our estimate of the within-class covariance to this
    /// larger value.  This means that the transform has to change to as
    /// to make this new, larger covariance unit.  And our between-class
    /// covariance in this space is now less.
  
    psi_.DivElements(within_class_covar);
    KALDI_LOG << "New value of Psi is " << psi_;
  
    within_class_covar.ApplyPow(-0.5);
    transform_.MulRowsVec(within_class_covar);
  
    ComputeDerivedVars();
  }
  
  void Plda::ApplyTransform(const Matrix<double> &in_transform) {
    KALDI_ASSERT(in_transform.NumRows() <= Dim()
      && in_transform.NumCols() == Dim());
  
    // Apply in_transform to mean_.
    Vector<double> mean_new(in_transform.NumRows());
    mean_new.AddMatVec(1.0, in_transform, kNoTrans, mean_, 0.0);
    mean_.Resize(in_transform.NumRows());
    mean_.CopyFromVec(mean_new);
  
    SpMatrix<double> between_var(in_transform.NumCols()),
                     within_var(in_transform.NumCols()),
                     psi_mat(in_transform.NumCols()),
                     between_var_new(Dim()),
                     within_var_new(Dim());
    Matrix<double> transform_invert(transform_);
  
    // Next, compute the between_var and within_var that existed
    // prior to diagonalization.
    psi_mat.AddDiagVec(1.0, psi_);
    transform_invert.Invert();
    within_var.AddMat2(1.0, transform_invert, kNoTrans, 0.0);
    between_var.AddMat2Sp(1.0, transform_invert, kNoTrans, psi_mat, 0.0);
  
    // Next, transform the variances using the input transformation.
    between_var_new.AddMat2Sp(1.0, in_transform, kNoTrans, between_var, 0.0);
    within_var_new.AddMat2Sp(1.0, in_transform, kNoTrans, within_var, 0.0);
  
    // Finally, we need to recompute psi_ and transform_. The remainder of
    // the code in this function  is a lightly modified copy of
    // PldaEstimator::GetOutput().
    Matrix<double> transform1(Dim(), Dim());
    ComputeNormalizingTransform(within_var_new, &transform1);
    // Now transform is a matrix that if we project with it,
    // within_var becomes unit.
    // between_var_proj is between_var after projecting with transform1.
    SpMatrix<double> between_var_proj(Dim());
    between_var_proj.AddMat2Sp(1.0, transform1, kNoTrans, between_var_new, 0.0);
  
    Matrix<double> U(Dim(), Dim());
    Vector<double> s(Dim());
    // Do symmetric eigenvalue decomposition between_var_proj = U diag(s) U^T,
    // where U is orthogonal.
    between_var_proj.Eig(&s, &U);
  
    KALDI_ASSERT(s.Min() >= 0.0);
    int32 n;
    s.ApplyFloor(0.0, &n);
    if (n > 0) {
      KALDI_WARN << "Floored " << n << " eigenvalues of between-class "
                 << "variance to zero.";
    }
    // Sort from greatest to smallest eigenvalue.
    SortSvd(&s, &U);
  
    // The transform U^T will make between_var_proj diagonal with value s
    // (i.e. U^T U diag(s) U U^T = diag(s)).  The final transform that
    // makes within_var unit and between_var diagonal is U^T transform1,
    // i.e. first transform1 and then U^T.
    transform_.Resize(Dim(), Dim());
    transform_.AddMatMat(1.0, U, kTrans, transform1, kNoTrans, 0.0);
    psi_.Resize(Dim());
    psi_.CopyFromVec(s);
    ComputeDerivedVars();
  }
  
  void PldaStats::AddSamples(double weight,
                             const Matrix<double> &group) {
    if (dim_ == 0) {
      Init(group.NumCols());
    } else {
      KALDI_ASSERT(dim_ == group.NumCols());
    }
    int32 n = group.NumRows(); // number of examples for this class
    Vector<double> *mean = new Vector<double>(dim_);
    mean->AddRowSumMat(1.0 / n, group);
  
    offset_scatter_.AddMat2(weight, group, kTrans, 1.0);
    // the following statement has the same effect as if we
    // had first subtracted the mean from each element of
    // the group before the statement above.
    offset_scatter_.AddVec2(-n * weight, *mean);
  
    class_info_.push_back(ClassInfo(weight, mean, n));
  
    num_classes_ ++;
    num_examples_ += n;
    class_weight_ += weight;
    example_weight_ += weight * n;
  
    sum_.AddVec(weight, *mean);
  }
  
  PldaStats::~PldaStats() {
    for (size_t i = 0; i < class_info_.size(); i++)
      delete class_info_[i].mean;
  }
  
  bool PldaStats::IsSorted() const {
    for (size_t i = 0; i + 1 < class_info_.size(); i++)
      if (class_info_[i+1] < class_info_[i])
        return false;
    return true;
  }
  
  void PldaStats::Init(int32 dim) {
    KALDI_ASSERT(dim_ == 0);
    dim_ = dim;
    num_classes_ = 0;
    num_examples_ = 0;
    class_weight_ = 0.0;
    example_weight_ = 0.0;
    sum_.Resize(dim);
    offset_scatter_.Resize(dim);
    KALDI_ASSERT(class_info_.empty());
  }
  
  
  PldaEstimator::PldaEstimator(const PldaStats &stats):
      stats_(stats) {
    KALDI_ASSERT(stats.IsSorted());
    InitParameters();
  }
  
  
  double PldaEstimator::ComputeObjfPart1() const {
    // Returns the part of the objf relating to offsets from the class means.
    // within_class_count equals the sum over the classes, of the weight of that
    // class (normally 1) times (1 - #examples) of that class, which equals the
    // rank of the covariance we're modeling.  We imagine that we're modeling (1 -
    // #examples) separate samples, each with the within-class covariance.. the
    // argument is a little complicated and involves an orthogonal complement of a
    // matrix whose first row computes the mean.
  
    double within_class_count = stats_.example_weight_ - stats_.class_weight_,
        within_logdet, det_sign;
    SpMatrix<double> inv_within_var(within_var_);
    inv_within_var.Invert(&within_logdet, &det_sign);
    KALDI_ASSERT(det_sign == 1 && "Within-class covariance is singular");
  
    double objf = -0.5 * (within_class_count * (within_logdet + M_LOG_2PI * Dim())
                          + TraceSpSp(inv_within_var, stats_.offset_scatter_));
    return objf;
  }
  
  double PldaEstimator::ComputeObjfPart2() const {
    double tot_objf = 0.0;
  
    int32 n = -1; // the number of examples for the current class
    SpMatrix<double> combined_inv_var(Dim());
    // combined_inv_var = (between_var_ + within_var_ / n)^{-1}
    double combined_var_logdet;
  
    for (size_t i = 0; i < stats_.class_info_.size(); i++) {
      const ClassInfo &info = stats_.class_info_[i];
      if (info.num_examples != n) {
        n = info.num_examples;
        // variance of mean of n examples is between-class + 1/n * within-class
        combined_inv_var.CopyFromSp(between_var_);
        combined_inv_var.AddSp(1.0 / n, within_var_);
        combined_inv_var.Invert(&combined_var_logdet);
      }
      Vector<double> mean (*(info.mean));
      mean.AddVec(-1.0 / stats_.class_weight_, stats_.sum_);
      tot_objf += info.weight * -0.5 * (combined_var_logdet + M_LOG_2PI * Dim()
                                        + VecSpVec(mean, combined_inv_var, mean));
    }
    return tot_objf;
  }
  
  double PldaEstimator::ComputeObjf() const {
    double ans1 = ComputeObjfPart1(),
        ans2 = ComputeObjfPart2(),
        ans = ans1 + ans2,
        example_weights = stats_.example_weight_,
        normalized_ans = ans / example_weights;
    KALDI_LOG << "Within-class objf per sample is " << (ans1 / example_weights)
              << ", between-class is " << (ans2 / example_weights)
              << ", total is " << normalized_ans;
    return normalized_ans;
  }
  
  void PldaEstimator::InitParameters() {
    within_var_.Resize(Dim());
    within_var_.SetUnit();
    between_var_.Resize(Dim());
    between_var_.SetUnit();
  }
  
  void PldaEstimator::ResetPerIterStats() {
    within_var_stats_.Resize(Dim());
    within_var_count_ = 0.0;
    between_var_stats_.Resize(Dim());
    between_var_count_ = 0.0;
  }
  
  void PldaEstimator::GetStatsFromIntraClass() {
    within_var_stats_.AddSp(1.0, stats_.offset_scatter_);
    // Note: in the normal case, the expression below will be equal to the sum
    // over the classes, of (1-n), where n is the #examples for that class.  That
    // is the rank of the scatter matrix that "offset_scatter_" has for that
    // class. [if weights other than 1.0 are used, it will be different.]
    within_var_count_ += (stats_.example_weight_ - stats_.class_weight_);
  }
  
  
  /**
     GetStatsFromClassMeans() is the more complicated part of PLDA estimation.
     Let's suppose the mean of a particular class is m, and suppose that
     that class had n examples.  We suppose that
       m ~ N(0, between_var_ + 1/n within_var_)
     i.e. m is Gaussian-distributed with zero mean and variance equal to the
     between-class variance plus 1/n times the within-class variance.  Now, m
     is observed (as stats_.class_info_[something].mean).  We're doing an E-M
     procedure where we treat m as the sum of two variables:
       m = x + y
     where
       x ~ N(0, between_var_)
       y ~ N(0, 1/n * within_var_)
     The distribution of x will contribute to the stats of between_var_, and
     y to within_var_.  Now, y = m - x, so we can focus on working out the
     distribution of x and then we can very simply get the distribution of y.
     The following expression also includes the likelihood of y as a function of
     x.  Note: the C is different from line to line.
  
     log p(x) = C - 0.5 ( x^T between_var^{-1} x  + (m-x)^T (1/n within_var)^{-1) (m-x) )
              = C - 0.5 x^T (between_var^{-1} + n within_var^{-1}) x + x^T z
  
              where z = n within_var^{-1} m, and we can write this as:
  
     log p(x) = C - 0.5 (x-w)^T (between_var^{-1} + n within_var^{-1}) (x-w)
  
      where x^T (between_var^{-1} + n within_var^{-1}) w = x^T z, i.e.
         (between_var^{-1} + n within_var^{-1}) w = z = n within_var^{-1} m, so
  
         w = (between_var^{-1} + n within_var^{-1})^{-1} * n within_var^{-1} m
  
      We can see that the distribution over x is Gaussian, with mean w and variance
       (between_var^{-1} + n within_var^{-1})^{-1}.
      The distribution over y is Gaussian with the same variance, and mean m - w.
      So the update to the between-var stats will be:
         between-var-stats += w w^T + (between_var^{-1} + n within_var^{-1})^{-1}.
      and the update to the within-var stats will be:
         within-var-stats += n ( (m-w) (m-w)^T (between_var^{-1} + n within_var^{-1})^{-1} ).
  
      The drawback of this formulation is that each time we encounter a different
      value of n (number of examples) we will have to do a different matrix
      inversion.  We'll try to improve on this later using a suitable transform.
   */
  
  void PldaEstimator::GetStatsFromClassMeans() {
    SpMatrix<double> between_var_inv(between_var_);
    between_var_inv.Invert();
    SpMatrix<double> within_var_inv(within_var_);
    within_var_inv.Invert();
    // mixed_var will equal (between_var^{-1} + n within_var^{-1})^{-1}.
    SpMatrix<double> mixed_var(Dim());
    int32 n = -1; // the current number of examples for the class.
  
    for (size_t i = 0; i < stats_.class_info_.size(); i++) {
      const ClassInfo &info = stats_.class_info_[i];
      double weight = info.weight;
      if (info.num_examples != n) {
        n = info.num_examples;
        mixed_var.CopyFromSp(between_var_inv);
        mixed_var.AddSp(n, within_var_inv);
        mixed_var.Invert();
      }
      Vector<double> m = *(info.mean); // the mean for this class.
      m.AddVec(-1.0 / stats_.class_weight_, stats_.sum_); // remove global mean
      Vector<double> temp(Dim()); // n within_var^{-1} m
      temp.AddSpVec(n, within_var_inv, m, 0.0);
      Vector<double> w(Dim()); // w, as defined in the comment.
      w.AddSpVec(1.0, mixed_var, temp, 0.0);
      Vector<double> m_w(m); // m - w
      m_w.AddVec(-1.0, w);
      between_var_stats_.AddSp(weight, mixed_var);
      between_var_stats_.AddVec2(weight, w);
      between_var_count_ += weight;
      within_var_stats_.AddSp(weight * n, mixed_var);
      within_var_stats_.AddVec2(weight * n, m_w);
      within_var_count_ += weight;
    }
  }
  
  void PldaEstimator::EstimateFromStats() {
    within_var_.CopyFromSp(within_var_stats_);
    within_var_.Scale(1.0 / within_var_count_);
    between_var_.CopyFromSp(between_var_stats_);
    between_var_.Scale(1.0 / between_var_count_);
  
    KALDI_LOG << "Trace of within-class variance is " << within_var_.Trace();
    KALDI_LOG << "Trace of between-class variance is " << between_var_.Trace();
  }
  
  
  void PldaEstimator::EstimateOneIter() {
    ResetPerIterStats();
    GetStatsFromIntraClass();
    GetStatsFromClassMeans();
    EstimateFromStats();
    KALDI_VLOG(2) << "Objective function is " << ComputeObjf();
  }
  
  
  void PldaEstimator::Estimate(const PldaEstimationConfig &config,
                               Plda *plda) {
    KALDI_ASSERT(stats_.example_weight_ > 0 && "Cannot estimate with no stats");
    for (int32 i = 0; i < config.num_em_iters; i++) {
      KALDI_LOG << "Plda estimation iteration " << i
                << " of " << config.num_em_iters;
      EstimateOneIter();
    }
    GetOutput(plda);
  }
  
  
  void PldaEstimator::GetOutput(Plda *plda) {
    plda->mean_ = stats_.sum_;
    plda->mean_.Scale(1.0 / stats_.class_weight_);
    KALDI_LOG << "Norm of mean of iVector distribution is "
              << plda->mean_.Norm(2.0);
  
    Matrix<double> transform1(Dim(), Dim());
    ComputeNormalizingTransform(within_var_, &transform1);
    // now transform is a matrix that if we project with it,
    // within_var_ becomes unit.
  
    // between_var_proj is between_var after projecting with transform1.
    SpMatrix<double> between_var_proj(Dim());
    between_var_proj.AddMat2Sp(1.0, transform1, kNoTrans, between_var_, 0.0);
  
    Matrix<double> U(Dim(), Dim());
    Vector<double> s(Dim());
    // Do symmetric eigenvalue decomposition between_var_proj = U diag(s) U^T,
    // where U is orthogonal.
    between_var_proj.Eig(&s, &U);
  
    KALDI_ASSERT(s.Min() >= 0.0);
    int32 n;
    s.ApplyFloor(0.0, &n);
    if (n > 0) {
      KALDI_WARN << "Floored " << n << " eigenvalues of between-class "
                 << "variance to zero.";
    }
    // Sort from greatest to smallest eigenvalue.
    SortSvd(&s, &U);
  
    // The transform U^T will make between_var_proj diagonal with value s
    // (i.e. U^T U diag(s) U U^T = diag(s)).  The final transform that
    // makes within_var_ unit and between_var_ diagonal is U^T transform1,
    // i.e. first transform1 and then U^T.
  
    plda->transform_.Resize(Dim(), Dim());
    plda->transform_.AddMatMat(1.0, U, kTrans, transform1, kNoTrans, 0.0);
    plda->psi_ = s;
  
    KALDI_LOG << "Diagonal of between-class variance in normalized space is " << s;
  
    if (GetVerboseLevel() >= 2) { // at higher verbose levels, do a self-test
                                  // (just tests that this function does what it
                                  // should).
      SpMatrix<double> tmp_within(Dim());
      tmp_within.AddMat2Sp(1.0, plda->transform_, kNoTrans, within_var_, 0.0);
      KALDI_ASSERT(tmp_within.IsUnit(0.0001));
      SpMatrix<double> tmp_between(Dim());
      tmp_between.AddMat2Sp(1.0, plda->transform_, kNoTrans, between_var_, 0.0);
      KALDI_ASSERT(tmp_between.IsDiagonal(0.0001));
      Vector<double> psi(Dim());
      psi.CopyDiagFromSp(tmp_between);
      AssertEqual(psi, plda->psi_);
    }
    plda->ComputeDerivedVars();
  }
  
  void PldaUnsupervisedAdaptor::AddStats(double weight,
                                         const Vector<double> &ivector) {
    if (mean_stats_.Dim() == 0) {
      mean_stats_.Resize(ivector.Dim());
      variance_stats_.Resize(ivector.Dim());
    }
    KALDI_ASSERT(weight >= 0.0);
    tot_weight_ += weight;
    mean_stats_.AddVec(weight, ivector);
    variance_stats_.AddVec2(weight, ivector);
  }
  
  void PldaUnsupervisedAdaptor::AddStats(double weight,
                                         const Vector<float> &ivector) {
    Vector<double> ivector_dbl(ivector);
    this->AddStats(weight, ivector_dbl);
  }
  
  void PldaUnsupervisedAdaptor::UpdatePlda(const PldaUnsupervisedAdaptorConfig &config,
                                           Plda *plda) const {
    KALDI_ASSERT(tot_weight_ > 0.0);
    int32 dim = mean_stats_.Dim();
    KALDI_ASSERT(dim == plda->Dim());
    Vector<double> mean(mean_stats_);
    mean.Scale(1.0 / tot_weight_);
    SpMatrix<double> variance(variance_stats_);
    variance.Scale(1.0 / tot_weight_);
    variance.AddVec2(-1.0, mean);  // Make it the uncentered variance.
  
    // mean_diff of the adaptation data from the training data.  We optionally add
    // this to our total covariance matrix
    Vector<double> mean_diff(mean);
    mean_diff.AddVec(-1.0, plda->mean_);
    KALDI_ASSERT(config.mean_diff_scale >= 0.0);
    variance.AddVec2(config.mean_diff_scale, mean_diff);
  
    // update the plda's mean data-member with our adaptation-data mean.
    plda->mean_.CopyFromVec(mean);
  
  
    // transform_model_ is a row-scaled version of plda->transform_ that
    // transforms into the space where the total covariance is 1.0.  Because
    // plda->transform_ transforms into a space where the within-class covar is
    // 1.0 and the the between-class covar is diag(plda->psi_), we need to scale
    // each dimension i by 1.0 / sqrt(1.0 + plda->psi_(i))
  
    Matrix<double> transform_mod(plda->transform_);
    for (int32 i = 0; i < dim; i++)
      transform_mod.Row(i).Scale(1.0 / sqrt(1.0 + plda->psi_(i)));
  
    // project the variance of the adaptation set into this space where
    // the total covariance is unit.
    SpMatrix<double> variance_proj(dim);
    variance_proj.AddMat2Sp(1.0, transform_mod, kNoTrans,
                            variance, 0.0);
  
    // Do eigenvalue decomposition of variance_proj; this will tell us the
    // directions in which the adaptation-data covariance is more than
    // the training-data covariance.
    Matrix<double> P(dim, dim);
    Vector<double> s(dim);
    variance_proj.Eig(&s, &P);
    SortSvd(&s, &P);
    KALDI_LOG << "Eigenvalues of adaptation-data total-covariance in space where "
              << "training-data total-covariance is unit, is: " << s;
  
    // W, B are the (within,between)-class covars in the space transformed by
    // transform_mod.
    SpMatrix<double> W(dim), B(dim);
    for (int32 i = 0; i < dim; i++) {
      W(i, i) =           1.0 / (1.0 + plda->psi_(i)),
      B(i, i) = plda->psi_(i) / (1.0 + plda->psi_(i));
    }
  
    // OK, so variance_proj (projected by transform_mod) is P diag(s) P^T.
    // Suppose that after transform_mod we project by P^T.  Then the adaptation-data's
    // variance would be P^T P diag(s) P^T P = diag(s), and the PLDA model's
    // within class variance would be P^T W P and its between-class variance would be
    // P^T B P.  We'd still have that W+B = I in this space.
    // First let's compute these projected variances... we call the "proj2" because
    // it's after the data has been projected twice (actually, transformed, as there is no
    // dimension loss), by transform_mod and then P^T.
  
    SpMatrix<double> Wproj2(dim), Bproj2(dim);
    Wproj2.AddMat2Sp(1.0, P, kTrans, W, 0.0);
    Bproj2.AddMat2Sp(1.0, P, kTrans, B, 0.0);
  
    Matrix<double> Ptrans(P, kTrans);
  
    SpMatrix<double> Wproj2mod(Wproj2), Bproj2mod(Bproj2);
  
    for (int32 i = 0; i < dim; i++) {
      // For this eigenvalue, compute the within-class covar projected with this direction,
      // and the same for between.
      BaseFloat within = Wproj2(i, i),
          between = Bproj2(i, i);
      KALDI_LOG << "For " << i << "'th eigenvalue, value is " << s(i)
                << ", within-class covar in this direction is " << within
                << ", between-class is " << between;
      if (s(i) > 1.0) {
        double excess_eig = s(i) - 1.0;
        double excess_within_covar = excess_eig * config.within_covar_scale,
            excess_between_covar = excess_eig * config.between_covar_scale;
        Wproj2mod(i, i) += excess_within_covar;
        Bproj2mod(i, i) += excess_between_covar;
      } /*
          Below I was considering a method like below, to try to scale up
          the dimensions that had less variance than expected in our sample..
          this didn't help, and actually when I set that power to +0.2 instead
          of -0.5 it gave me an improvement on sre08.  But I'm not sure
          about this.. it just doesn't seem right.
        else {
        BaseFloat scale = pow(std::max(1.0e-10, s(i)), -0.5);
        BaseFloat max_scale = 10.0;  // I'll make this configurable later.
        scale = std::min(scale, max_scale);
        Ptrans.Row(i).Scale(scale);
        } */
    }
  
    // combined transform "transform_mod" and then P^T that takes us to the space
    // where {W,B}proj2{,mod} are.
    Matrix<double> combined_trans(dim, dim);
    combined_trans.AddMatMat(1.0, Ptrans, kNoTrans,
                             transform_mod, kNoTrans, 0.0);
    Matrix<double> combined_trans_inv(combined_trans);  // ... and its inverse.
    combined_trans_inv.Invert();
  
    // Wmod and Bmod are as Wproj2 and Bproj2 but taken back into the original
    // iVector space.
    SpMatrix<double> Wmod(dim), Bmod(dim);
    Wmod.AddMat2Sp(1.0, combined_trans_inv, kNoTrans, Wproj2mod, 0.0);
    Bmod.AddMat2Sp(1.0, combined_trans_inv, kNoTrans, Bproj2mod, 0.0);
  
    TpMatrix<double> C(dim);
    // Do Cholesky Wmod = C C^T.  Now if we use C^{-1} as a transform, we have
    // C^{-1} W C^{-T} = I, so it makes the within-class covar unit.
    C.Cholesky(Wmod);
    TpMatrix<double> Cinv(C);
    Cinv.Invert();
  
    // Bmod_proj is Bmod projected by Cinv.
    SpMatrix<double> Bmod_proj(dim);
    Bmod_proj.AddTp2Sp(1.0, Cinv, kNoTrans, Bmod, 0.0);
    Vector<double> psi_new(dim);
    Matrix<double> Q(dim, dim);
    // Do symmetric eigenvalue decomposition of Bmod_proj, so
    // Bmod_proj = Q diag(psi_new) Q^T
    Bmod_proj.Eig(&psi_new, &Q);
    SortSvd(&psi_new, &Q);
    // This means that if we use Q^T as a transform, then Q^T Bmod_proj Q =
    // diag(psi_new), hence Q^T diagonalizes Bmod_proj (while leaving the
    // within-covar unit).
    // The final transform we want, that projects from our original
    // space to our newly normalized space, is:
    // first Cinv, then Q^T, i.e. the
    // matrix Q^T Cinv.
    Matrix<double> final_transform(dim, dim);
    final_transform.AddMatTp(1.0, Q, kTrans, Cinv, kNoTrans, 0.0);
  
    KALDI_LOG << "Old diagonal of between-class covar was: "
              << plda->psi_ << ", new diagonal is "
              << psi_new;
    plda->transform_.CopyFromMat(final_transform);
    plda->psi_.CopyFromVec(psi_new);
  }
  
  } // namespace kaldi