fit_ellipsoid_3d.py

Go to the documentation of this file.

from __future__ import print_function
from sympy import *
import sys

# Lets work this out symbolically first and perform python unit test.

# Equation for the ellipsoid can be written sxx*x^2 + syy*y^2 + szz*z^2 + syz*y*z + sxz*x*z + sxy*x*y = 1,

sxx,syy,szz,syz,sxz,sxy=symbols("sxx,syy,szz,syz,sxz,sxy")
S=Matrix([[sxx],[syy],[szz],[syz],[sxz],[sxy]])

x,y,z=symbols('x,y,z')

P=Matrix([[x**2, y**2, z**2, y*z, x*z, x*y]])

A = P.transpose()*P
b = P.transpose()*Matrix([[1]])

# Sphere of radius sqrt(3). Sxx=1/h^2, therefore Sxx=1/3
sphere = {x:1, y:1, z:1,
          sxx:Rational(1,3),syy:Rational(1, 3),szz:Rational(1, 3),syz:0,sxz:0,sxy:0}

if (A*S).evalf(subs=sphere)==b.evalf(subs=sphere):
    print("pass")
else:
    print("fail")
    sys.exit(-1)

# Move onto code generation.

pyname=sys.argv[0].split('/')[-1]

# Write source file
cxxname=pyname[:-3]+".cpp"

src="""
/* Start of code generated by %s. Warning - be careful about modifying
   any of the generated code directly.  Any changes/fixes should be done
   in the code generation script generation.
 */
"""%(pyname)

src += """
void fit_ellipsoid(int i, real_t *sm){
  Eigen::Matrix<double, 6, 6> A = Eigen::Matrix<real_t, 6, 6>::Zero(6,6);
  Eigen::Matrix<double, 6, 1> b = Eigen::Matrix<real_t, 6, 1>::Zero(6);

  std::vector<index_t> nodes = _mesh->NNList[i];
  nodes.push_back(i);

  for(typename std::vector<index_t>::const_iterator it=nodes.begin();it!=nodes.end();++it){
    const real_t *X0=_mesh->get_coords(*it);
    real_t x0=X0[0], y0=X0[1], z0=X0[2];
    assert(std::isfinite(x0));
    assert(std::isfinite(y0));
    assert(std::isfinite(z0));

    for(typename std::vector<index_t>::const_iterator n=_mesh->NNList[*it].begin();n!=_mesh->NNList[*it].end();++n){
      if(*n<=*it)
    continue;
      
      const real_t *X=_mesh->get_coords(*n);
      real_t x=X[0]-x0, y=X[1]-y0, z=X[2]-z0;

      assert(std::isfinite(x));
      assert(std::isfinite(y));
      assert(std::isfinite(z));
      if(x<0){
        x*=-1;
        y*=-1;
        z*=-1;
      }
      """

for i in range(6):
    src+="      "
    for j in range(6):
        src+="A[%d]+="%(i*6+j)+ccode(A[i,j], contract=False)+"; "
    src+="\n"

for i in range(6):
    src+="b[%d]+="%(i)+ccode(b[i], contract=False)+";\n"

src+="""
    }
  }

  Eigen::Matrix<double, 6, 1> S = Eigen::Matrix<real_t, 6, 1>::Zero(6);

  A.svd().solve(b, &S);

  if(_mesh->NNList[i].size()>=6){
    sm[0] = S[0]; sm[1] = S[5]; sm[2] = S[4];
                  sm[3] = S[1]; sm[4] = S[3];
                                sm[5] = S[2];
  }else{
    sm[0] = S[0]; sm[1] = 0;    sm[2] = 0;
                  sm[3] = S[1]; sm[4] = 0;
                                sm[5] = S[2];
  }

  return;
}

/* End of code generated by %s. Warning - be careful about
   modifying any of the generated code directly.  Any changes/fixes
   should be done in the code generation script generation.*/\n"""%pyname

cxxfile = open(cxxname, "w")
cxxfile.write(src)
cxxfile.close()