void StackedRuizEquil
( DistSparseMatrix<Field>& A,
DistSparseMatrix<Field>& B,
DistMultiVec<Base<Field>>& dRowA,
DistMultiVec<Base<Field>>& dRowB,
DistMultiVec<Base<Field>>& dCol,
bool progress )
{
EL_DEBUG_CSE
typedef Base<Field> Real;
const Int mA = A.Height();
const Int mB = B.Height();
const Int n = A.Width();
mpi::Comm comm = A.Comm();
dRowA.SetComm( comm );
dRowB.SetComm( comm );
dCol.SetComm( comm );
Ones( dRowA, mA, 1 );
Ones( dRowB, mB, 1 );
Ones( dCol, n, 1 );
// TODO(poulson): Expose to control structure
// For, simply hard-code a small number of iterations
const Int maxIter = 4;
DistMultiVec<Real> scales(comm), maxAbsValsB(comm);
auto& scalesLoc = scales.Matrix();
auto& maxAbsValsBLoc = maxAbsValsB.Matrix();
const Int localHeight = scalesLoc.Height();
const Int indent = PushIndent();
for( Int iter=0; iter<maxIter; ++iter )
{
// Rescale the columns
// -------------------
ColumnMaxNorms( A, scales );
ColumnMaxNorms( B, maxAbsValsB );
for( Int jLoc=0; jLoc<localHeight; ++jLoc )
scalesLoc(jLoc) = Max(scalesLoc(jLoc),maxAbsValsBLoc(jLoc));
EntrywiseMap( scales, MakeFunction(DampScaling<Real>) );
DiagonalScale( LEFT, NORMAL, scales, dCol );
DiagonalSolve( RIGHT, NORMAL, scales, A );
DiagonalSolve( RIGHT, NORMAL, scales, B );
// Rescale the rows
// ----------------
RowMaxNorms( A, scales );
EntrywiseMap( scales, MakeFunction(DampScaling<Real>) );
DiagonalScale( LEFT, NORMAL, scales, dRowA );
DiagonalSolve( LEFT, NORMAL, scales, A );
RowMaxNorms( B, scales );
EntrywiseMap( scales, MakeFunction(DampScaling<Real>) );
DiagonalScale( LEFT, NORMAL, scales, dRowB );
DiagonalSolve( LEFT, NORMAL, scales, B );
}
SetIndent( indent );
}
void SymmetricRuizEquil
( DistSparseMatrix<F>& A,
DistMultiVec<Base<F>>& d,
Int maxIter, bool progress )
{
DEBUG_CSE
typedef Base<F> Real;
const Int n = A.Height();
mpi::Comm comm = A.Comm();
d.SetComm( comm );
Ones( d, n, 1 );
DistMultiVec<Real> scales(comm);
const Int indent = PushIndent();
for( Int iter=0; iter<maxIter; ++iter )
{
// Rescale the columns (and rows)
// ------------------------------
ColumnMaxNorms( A, scales );
EntrywiseMap( scales, function<Real(Real)>(DampScaling<Real>) );
EntrywiseMap( scales, function<Real(Real)>(SquareRootScaling<Real>) );
DiagonalScale( LEFT, NORMAL, scales, d );
SymmetricDiagonalSolve( scales, A );
}
SetIndent( indent );
}
void RowMaxNorms( const DistMultiVec<F>& A, DistMultiVec<Base<F>>& norms )
{
DEBUG_CSE
norms.SetComm( A.Comm() );
norms.Resize( A.Height(), 1 );
RowMaxNorms( A.LockedMatrix(), norms.Matrix() );
}
void SOCSquareRoot
( const DistMultiVec<Real>& x,
DistMultiVec<Real>& xRoot,
const DistMultiVec<Int>& orders,
const DistMultiVec<Int>& firstInds, Int cutoff )
{
DEBUG_ONLY(CSE cse("SOCSquareRoot"))
DistMultiVec<Real> d(x.Comm());
SOCDets( x, d, orders, firstInds );
ConeBroadcast( d, orders, firstInds );
auto roots = x;
ConeBroadcast( roots, orders, firstInds );
const Int localHeight = x.LocalHeight();
xRoot.SetComm( x.Comm() );
Zeros( xRoot, x.Height(), 1 );
for( Int iLoc=0; iLoc<localHeight; ++iLoc )
{
const Int i = x.GlobalRow(iLoc);
const Real x0 = roots.GetLocal(iLoc,0);
const Real det = d.GetLocal(iLoc,0);
const Real eta0 = Sqrt(x0+Sqrt(det))/Sqrt(Real(2));
if( i == firstInds.GetLocal(iLoc,0) )
xRoot.SetLocal( iLoc, 0, eta0 );
else
xRoot.SetLocal( iLoc, 0, x.GetLocal(iLoc,0)/(2*eta0) );
}
}
void EntrywiseMap
( const DistMultiVec<S>& A,
DistMultiVec<T>& B,
function<T(S)> func )
{
DEBUG_CSE
B.SetComm( A.Comm() );
B.Resize( A.Height(), A.Width() );
EntrywiseMap( A.LockedMatrix(), B.Matrix(), func );
}
void RuizEquil
( DistSparseMatrix<Field>& A,
DistMultiVec<Base<Field>>& dRow,
DistMultiVec<Base<Field>>& dCol,
bool progress )
{
EL_DEBUG_CSE
typedef Base<Field> Real;
const Int m = A.Height();
const Int n = A.Width();
mpi::Comm comm = A.Comm();
dRow.SetComm( comm );
dCol.SetComm( comm );
Ones( dRow, m, 1 );
Ones( dCol, n, 1 );
// TODO(poulson): Expose to control structure
// For, simply hard-code a small number of iterations
const Int maxIter = 4;
DistMultiVec<Real> scales(comm);
const Int indent = PushIndent();
for( Int iter=0; iter<maxIter; ++iter )
{
// Rescale the columns
// -------------------
ColumnMaxNorms( A, scales );
EntrywiseMap( scales, MakeFunction(DampScaling<Real>) );
DiagonalScale( LEFT, NORMAL, scales, dCol );
DiagonalSolve( RIGHT, NORMAL, scales, A );
// Rescale the rows
// ----------------
RowMaxNorms( A, scales );
EntrywiseMap( scales, MakeFunction(DampScaling<Real>) );
DiagonalScale( LEFT, NORMAL, scales, dRow );
DiagonalSolve( LEFT, NORMAL, scales, A );
}
SetIndent( indent );
}
void NesterovTodd
( const DistMultiVec<Real>& s,
const DistMultiVec<Real>& z,
DistMultiVec<Real>& w )
{
DEBUG_CSE
w.SetComm( s.Comm() );
w.Resize( s.Height(), 1 );
const Real* sBuf = s.LockedMatrix().LockedBuffer();
const Real* zBuf = z.LockedMatrix().LockedBuffer();
Real* wBuf = w.Matrix().Buffer();
const Int localHeight = w.LocalHeight();
for( Int iLoc=0; iLoc<localHeight; ++iLoc )
wBuf[iLoc] = Sqrt(sBuf[iLoc]/zBuf[iLoc]);
}
void RowMaxNorms( const DistSparseMatrix<F>& A, DistMultiVec<Base<F>>& norms )
{
DEBUG_CSE
typedef Base<F> Real;
const Int localHeight = A.LocalHeight();
const F* valBuf = A.LockedValueBuffer();
const Int* offsetBuf = A.LockedOffsetBuffer();
norms.SetComm( A.Comm() );
norms.Resize( A.Height(), 1 );
auto& normsLoc = norms.Matrix();
for( Int iLoc=0; iLoc<localHeight; ++iLoc )
{
Real rowMax = 0;
const Int offset = offsetBuf[iLoc];
const Int numConn = offsetBuf[iLoc+1] - offset;
for( Int e=offset; e<offset+numConn; ++e )
rowMax = Max(rowMax,Abs(valBuf[e]));
normsLoc(iLoc) = rowMax;
}
}
void RowTwoNorms( const DistSparseMatrix<F>& A, DistMultiVec<Base<F>>& norms )
{
DEBUG_CSE
typedef Base<F> Real;
const Int localHeight = A.LocalHeight();
const F* valBuf = A.LockedValueBuffer();
const Int* offsetBuf = A.LockedOffsetBuffer();
norms.SetComm( A.Comm() );
norms.Resize( A.Height(), 1 );
auto& normLoc = norms.Matrix();
for( Int iLoc=0; iLoc<localHeight; ++iLoc )
{
Real scale = 0;
Real scaledSquare = 1;
const Int offset = offsetBuf[iLoc];
const Int numConn = offsetBuf[iLoc+1] - offset;
for( Int e=offset; e<offset+numConn; ++e )
UpdateScaledSquare( valBuf[e], scale, scaledSquare );
normLoc(iLoc) = scale*Sqrt(scaledSquare);
}
}
