wehub-resource-sync
336 lines
13 KiB
C++
336 lines
13 KiB
C++
/* ******************************************************************************
|
|
*
|
|
*
|
|
* This program and the accompanying materials are made available under the
|
|
* terms of the Apache License, Version 2.0 which is available at
|
|
* https://www.apache.org/licenses/LICENSE-2.0.
|
|
*
|
|
* See the NOTICE file distributed with this work for additional
|
|
* information regarding copyright ownership.
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
* License for the specific language governing permissions and limitations
|
|
* under the License.
|
|
*
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
******************************************************************************/
|
|
|
|
//
|
|
// @author Yurii Shyrma (iuriish@yahoo.com)
|
|
//
|
|
#include <helpers/EigenValsAndVecs.h>
|
|
#include <helpers/FullPivLU.h>
|
|
#include <helpers/HessenbergAndSchur.h>
|
|
#include <helpers/MmulHelper.h>
|
|
#include <helpers/Sqrtm.h>
|
|
#include <ops/declarable/helpers/lup.h>
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void sqrtmQuasiTrianDiag(NDArray& matrixT, NDArray& sqrtT) {
|
|
const int rows = matrixT.sizeAt(0);
|
|
|
|
for (int i = 0; i < rows; i++) {
|
|
if (i == rows - 1 || matrixT.t<T>(i + 1, i) == (T)0) {
|
|
const auto elemT = matrixT.t<T>(i, i);
|
|
if (elemT < (T)0)
|
|
THROW_EXCEPTION(
|
|
"ops::helpers::Sqrtm::sqrtmQuasiTrianDiag: can't take sqrt of negative diagonal element of T matrix !");
|
|
sqrtT.r<T>(i, i) = math::sd_sqrt<T, T>(elemT);
|
|
} else {
|
|
NDArray *esViewPtr = matrixT({i, i + 2, i, i + 2}, true);
|
|
EigenValsAndVecs<T> es(*esViewPtr); // es._Vecs {2,2,2}, es._Vals{2,2}
|
|
delete esViewPtr;
|
|
|
|
NDArray& vecs = es._Vecs;
|
|
NDArray& vals = es._Vals;
|
|
|
|
const T& vecsReal00 = vecs.t<T>(0, 0, 0);
|
|
const T& vecsImag00 = vecs.t<T>(0, 0, 1);
|
|
const T& vecsReal01 = vecs.t<T>(0, 1, 0);
|
|
const T& vecsImag01 = vecs.t<T>(0, 1, 1);
|
|
const T& vecsReal10 = vecs.t<T>(1, 0, 0);
|
|
const T& vecsImag10 = vecs.t<T>(1, 0, 1);
|
|
const T& vecsReal11 = vecs.t<T>(1, 1, 0);
|
|
const T& vecsImag11 = vecs.t<T>(1, 1, 1);
|
|
|
|
// es.eigenvalues().cwiseSqrt().asDiagonal()
|
|
T eigenValsSqrt[2][2];
|
|
eigenValsSqrt[0][0] = vals.t<T>(0, 0);
|
|
eigenValsSqrt[0][1] = vals.t<T>(0, 1);
|
|
eigenValsSqrt[1][0] = vals.t<T>(1, 0);
|
|
eigenValsSqrt[1][1] = vals.t<T>(1, 1);
|
|
EigenValsAndVecs<T>::sqrtComplexNum(eigenValsSqrt[0][0], eigenValsSqrt[0][1]);
|
|
EigenValsAndVecs<T>::sqrtComplexNum(eigenValsSqrt[1][0], eigenValsSqrt[1][1]);
|
|
|
|
// es.eigenvectors() * es.eigenvalues().cwiseSqrt().asDiagonal()
|
|
T vecsElem[2][2][2];
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsReal00, vecsImag00, eigenValsSqrt[0][0], eigenValsSqrt[0][1],
|
|
vecsElem[0][0][0], vecsElem[0][0][1]);
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsReal01, vecsImag01, eigenValsSqrt[1][0], eigenValsSqrt[1][1],
|
|
vecsElem[0][1][0], vecsElem[0][1][1]);
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsReal10, vecsImag10, eigenValsSqrt[0][0], eigenValsSqrt[0][1],
|
|
vecsElem[1][0][0], vecsElem[1][0][1]);
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsReal11, vecsImag11, eigenValsSqrt[1][0], eigenValsSqrt[1][1],
|
|
vecsElem[1][1][0], vecsElem[1][1][1]);
|
|
|
|
// es.eigenvectors().inverse()
|
|
T vecsElemInv[2][2][2];
|
|
|
|
T tempReal, tempImag, divisorReal, divisorImag;
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsReal00, vecsImag00, vecsReal11, vecsImag11, divisorReal,
|
|
divisorImag);
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsReal01, vecsImag01, vecsReal10, vecsImag10, tempReal, tempImag);
|
|
divisorReal -= tempReal;
|
|
divisorImag -= tempImag;
|
|
|
|
EigenValsAndVecs<T>::divideComplexNums(vecsReal11, vecsImag11, divisorReal, divisorImag, vecsElemInv[0][0][0],
|
|
vecsElemInv[0][0][1]);
|
|
EigenValsAndVecs<T>::divideComplexNums(-vecsReal01, -vecsImag01, divisorReal, divisorImag, vecsElemInv[0][1][0],
|
|
vecsElemInv[0][1][1]);
|
|
EigenValsAndVecs<T>::divideComplexNums(-vecsReal10, -vecsImag10, divisorReal, divisorImag, vecsElemInv[1][0][0],
|
|
vecsElemInv[1][0][1]);
|
|
EigenValsAndVecs<T>::divideComplexNums(vecsReal00, vecsImag00, divisorReal, divisorImag, vecsElemInv[1][1][0],
|
|
vecsElemInv[1][1][1]);
|
|
|
|
// result
|
|
T result[2][2][2];
|
|
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][0][0], vecsElem[0][0][1], vecsElemInv[0][0][0],
|
|
vecsElemInv[0][0][1], tempReal, tempImag);
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][1][0], vecsElem[0][1][1], vecsElemInv[1][0][0],
|
|
vecsElemInv[1][0][1], result[0][0][0], result[0][0][1]);
|
|
result[0][0][0] += tempReal;
|
|
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][0][0], vecsElem[0][0][1], vecsElemInv[0][1][0],
|
|
vecsElemInv[0][1][1], tempReal, tempImag);
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][1][0], vecsElem[0][1][1], vecsElemInv[1][1][0],
|
|
vecsElemInv[1][1][1], result[0][1][0], result[0][1][1]);
|
|
result[0][1][0] += tempReal;
|
|
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][0][0], vecsElem[1][0][1], vecsElemInv[0][0][0],
|
|
vecsElemInv[0][0][1], tempReal, tempImag);
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][1][0], vecsElem[1][1][1], vecsElemInv[1][0][0],
|
|
vecsElemInv[1][0][1], result[1][0][0], result[1][0][1]);
|
|
result[1][0][0] += tempReal;
|
|
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][0][0], vecsElem[1][0][1], vecsElemInv[0][1][0],
|
|
vecsElemInv[0][1][1], tempReal, tempImag);
|
|
EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][1][0], vecsElem[1][1][1], vecsElemInv[1][1][0],
|
|
vecsElemInv[1][1][1], result[1][1][0], result[1][1][1]);
|
|
result[1][1][0] += tempReal;
|
|
|
|
sqrtT.r<T>(i, i) = result[0][0][0];
|
|
sqrtT.r<T>(i, i + 1) = result[0][1][0];
|
|
sqrtT.r<T>(i + 1, i) = result[1][0][0];
|
|
sqrtT.r<T>(i + 1, i + 1) = result[1][1][0];
|
|
|
|
++i;
|
|
}
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
// all matrices are {2,2} here
|
|
template <typename T>
|
|
static void sqrtmQuasiTrianAuxEq(NDArray& A, NDArray& B, NDArray& C, NDArray& X) {
|
|
std::vector<LongType> tempShape = {4,4};
|
|
NDArray tempMatrix(A.ordering(),tempShape, A.dataType(), A.getContext());
|
|
|
|
tempMatrix.r<T>(0, 0) = A.t<T>(0, 0) + B.t<T>(0, 0);
|
|
tempMatrix.r<T>(1, 1) = A.t<T>(0, 0) + B.t<T>(1, 1);
|
|
tempMatrix.r<T>(2, 2) = A.t<T>(1, 1) + B.t<T>(0, 0);
|
|
tempMatrix.r<T>(3, 3) = A.t<T>(1, 1) + B.t<T>(1, 1);
|
|
tempMatrix.r<T>(0, 1) = B.t<T>(1, 0);
|
|
tempMatrix.r<T>(0, 2) = A.t<T>(0, 1);
|
|
tempMatrix.r<T>(1, 0) = B.t<T>(0, 1);
|
|
tempMatrix.r<T>(1, 3) = A.t<T>(0, 1);
|
|
tempMatrix.r<T>(2, 0) = A.t<T>(1, 0);
|
|
tempMatrix.r<T>(2, 3) = B.t<T>(1, 0);
|
|
tempMatrix.r<T>(3, 1) = A.t<T>(1, 0);
|
|
tempMatrix.r<T>(3, 2) = B.t<T>(0, 1);
|
|
tempMatrix.r<T>(0, 3) = (T)0;
|
|
tempMatrix.r<T>(1, 2) = (T)0;
|
|
tempMatrix.r<T>(2, 1) = (T)0;
|
|
tempMatrix.r<T>(3, 0) = (T)0;
|
|
|
|
std::vector<LongType> resultShape = {4,1};
|
|
NDArray result(A.ordering(), resultShape, A.dataType(), A.getContext());
|
|
result.r<T>(0, 0) = C.t<T>(0, 0);
|
|
result.r<T>(1, 0) = C.t<T>(0, 1);
|
|
result.r<T>(2, 0) = C.t<T>(1, 0);
|
|
result.r<T>(3, 0) = C.t<T>(1, 1);
|
|
|
|
FullPivLU<T>::solve(tempMatrix, result, result);
|
|
|
|
X.r<T>(0, 0) = result.t<T>(0);
|
|
X.r<T>(0, 1) = result.t<T>(1);
|
|
X.r<T>(1, 0) = result.t<T>(2);
|
|
X.r<T>(1, 1) = result.t<T>(3);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void sqrtmQuasiTrianOffDiag(NDArray& matrixT, NDArray& sqrtT) {
|
|
const int rows = matrixT.sizeAt(0);
|
|
|
|
for (int j = 1; j < rows; j++) {
|
|
if (matrixT.t<T>(j, j - 1) != (T)0) continue;
|
|
|
|
for (int i = j - 1; i >= 0; i--) {
|
|
if (i > 0 && matrixT.t<T>(i, i - 1) != (T)0) continue;
|
|
|
|
const bool iBlockIs2x2 = (i < rows - 1) && (matrixT.t<T>(i + 1, i) != (T)0);
|
|
const bool jBlockIs2x2 = (j < rows - 1) && (matrixT.t<T>(j + 1, j) != (T)0);
|
|
|
|
if (iBlockIs2x2 && jBlockIs2x2) {
|
|
NDArray *APtr = sqrtT({i, i + 2, i, i + 2}, true);
|
|
NDArray A = *APtr;
|
|
delete APtr;
|
|
|
|
NDArray *BPtr = sqrtT({j, j + 2, j, j + 2}, true);
|
|
NDArray B = *BPtr;
|
|
delete BPtr;
|
|
|
|
NDArray *XPtr = matrixT({i, i + 2, j, j + 2}, true);
|
|
NDArray X = *XPtr;
|
|
delete XPtr;
|
|
|
|
if (j - i > 2) {
|
|
NDArray *leftPtr = sqrtT({i, i + 2, i + 2, j}, true);
|
|
NDArray *rightPtr = sqrtT({i + 2, j, j, j + 2}, true);
|
|
auto mul = mmul(*leftPtr, *rightPtr);
|
|
X -= *mul;
|
|
delete leftPtr;
|
|
delete rightPtr;
|
|
delete mul;
|
|
}
|
|
|
|
sqrtmQuasiTrianAuxEq<T>(A, B, X, X);
|
|
|
|
sqrtT.syncToDevice();
|
|
NDArray *assignPtr = sqrtT({i, i + 2, j, j + 2}, true);
|
|
assignPtr->assign(&X);
|
|
delete assignPtr;
|
|
} else if (iBlockIs2x2 && !jBlockIs2x2) {
|
|
NDArray *rhsPtr = matrixT({i, i + 2, j, j + 1}, true);
|
|
NDArray rhs = *rhsPtr;
|
|
delete rhsPtr;
|
|
|
|
if (j - i > 2) {
|
|
NDArray *leftPtr = sqrtT({i, i + 2, i + 2, j}, true);
|
|
NDArray *rightPtr = sqrtT({i + 2, j, j, j + 1}, true);
|
|
auto mul = mmul(*leftPtr, *rightPtr);
|
|
rhs -= *mul;
|
|
delete leftPtr;
|
|
delete rightPtr;
|
|
delete mul;
|
|
}
|
|
|
|
std::vector<LongType> aShape = {2,2};
|
|
NDArray A(matrixT.ordering(), aShape, matrixT.dataType(), matrixT.getContext());
|
|
A.r<T>(0, 0) = A.r<T>(1, 1) = sqrtT.t<T>(j, j);
|
|
A.r<T>(0, 1) = A.r<T>(1, 0) = T(0);
|
|
|
|
NDArray *addPtr = sqrtT({i, i + 2, i, i + 2}, true);
|
|
A += *addPtr;
|
|
delete addPtr;
|
|
|
|
FullPivLU<T>::solve(A, rhs, rhs);
|
|
|
|
// sqrtT.syncToDevice();
|
|
NDArray *assignPtr = sqrtT({i, i + 2, j, j + 1}, true);
|
|
assignPtr->assign(&rhs);
|
|
delete assignPtr;
|
|
} else if (!iBlockIs2x2 && jBlockIs2x2) {
|
|
NDArray *rhsPtr = matrixT({i, i + 1, j, j + 2}, true);
|
|
NDArray rhs = *rhsPtr;
|
|
delete rhsPtr;
|
|
|
|
if (j - i > 1) {
|
|
NDArray *leftPtr = sqrtT({i, i + 1, i + 1, j}, true);
|
|
NDArray *rightPtr = sqrtT({i + 1, j, j, j + 2}, true);
|
|
auto mul = mmul(*leftPtr, *rightPtr);
|
|
rhs -= *mul;
|
|
delete leftPtr;
|
|
delete rightPtr;
|
|
delete mul;
|
|
}
|
|
|
|
std::vector<LongType> aShape = {2,2};
|
|
NDArray A(matrixT.ordering(),aShape, matrixT.dataType(), matrixT.getContext());
|
|
A.r<T>(0, 0) = A.r<T>(1, 1) = sqrtT.t<T>(i, i);
|
|
A.r<T>(0, 1) = A.r<T>(1, 0) = T(0);
|
|
|
|
NDArray *addPtr = sqrtT({j, j + 2, j, j + 2}, true);
|
|
NDArray *add = addPtr->transpose();
|
|
delete addPtr;
|
|
A += *add;
|
|
delete add;
|
|
|
|
NDArray *rhsT = rhs.transpose();
|
|
FullPivLU<T>::solve(A, *rhsT, *rhsT);
|
|
|
|
// sqrtT.syncToDevice();
|
|
NDArray *assignPtr = sqrtT({i, i + 1, j, j + 2}, true);
|
|
assignPtr->assign(&rhs);
|
|
delete assignPtr;
|
|
delete rhsT;
|
|
} else if (!iBlockIs2x2 && !jBlockIs2x2) {
|
|
NDArray *leftPtr = sqrtT({i, i + 1, i + 1, j});
|
|
NDArray *rightPtr = sqrtT({i + 1, j, j, j + 1});
|
|
auto mul = mmul(*leftPtr, *rightPtr);
|
|
T temp = mul->t<T>(0); // dot
|
|
delete leftPtr;
|
|
delete rightPtr;
|
|
delete mul;
|
|
|
|
sqrtT.r<T>(i, j) = (matrixT.t<T>(i, j) - temp) / (sqrtT.t<T>(i, i) + sqrtT.t<T>(j, j));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
void Sqrtm<T>::calc(NDArray& in, NDArray& out) {
|
|
if (in.rankOf() != 2 || in.sizeAt(0) != in.sizeAt(1))
|
|
THROW_EXCEPTION("ops::helpers::Sqrtm::calc: input matrix must have rank 2 and be square !");
|
|
if (!out.isSameShape(in))
|
|
THROW_EXCEPTION("ops::helpers::Sqrtm::calc: output matrix must have the same shape as input one!");
|
|
|
|
if (in.lengthOf() == 1) {
|
|
out.r<T>(0) = math::sd_sqrt<T, T>(in.t<T>(0));
|
|
return;
|
|
}
|
|
|
|
Schur<T> schur(in);
|
|
|
|
|
|
NDArray *inULike = in.ulike();
|
|
NDArray sqrtT = *inULike;
|
|
sqrtT.nullify();
|
|
|
|
sqrtmQuasiTrianDiag<T>(*schur.t, sqrtT);
|
|
sqrtmQuasiTrianOffDiag<T>(*schur.t, sqrtT);
|
|
|
|
NDArray *second = schur.u->transpose();
|
|
// out = U * sqrtT * U^T;
|
|
NDArray *temp = mmul(sqrtT, *second);
|
|
MmulHelper::mmul(schur.u, temp, &out);
|
|
delete inULike;
|
|
delete second;
|
|
delete temp;
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( class Sqrtm, , SD_FLOAT_TYPES);
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|