// @(#)root/matrix:$Name:  $:$Id: TDecompChol.cxx,v 1.17 2005/09/03 13:12:35 brun Exp $
// Authors: Fons Rademakers, Eddy Offermann  Dec 2003

/*************************************************************************
 * Copyright (C) 1995-2000, Rene Brun and Fons Rademakers.               *
 * All rights reserved.                                                  *
 *                                                                       *
 * For the licensing terms see $ROOTSYS/LICENSE.                         *
 * For the list of contributors see $ROOTSYS/README/CREDITS.             *
 *************************************************************************/

///////////////////////////////////////////////////////////////////////////
//                                                                       //
// Cholesky Decomposition class                                          //
//                                                                       //
// Decompose a symmetric, positive definite matrix A = U^T * U           //
//                                                                       //
// where U is a upper triangular matrix                                  //
//                                                                       //
// The decomposition fails if a diagonal element of fU is <= 0, the      //
// matrix is not positive negative . The matrix fU is made invalid .     // 
//                                                                       //
// fU has the same index range as A .                                    //
//                                                                       //
///////////////////////////////////////////////////////////////////////////

#include "TDecompChol.h"

ClassImp(TDecompChol)

//______________________________________________________________________________
 TDecompChol::TDecompChol(Int_t nrows)
{
  fU.ResizeTo(nrows,nrows);
}

//______________________________________________________________________________
 TDecompChol::TDecompChol(Int_t row_lwb,Int_t row_upb)
{
  const Int_t nrows = row_upb-row_lwb+1;
  fRowLwb = row_lwb;
  fColLwb = row_lwb;
  fU.ResizeTo(row_lwb,row_lwb+nrows-1,row_lwb,row_lwb+nrows-1);
}

//______________________________________________________________________________
 TDecompChol::TDecompChol(const TMatrixDSym &a,Double_t tol)
{
  Assert(a.IsValid());

  SetBit(kMatrixSet);
  fCondition = a.Norm1();
  fTol = a.GetTol();
  if (tol > 0)
    fTol = tol;

  fRowLwb = a.GetRowLwb();
  fColLwb = a.GetColLwb();
  fU.ResizeTo(a);
  fU = a;
}

//______________________________________________________________________________
 TDecompChol::TDecompChol(const TMatrixD &a,Double_t tol)
{
  Assert(a.IsValid());

  if (a.GetNrows() != a.GetNcols() || a.GetRowLwb() != a.GetColLwb()) {
    Error("TDecompChol(const TMatrixD &","matrix should be square");
    return;
  }

  SetBit(kMatrixSet);
  fCondition = a.Norm1();
  fTol = a.GetTol();
  if (tol > 0)
    fTol = tol;

  fRowLwb = a.GetRowLwb();
  fColLwb = a.GetColLwb();
  fU.ResizeTo(a);
  fU = a;
}

//______________________________________________________________________________
 TDecompChol::TDecompChol(const TDecompChol &another) : TDecompBase(another)
{
  *this = another;
}

//______________________________________________________________________________
 Bool_t TDecompChol::Decompose()
{
  if ( !TestBit(kMatrixSet) )
    return kFALSE;

  Int_t i,j,icol,irow;
  const Int_t     n  = fU.GetNrows();
        Double_t *pU = fU.GetMatrixArray();
  for (icol = 0; icol < n; icol++) {
    const Int_t rowOff = icol*n;
   
    //Compute fU(j,j) and test for non-positive-definiteness.
    Double_t ujj = pU[rowOff+icol];
    for (irow = 0; irow < icol; irow++) {
      const Int_t pos_ij = irow*n+icol;
      ujj -= pU[pos_ij]*pU[pos_ij];
    }
    if (ujj <= 0) {
      Error("Decompose()","matrix not positive definite");
      return kFALSE;
    }
    ujj = TMath::Sqrt(ujj);
    pU[rowOff+icol] = ujj;

    if (icol < n-1) {
      for (j = icol+1; j < n; j++) {
        for (i = 0; i < icol; i++) {
          const Int_t rowOff2 = i*n;
          pU[rowOff+j] -= pU[rowOff2+j]*pU[rowOff2+icol];
        }
      }
      for (j = icol+1; j < n; j++)
        pU[rowOff+j] /= ujj;
    }
  }

  for (irow = 0; irow < n; irow++) {
    const Int_t rowOff = irow*n;
    for (icol = 0; icol < irow; icol++)
      pU[rowOff+icol] = 0.;
  }

  SetBit(kDecomposed);

  return kTRUE;
}

//______________________________________________________________________________
 const TMatrixDSym TDecompChol::GetMatrix()
{
// Reconstruct the original matrix using the decomposition parts

  if (TestBit(kSingular)) {
    TMatrixDSym tmp; tmp.Invalidate();
    return tmp;
  }
  if ( !TestBit(kDecomposed) ) {
    if (!Decompose()) {
      TMatrixDSym tmp; tmp.Invalidate();
      return tmp;
    }
  }

  return TMatrixDSym(TMatrixDSym::kAtA,fU);
}

//______________________________________________________________________________
 void TDecompChol::SetMatrix(const TMatrixDSym &a)
{
  Assert(a.IsValid());
  
  ResetStatus();
  if (a.GetNrows() != a.GetNcols() || a.GetRowLwb() != a.GetColLwb()) {
    Error("SetMatrix(const TMatrixDSym &","matrix should be square");
    return;
  } 
  
  SetBit(kMatrixSet);
  fCondition = -1.0;
    
  fRowLwb = a.GetRowLwb();
  fColLwb = a.GetColLwb();
  fU.ResizeTo(a);
  fU = a;
}

//______________________________________________________________________________
 Bool_t TDecompChol::Solve(TVectorD &b)
{
// Solve equations Ax=b assuming A has been factored by Cholesky. The factor U is
// assumed to be in upper triang of fU. fTol is used to determine if diagonal
// element is zero. The solution is returned in b.

  Assert(b.IsValid());
  if (TestBit(kSingular)) {
    b.Invalidate();
    return kFALSE;
  }
  if ( !TestBit(kDecomposed) ) {
    if (!Decompose()) {
      b.Invalidate();
      return kFALSE;
    }
  }

  if (fU.GetNrows() != b.GetNrows() || fU.GetRowLwb() != b.GetLwb()) {
    Error("Solve(TVectorD &","vector and matrix incompatible");
    b.Invalidate();
    return kFALSE;
  }

  const Int_t n = fU.GetNrows();

  const Double_t *pU = fU.GetMatrixArray();
        Double_t *pb = b.GetMatrixArray();

  Int_t i;
  // step 1: Forward substitution on U^T
  for (i = 0; i < n; i++) {
    const Int_t off_i = i*n;
    if (pU[off_i+i] < fTol)
    {
      Error("Solve(TVectorD &b)","u[%d,%d]=%.4e < %.4e",i,i,pU[off_i+i],fTol);
      b.Invalidate();
      return kFALSE;
    }
    Double_t r = pb[i];
    for (Int_t j = 0; j < i; j++) {
      const Int_t off_j = j*n;
      r -= pU[off_j+i]*pb[j];
    }
    pb[i] = r/pU[off_i+i];
  }

 // step 2: Backward substitution on U
  for (i = n-1; i >= 0; i--) {
    const Int_t off_i = i*n;
    Double_t r = pb[i];
    for (Int_t j = i+1; j < n; j++)
      r -= pU[off_i+j]*pb[j];
    pb[i] = r/pU[off_i+i];
  }

  return kTRUE;
}

//______________________________________________________________________________
 Bool_t TDecompChol::Solve(TMatrixDColumn &cb)
{ 
  TMatrixDBase *b = const_cast<TMatrixDBase *>(cb.GetMatrix());
  Assert(b->IsValid());
  if (TestBit(kSingular)) {
    b->Invalidate();
    return kFALSE;
  }
  if ( !TestBit(kDecomposed) ) {
    if (!Decompose()) {
      b->Invalidate();
      return kFALSE;
    }
  }

  if (fU.GetNrows() != b->GetNrows() || fU.GetRowLwb() != b->GetRowLwb())
  { 
    Error("Solve(TMatrixDColumn &cb","vector and matrix incompatible");
    b->Invalidate();
    return kFALSE; 
  }
      
  const Int_t n = fU.GetNrows();
    
  const Double_t *pU  = fU.GetMatrixArray();
        Double_t *pcb = cb.GetPtr();
  const Int_t     inc = cb.GetInc(); 
  
  Int_t i;
  // step 1: Forward substitution U^T
  for (i = 0; i < n; i++) { 
    const Int_t off_i  = i*n;
    const Int_t off_i2 = i*inc;
    if (pU[off_i+i] < fTol)
    {
      Error("Solve(TMatrixDColumn &cb)","u[%d,%d]=%.4e < %.4e",i,i,pU[off_i+i],fTol);
      b->Invalidate();
      return kFALSE;
    }
    Double_t r = pcb[off_i2];
    for (Int_t j = 0; j < i; j++) {
      const Int_t off_j = j*n;
      r -= pU[off_j+i]*pcb[j*inc];
    }
    pcb[off_i2] = r/pU[off_i+i];
  }

  // step 2: Backward substitution U
  for (i = n-1; i >= 0; i--) {
    const Int_t off_i  = i*n;
    const Int_t off_i2 = i*inc;
    Double_t r = pcb[off_i2];
    for (Int_t j = i+1; j < n; j++)
      r -= pU[off_i+j]*pcb[j*inc];
    pcb[off_i2] = r/pU[off_i+i];
  }

  return kTRUE;
}

//______________________________________________________________________________
 void TDecompChol::Det(Double_t &d1,Double_t &d2)
{
  // determinant is square of diagProd of cholesky factor

  if ( !TestBit(kDetermined) ) {
    if ( !TestBit(kDecomposed) )
      Decompose();
    TDecompBase::Det(d1,d2);
    // square det as calculated by above
    fDet1 *= fDet1;
    fDet2 += fDet2;
    SetBit(kDetermined);
  }
  d1 = fDet1;
  d2 = fDet2;
}

//______________________________________________________________________________
 void TDecompChol::Invert(TMatrixDSym &inv)
{
  // For a symmetric matrix A(m,m), its inverse A_inv(m,m) is returned .

  if (inv.GetNrows() != GetNrows() || inv.GetRowLwb() != GetRowLwb()) {
    Error("Invert(TMatrixDSym &","Input matrix has wrong shape");
    inv.Invalidate();
    return;
  }

  inv.UnitMatrix();

  const Int_t colLwb = inv.GetColLwb();
  const Int_t colUpb = inv.GetColUpb();
  Bool_t status = kTRUE;
  for (Int_t icol = colLwb; icol <= colUpb && status; icol++) {
    TMatrixDColumn b(inv,icol);
    status &= Solve(b);
  }
  
  if (!status)
    inv.Invalidate();
}

//______________________________________________________________________________
 TMatrixDSym TDecompChol::Invert()
{ 
  // For a symmetric matrix A(m,m), its inverse A_inv(m,m) is returned .

  const Int_t rowLwb = GetRowLwb();
  const Int_t rowUpb = rowLwb+GetNrows()-1;

  TMatrixDSym inv(rowLwb,rowUpb);
  inv.UnitMatrix();
  Invert(inv);

  return inv;
}

//______________________________________________________________________________
 void TDecompChol::Print(Option_t *opt) const
{
  TDecompBase::Print(opt);
  fU.Print("fU");
}

//______________________________________________________________________________
TDecompChol &TDecompChol::operator=(const TDecompChol &source)
{ 
  if (this != &source) {
    TDecompBase::operator=(source);
    fU.ResizeTo(source.fU);
    fU = source.fU;
  }
  return *this;
}     

//______________________________________________________________________________
TVectorD NormalEqn(const TMatrixD &A,const TVectorD &b)
{
  // Solve min {(A . x - b)^T (A . x - b)} for vector x where
  //   A : (m x n) matrix, m >= n
  //   b : (m)     vector
  //   x : (n)     vector

  TDecompChol ch(TMatrixDSym(TMatrixDBase::kAtA,A));
  Bool_t ok;
  return ch.Solve(TMatrixD(TMatrixDBase::kTransposed,A)*b,ok);
}

//______________________________________________________________________________
TVectorD NormalEqn(const TMatrixD &A,const TVectorD &b,const TVectorD &std)
{
  // Solve min {(A . x - b)^T W (A . x - b)} for vector x where
  //   A : (m x n) matrix, m >= n
  //   b : (m)     vector
  //   x : (n)     vector
  //   W : (m x m) weight matrix with W(i,j) = 1/std(i)^2  for i == j
  //                                         = 0           fir i != j

  if (!AreCompatible(b,std)) {
    ::Error("NormalEqn","vectors b and std are not compatible");
    TVectorD tmp; tmp.Invalidate();
    return tmp;
  }

  TMatrixD mAw = A;
  TVectorD mBw = b;
  for (Int_t irow = 0; irow < A.GetNrows(); irow++) {
    TMatrixDRow(mAw,irow) *= 1/std(irow);
    mBw(irow) /= std(irow);
  }
  TDecompChol ch(TMatrixDSym(TMatrixDBase::kAtA,mAw));
  Bool_t ok;
  return ch.Solve(TMatrixD(TMatrixDBase::kTransposed,mAw)*mBw,ok);
}

//______________________________________________________________________________
TMatrixD NormalEqn(const TMatrixD &A,const TMatrixD &B)
{
  // Solve min {(A . X_j - B_j)^T (A . X_j - B_j)} for each column j in
  // B and X
  //   A : (m x n ) matrix, m >= n
  //   B : (m x nb) matrix, nb >= 1
  //  mX : (n x nb) matrix

  TDecompChol ch(TMatrixDSym(TMatrixDBase::kAtA,A));
  TMatrixD mX(A,TMatrixDBase::kTransposeMult,B);
  ch.MultiSolve(mX);
  return mX;
}

//______________________________________________________________________________
TMatrixD NormalEqn(const TMatrixD &A,const TMatrixD &B,const TVectorD &std)
{
  // Solve min {(A . X_j - B_j)^T W (A . X_j - B_j)} for each column j in
  // B and X
  //   A : (m x n ) matrix, m >= n
  //   B : (m x nb) matrix, nb >= 1
  //  mX : (n x nb) matrix
  //   W : (m x m) weight matrix with W(i,j) = 1/std(i)^2  for i == j
  //                                         = 0           fir i != j

  TMatrixD mAw = A;
  TMatrixD mBw = B;
  for (Int_t irow = 0; irow < A.GetNrows(); irow++) {
    TMatrixDRow(mAw,irow) *= 1/std(irow);
    TMatrixDRow(mBw,irow) *= 1/std(irow);
  }

  TDecompChol ch(TMatrixDSym(TMatrixDBase::kAtA,mAw));
  TMatrixD mX(mAw,TMatrixDBase::kTransposeMult,mBw);
  ch.MultiSolve(mX);
  return mX;
}