rotation.h

Go to the documentation of this file.

#ifndef HPP_ROTATION_H
#define HPP_ROTATION_H

#include "mpi.h"
#include <hpp/config.h>
#include <hpp/tensor.h>
#include <hpp/mpiUtils.h>
#include <hpp/external/ISOI/grid_generation.h>

namespace hpp
{

// Types of symmetry in Schoenflies notation
enum SymmetryType {
    SYMMETRY_TYPE_NONE,
    SYMMETRY_TYPE_C4
};
    
template <typename T>
struct EulerAngles {
    T alpha = 0;
    T beta = 0;
    T gamma = 0;
    
    HPP_NVCC_ONLY(__host__ __device__) EulerAngles() {;}
    HPP_NVCC_ONLY(__host__ __device__) EulerAngles(T alpha, T beta, T gamma) :
    alpha(alpha), beta(beta), gamma(gamma) {;}
    
    // Getters/setters (mainly intended for Python interface)
    T getAlpha() const {return alpha;}
    T getBeta() const {return beta;}
    T getGamma() const {return gamma;}
    void setAlpha(const T& alpha) {this->alpha = alpha;}
    void setBeta(const T& beta) {this->beta = beta;}
    void setGamma(const T& gamma) {this->gamma = gamma;}
};

template <typename T>
MPI_Datatype getEulerAnglesTypeMPI() {
    MPI_Datatype dtype;
    MPI_Type_contiguous(3, MPIType<T>(), &dtype);
    MPI_Type_commit(&dtype);
    return dtype;
}

template <typename T>
EulerAngles<T> polarToEuler(T theta, T phi) {
    EulerAngles<T> angles;
    // Rotate about z axis until x axis is aligned with y axis
    angles.gamma = M_PI/2;
    
    // Rotate about x axis to get the zenithal angle correct
    angles.beta = M_PI/2-phi;
    
    // Rotate about z axis to get the azimuthal angle correct
    angles.alpha = theta;
    
    // Return
    return angles;
}

template <typename T>
std::ostream& operator<<(std::ostream& out, const EulerAngles<T>& angles)
{
    out << "[";
    out << angles.alpha << ",";
    out << angles.beta << ",";
    out << angles.gamma;
    out << "]";
    return out;
}

template <typename T>
bool operator==(const EulerAngles<T>& l, const EulerAngles<T>& r) {
    if (l.alpha != r.alpha) return false;
    if (l.beta != r.beta) return false;
    if (l.gamma != r.gamma) return false;
    
    // All checks passed
    return true;
}

template <typename T>
bool operator!=(const EulerAngles<T>& l, const EulerAngles<T>& r) {
    return !(l==r);
}

template <typename T>
Tensor2<T> EulerZXZRotationMatrix(T alpha, T beta, T gamma) {
    Tensor2<T> R(3,3);
    T c1 = std::cos(alpha);
    T c2 = std::cos(beta);
    T c3 = std::cos(gamma);
    T s1 = std::sin(alpha);
    T s2 = std::sin(beta);
    T s3 = std::sin(gamma);
    R(0,0) = c1*c3 - c2*s1*s3;
    R(0,1) = -c1*s3 - c2*c3*s1;
    R(0,2) = s1*s2;
    R(1,0) = c3*s1 + c1*c2*s3;
    R(1,1) = c1*c2*c3 - s1*s3;
    R(1,2) = -c1*s2;
    R(2,0) = s2*s3;
    R(2,1) = c3*s2;
    R(2,2) = c2;
    return R;
}

template <typename T>
Tensor2<T> EulerZXZRotationMatrix(EulerAngles<T> angle) {
    return EulerZXZRotationMatrix(angle.alpha, angle.beta, angle.gamma);
}

template <typename T>
EulerAngles<T> getEulerZXZAngles(Tensor2<T> R)
{      
    EulerAngles<T> angle;
    
    // Angle beta
    angle.beta = std::acos(R(2,2));
    
    if (angle.beta > 1e3*std::numeric_limits<T>::epsilon()) {
        // The other 2 angles
        angle.alpha = std::atan2(R(0,2),-R(1,2));
        angle.gamma = std::atan2(R(2,0),R(2,1));
    }
    else {
        // Singular case
        angle.beta = 0.0;
        T alphaPlusGamma = std::atan2(-R(0,1), R(0,0));
        
        // Not uniquely determined, so just pick a combination
        angle.alpha = alphaPlusGamma/2.0;
        angle.gamma =  alphaPlusGamma/2.0;
    }
    
    // Correct the angle ranges if necessary
    if (angle.alpha < 0) angle.alpha += 2*M_PI;
    if (angle.gamma < 0) angle.gamma += 2*M_PI;
    
    // Return
    return angle;
}

// RANDOM ROTATIONS //

// TRANSFORMATIONS
inline std::mt19937 makeMt19937(bool defaultSeed) {
    if (defaultSeed) {
        return std::mt19937();
    }
    else {
        std::random_device rd;
        return std::mt19937(rd());
    }
}

// Arvo 1992 method
template <typename T>
void rotationTensorFrom3UniformRandomsArvo1992(Tensor2<T>& A, T x1, T x2, T x3) {    
    // Static to avoid expensive alloc/dealloc for the tens/hundreds of millions regime
    static std::vector<T> v(3);
    static Tensor2<T> R(3,3);
    static Tensor2<T> H(3,3);
    
    // Random rotation about vertical axis
    R(0,0) = std::cos(2*M_PI*x1);
    R(1,1) = R(0,0);
    R(0,1) = std::sin(2*M_PI*x1);
    R(1,0) = -R(0,1);
    R(2,2) = 1.0;
    
    // Random rotation of vertical axis    
    v[0] = std::cos(2*M_PI*x2)*std::sqrt(x3);
    v[1] = std::sin(2*M_PI*x2)*std::sqrt(x3);
    v[2] = std::sqrt(1-x3);
    
    // Householder reflection
    outerInPlace(v, v, H);
    H *= (T)2.0;
    H(0,0) -= 1.0;
    H(1,1) -= 1.0;
    H(2,2) -= 1.0;
    
    // Final matrix
    A = H*R;
}

template <typename T>
void randomRotationTensorArvo1992(std::mt19937& gen, Tensor2<T>& A) {    
    // Static to avoid expensive alloc/dealloc for the tens/hundreds of millions regime
    static std::vector<T> v(3);
    static Tensor2<T> R(3,3);
    static Tensor2<T> H(3,3);
    
    // double is specified here, as if it is not, the values out of the
    // distribution are different for float and double, even with the same
    // default seed
    static std::uniform_real_distribution<double> dist(0.0,1.0);
    
    // Random samples
    T x1 = (T)dist(gen);
    T x2 = (T)dist(gen);
    T x3 = (T)dist(gen);
    
    // Generate 
    rotationTensorFrom3UniformRandomsArvo1992(A, x1, x2, x3);
}

template <typename T>
void randomRotationTensorArvo1992(Tensor2<T>& A, bool defaultSeed=false) {
    // Create generator once
    static std::mt19937 gen = makeMt19937(defaultSeed);
    
    // Generate
    randomRotationTensorArvo1992(gen, A);
}

template <typename T>
void rotationTensorFrom3UniformRandoms(Tensor2<T>& A, T x1, T x2, T x3) {
    rotationTensorFrom3UniformRandomsArvo1992(A, x1, x2, x3);
} 

template <typename T>
Tensor2<T> randomRotationTensorStewart1980(unsigned int dim, bool defaultSeed=false) {
    // Generator
    static std::mt19937 gen = makeMt19937(defaultSeed);
    
    // Normal distribution to draw from
    std::normal_distribution<T> dist(0.0,1.0);   
    
    // Construction
    Tensor2<T> H = identityTensor2<T>(dim);
    std::vector<T> D = ones<T>(dim);
    for (unsigned int n=1; n<dim; n++) {
        std::valarray<T> x(dim-n+1);
        for (unsigned int i=0; i<dim-n+1; i++) {
            x[i] = dist(gen);
        }
        D[n-1] = std::copysign(1.0, x[0]);
        x[0] -= D[n-1]*std::sqrt((x*x).sum());
        
        // Householder transformation
        Tensor2<T> Hx = identityTensor2<T>(dim-n+1) - (T)2.0*outer(x, x)/(x*x).sum();
        Tensor2<T> mat = identityTensor2<T>(dim);
        for (unsigned int i=n-1; i<dim; i++) {
            for (unsigned int j=n-1; j<dim; j++) {
                mat(i,j) = Hx(i-n+1,j-n+1);
            }
        }
        H = H*mat;
    }
    
    // Fix the last sign such that the determinant is 1
    D[dim-1] = -prod(D);
    H = (D*H.trans()).trans();
    
    // Return
    return H;
}

template <typename T>
void randomRotationTensorInPlace(unsigned int dim, Tensor2<T>& A, bool defaultSeed=false) {
    // Generate tensor
    if (dim == 3) {
        randomRotationTensorArvo1992<T>(A, defaultSeed);
    }
    else {
        A = randomRotationTensorStewart1980<T>(dim, defaultSeed);
    }
    
    // Check the tensor
    #ifdef DEBUG_BUILD
        if (!(A.isRotationMatrix())) {
            std::cerr << "Warning: random rotation tensor didn't pass check. Re-generating." << std::endl;
            randomRotationTensorInPlace<T>(dim, A, defaultSeed);
        }
    #endif
}

template <typename T>
Tensor2<T> randomRotationTensor(unsigned int dim, bool defaultSeed=false) {
    // Generate tensor
    Tensor2<T> A(dim,dim);
    randomRotationTensorInPlace<T>(dim, A, defaultSeed);
    return A;
}

// ROTATION SPACES //

template <typename T>
EulerAngles<T> quaternionToEulerAngles(isoi::Quaternion& q) {
    EulerAngles<T> angles;
    auto q0 = q.a;
    auto q1 = q.b;
    auto q2 = q.c;
    auto q3 = q.d;
    angles.alpha = std::atan2(2*(q0*q1+q2*q3), 1-2*(std::pow(q1,2.)+std::pow(q2,2.)));
    angles.beta = std::asin(2*(q0*q2-q3*q1));
    angles.gamma = std::atan2(2*(q0*q3+q1*q2), 1-2*(std::pow(q2,2.)+std::pow(q3,2.)));
    return angles;
}

template <typename T>
class SO3Discrete {
public:
    SO3Discrete() {;}
    SO3Discrete(unsigned int resolution, SymmetryType symmetryType = SYMMETRY_TYPE_NONE);
    isoi::Quaternion getQuat(unsigned int i) {return quatList[i];}
    EulerAngles<T> getEulerAngle(unsigned int i) {return eulerAngleList[i];}
    unsigned int size() {return quatList.size();}
private:
    SymmetryType symmetryType;
    std::vector<isoi::Quaternion> quatList;
    std::vector<EulerAngles<T>> eulerAngleList;
};

} //END NAMESPACE HPP

#endif /* HPP_ROTATION_H */