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 SOCSquareRoot
( const Matrix<Real>& x,
Matrix<Real>& xRoot,
const Matrix<Int>& orders,
const Matrix<Int>& firstInds )
{
DEBUG_ONLY(CSE cse("SOCSquareRoot"))
Matrix<Real> d;
SOCDets( x, d, orders, firstInds );
ConeBroadcast( d, orders, firstInds );
const Int height = x.Height();
Zeros( xRoot, height, 1 );
for( Int i=0; i<height; )
{
const Int order = orders.Get(i,0);
const Int firstInd = firstInds.Get(i,0);
if( i != firstInd )
LogicError("Inconsistency in orders and firstInds");
const Real eta0 = Sqrt(x.Get(i,0)+Sqrt(d.Get(i,0)))/Sqrt(Real(2));
xRoot.Set( i, 0, eta0 );
for( Int k=1; k<order; ++k )
xRoot.Set( i+k, 0, x.Get(i+k,0)/(2*eta0) );
i += order;
}
}
void SOCSquareRoot
( const ElementalMatrix<Real>& xPre,
ElementalMatrix<Real>& xRootPre,
const ElementalMatrix<Int>& ordersPre,
const ElementalMatrix<Int>& firstIndsPre,
Int cutoff )
{
DEBUG_ONLY(CSE cse("SOCSquareRoot"))
AssertSameGrids( xPre, xRootPre, ordersPre, firstIndsPre );
ProxyCtrl ctrl;
ctrl.colConstrain = true;
ctrl.colAlign = 0;
auto xPtr = ReadProxy<Real,VC,STAR>(&xPre,ctrl);
auto xRootPtr = WriteProxy<Real,VC,STAR>(&xRootPre,ctrl);
auto ordersPtr = ReadProxy<Int,VC,STAR>(&ordersPre,ctrl);
auto firstIndsPtr = ReadProxy<Int,VC,STAR>(&firstIndsPre,ctrl);
auto& x = *xPtr;
auto& xRoot = *xRootPtr;
auto& orders = *ordersPtr;
auto& firstInds = *firstIndsPtr;
DistMatrix<Real,VC,STAR> d(x.Grid());
SOCDets( x, d, orders, firstInds );
ConeBroadcast( d, orders, firstInds );
auto roots = x;
ConeBroadcast( roots, orders, firstInds );
const Int localHeight = x.LocalHeight();
xRoot.SetGrid( x.Grid() );
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 SOCApply
( const Matrix<Real>& x,
const Matrix<Real>& y,
Matrix<Real>& z,
const Matrix<Int>& orders,
const Matrix<Int>& firstInds )
{
DEBUG_ONLY(CSE cse("SOCApply"))
SOCDots( x, y, z, orders, firstInds );
auto xRoots = x;
auto yRoots = y;
ConeBroadcast( xRoots, orders, firstInds );
ConeBroadcast( yRoots, orders, firstInds );
const Int height = x.Height();
for( Int i=0; i<height; ++i )
if( i != firstInds.Get(i,0) )
z.Update( i, 0, xRoots.Get(i,0)*y.Get(i,0) +
yRoots.Get(i,0)*x.Get(i,0) );
}
void SOCApply
( const AbstractDistMatrix<Real>& xPre,
const AbstractDistMatrix<Real>& yPre,
AbstractDistMatrix<Real>& zPre,
const AbstractDistMatrix<Int>& ordersPre,
const AbstractDistMatrix<Int>& firstIndsPre,
Int cutoff )
{
DEBUG_ONLY(CSE cse("SOCApply"))
AssertSameGrids( xPre, yPre, zPre, ordersPre, firstIndsPre );
ProxyCtrl ctrl;
ctrl.colConstrain = true;
ctrl.colAlign = 0;
auto xPtr = ReadProxy<Real,VC,STAR>(&xPre,ctrl);
auto yPtr = ReadProxy<Real,VC,STAR>(&yPre,ctrl);
auto zPtr = WriteProxy<Real,VC,STAR>(&zPre,ctrl);
auto ordersPtr = ReadProxy<Int,VC,STAR>(&ordersPre,ctrl);
auto firstIndsPtr = ReadProxy<Int,VC,STAR>(&firstIndsPre,ctrl);
auto& x = *xPtr;
auto& y = *yPtr;
auto& z = *zPtr;
auto& orders = *ordersPtr;
auto& firstInds = *firstIndsPtr;
SOCDots( x, y, z, orders, firstInds );
auto xRoots = x;
auto yRoots = y;
ConeBroadcast( xRoots, orders, firstInds );
ConeBroadcast( yRoots, orders, firstInds );
const Int localHeight = x.LocalHeight();
for( Int iLoc=0; iLoc<localHeight; ++iLoc )
{
const Int i = x.GlobalRow(iLoc);
const Int firstInd = firstInds.GetLocal(iLoc,0);
if( i != firstInd )
z.UpdateLocal
( iLoc, 0, xRoots.GetLocal(iLoc,0)*y.GetLocal(iLoc,0) +
yRoots.GetLocal(iLoc,0)*x.GetLocal(iLoc,0) );
}
}
void SOCApply
( const DistMultiVec<Real>& x,
const DistMultiVec<Real>& y,
DistMultiVec<Real>& z,
const DistMultiVec<Int>& orders,
const DistMultiVec<Int>& firstInds, Int cutoff )
{
DEBUG_ONLY(CSE cse("SOCApply"))
SOCDots( x, y, z, orders, firstInds );
auto xRoots = x;
auto yRoots = y;
ConeBroadcast( xRoots, orders, firstInds );
ConeBroadcast( yRoots, orders, firstInds );
const Int localHeight = x.LocalHeight();
for( Int iLoc=0; iLoc<localHeight; ++iLoc )
{
const Int i = x.GlobalRow(iLoc);
const Int firstInd = firstInds.GetLocal(iLoc,0);
if( i != firstInd )
z.UpdateLocal
( iLoc, 0, xRoots.GetLocal(iLoc,0)*y.GetLocal(iLoc,0) +
yRoots.GetLocal(iLoc,0)*x.GetLocal(iLoc,0) );
}
}
