.

include/sctl/boundary_quadrature.hpp

@@ -5,6 +5,8 @@
 #include <mutex>
 #include <atomic>
 #include <tuple>
+#include <Eigen/Core>
+#include <LBFGS.h>
 
 namespace SCTL_NAMESPACE {
 
@@ -2180,6 +2182,7 @@ template <class Real, Integer ORDER=10> class Stellarator {
       return err;
     }
 
+  public:
     static Vector<ElemBasis> compute_dot_prod(const Vector<ElemBasis>& A, const Vector<ElemBasis>& B) {
       const Long Nelem = A.Dim() / COORD_DIM;
       const Long Nnodes = ElemBasis::Size();
@@ -2834,7 +2837,6 @@ template <class Real, Integer ORDER=10> class Stellarator {
       }
     }
 
-  public:
     Stellarator(const Vector<Long>& NtNp = Vector<Long>()) {
       NtNp_ = NtNp;
       Long Nsurf = NtNp_.Dim() / 2;
@@ -3703,6 +3705,9 @@ template <class Real, Integer ORDER=10> class Stellarator {
       return compute_g(Svec, pressure);
     }
 
+
+
+
     static void test() {
       constexpr Integer order_singular = 15;
       constexpr Integer order_direct = 35;
@@ -3732,153 +3737,164 @@ template <class Real, Integer ORDER=10> class Stellarator {
       }
 
       { // find equilibrium flux surfaces
-        const Long Nnodes = ElemBasis::Size();
-        const Long Nelem = S.NElem();
-        const Long Nsurf = S.Nsurf();
+        auto filter = [](const Stellarator<Real,ORDER>& S, Vector<ElemBasis>& f) {
+          auto cheb2grid = [] (const Vector<ElemBasis>& X, Long Mt, Long Mp, Long Nt, Long Np) {
+            const Long dof = X.Dim() / (Mt * Mp);
+            SCTL_ASSERT(X.Dim() == Mt * Mp *dof);
 
-        Long iter = 0;
-        Real dt = 0.1;
-        while (1) { // time-step
-          Vector<ElemBasis> normal, area_elem;
-          compute_norm_area_elem(S, normal, area_elem);
-          Vector<ElemBasis> dgdnu = compute_gradient(S, pressure, flux_tor, flux_pol)*(-1);
-          //Vector<ElemBasis> dgdnu = compute_pressure_jump(S, pressure, flux_tor, flux_pol)*(-1);
-
-          Vector<ElemBasis> dXdt(dgdnu.Dim()*COORD_DIM);
-          { // Set dXdt
-            dXdt = 0;
-            for (Long i = 0; i < S.ElemDsp(Nsurf-1); i++) {
-              for (Long j = 0; j < Nnodes; j++) {
-                dXdt[i*COORD_DIM+0][j] = normal[i*COORD_DIM+0][j] * dgdnu[i][j];
-                dXdt[i*COORD_DIM+1][j] = normal[i*COORD_DIM+1][j] * dgdnu[i][j];
-                dXdt[i*COORD_DIM+2][j] = normal[i*COORD_DIM+2][j] * dgdnu[i][j];
-              }
-            }
-          }
-          for (Long i = 0; i < S.Nsurf(); i++) { // filter dXdt
-            const Long Nt0 = S.NTor(i);
-            const Long Np0 = S.NPol(i);
-            const Long Nelem = Nt0 * Np0;
-            const Long offset = S.ElemDsp(i);
+            Vector<Real> Xf(dof*Nt*Np); Xf = 0;
             const Long Nnodes = ElemBasis::Size();
-            const Long INTERP_ORDER = 12;
-            const Long Nt = Nt0*ORDER/5;
-            const Long Np = Np0*ORDER/5;
-            for (Long k = 0; k < COORD_DIM; k++) {
-              Matrix<Real> M(Nt, Np); M = 0;
-              const auto& quad_wts = ElemBasis::QuadWts();
-              const Matrix<Real>& Mnodes = Basis<Real,1,ORDER>::Nodes();
-              for (Long tt = 0; tt < Nt0; tt++) { // Set M
-                for (Long pp = 0; pp < Np0; pp++) {
-                  for (Long t = 0; t < ORDER; t++) {
-                    for (Long p = 0; p < ORDER; p++) {
-                      Real theta = (tt + Mnodes[0][t]) / Nt0;
-                      Real phi   = (pp + Mnodes[0][p]) / Np0;
-                      Long i = (Long)(theta * Nt);
-                      Long j = (Long)(phi   * Np);
-                      Real x = theta * Nt - i;
-                      Real y = phi   * Np - j;
-
-                      Long elem_idx = tt * Np0 + pp;
-                      Long node_idx = p * ORDER + t;
-
-                      Vector<Real> Interp0(INTERP_ORDER);
-                      Vector<Real> Interp1(INTERP_ORDER);
-                      { // Set Interp0, Interp1
-                        auto node = [] (Long i) {
-                          return (Real)i - (INTERP_ORDER-1)/2;
-                        };
-                        for (Long i = 0; i < INTERP_ORDER; i++) {
-                          Real wt_x = 1, wt_y = 1;
-                          for (Long j = 0; j < INTERP_ORDER; j++) {
-                            if (j != i) {
-                              wt_x *= (x - node(j)) / (node(i) - node(j));
-                              wt_y *= (y - node(j)) / (node(i) - node(j));
-                            }
-                            Interp0[i] = wt_x;
-                            Interp1[i] = wt_y;
-                          }
-                        }
+            const Matrix<Real>& Mnodes = Basis<Real,1,ORDER>::Nodes();
+            for (Long t = 0; t < Nt; t++) {
+              for (Long p = 0; p < Np; p++) {
+                Real theta = t / (Real)Nt;
+                Real phi   = p / (Real)Np;
+                Long i = (Long)(theta * Mt);
+                Long j = (Long)(phi   * Mp);
+                Real x = theta * Mt - i;
+                Real y = phi   * Mp - j;
+                Long elem_idx = i * Mp + j;
+
+                Vector<Real> Interp0(ORDER);
+                Vector<Real> Interp1(ORDER);
+                { // Set Interp0, Interp1
+                  auto node = [&Mnodes] (Long i) {
+                    return Mnodes[0][i];
+                  };
+                  for (Long i = 0; i < ORDER; i++) {
+                    Real wt_x = 1, wt_y = 1;
+                    for (Long j = 0; j < ORDER; j++) {
+                      if (j != i) {
+                        wt_x *= (x - node(j)) / (node(i) - node(j));
+                        wt_y *= (y - node(j)) / (node(i) - node(j));
                       }
+                      Interp0[i] = wt_x;
+                      Interp1[i] = wt_y;
+                    }
+                  }
+                }
 
-                      for (Long ii = 0; ii < INTERP_ORDER; ii++) {
-                        for (Long jj = 0; jj < INTERP_ORDER; jj++) {
-                          Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
-                          Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
-                          M[idx_i][idx_j] += dXdt[(offset+elem_idx)*COORD_DIM+k][node_idx] * quad_wts[node_idx] * Interp0[ii] * Interp1[jj] / (Nt0 * Np0) * (Nt * Np);
-                        }
-                      }
+                for (Long ii = 0; ii < ORDER; ii++) {
+                  for (Long jj = 0; jj < ORDER; jj++) {
+                    Long node_idx = jj * ORDER + ii;
+                    for (Long k = 0; k < dof; k++) {
+                      Xf[(k*Nt+t)*Np+p] += X[elem_idx*dof+k][node_idx] * Interp0[ii] * Interp1[jj];
                     }
                   }
                 }
               }
-
-              for (Long tt = 0; tt < Nt0; tt++) {
-                for (Long pp = 0; pp < Np0; pp++) {
-                  for (Long t = 0; t < ORDER; t++) {
-                    for (Long p = 0; p < ORDER; p++) {
-                      Matrix<Real> Mnodes = Basis<Real,1,ORDER>::Nodes();
-                      Real theta = (tt + Mnodes[0][t]) / Nt0;
-                      Real phi   = (pp + Mnodes[0][p]) / Np0;
-                      Long i = (Long)(theta * Nt);
-                      Long j = (Long)(phi   * Np);
-                      Real x = theta * Nt - i;
-                      Real y = phi   * Np - j;
-
-                      Vector<Real> Interp0(INTERP_ORDER);
-                      Vector<Real> Interp1(INTERP_ORDER);
-                      { // Set Interp0, Interp1
-                        auto node = [] (Long i) {
-                          return (Real)i - (INTERP_ORDER-1)/2;
-                        };
-                        for (Long i = 0; i < INTERP_ORDER; i++) {
-                          Real wt_x = 1, wt_y = 1;
-                          for (Long j = 0; j < INTERP_ORDER; j++) {
-                            if (j != i) {
-                              wt_x *= (x - node(j)) / (node(i) - node(j));
-                              wt_y *= (y - node(j)) / (node(i) - node(j));
-                            }
-                            Interp0[i] = wt_x;
-                            Interp1[i] = wt_y;
+            }
+            return Xf;
+          };
+          auto grid2cheb = [] (const Vector<Real>& Xf, Long Nt, Long Np, Long Mt, Long Mp) {
+            Long dof = Xf.Dim() / (Nt*Np);
+            SCTL_ASSERT(Xf.Dim() == dof*Nt*Np);
+            Vector<ElemBasis> X(Mt*Mp*dof);
+            constexpr Integer INTERP_ORDER = 12;
+
+            for (Long tt = 0; tt < Mt; tt++) {
+              for (Long pp = 0; pp < Mp; pp++) {
+                for (Long t = 0; t < ORDER; t++) {
+                  for (Long p = 0; p < ORDER; p++) {
+                    Matrix<Real> Mnodes = Basis<Real,1,ORDER>::Nodes();
+                    Real theta = (tt + Mnodes[0][t]) / Mt;
+                    Real phi   = (pp + Mnodes[0][p]) / Mp;
+                    Long i = (Long)(theta * Nt);
+                    Long j = (Long)(phi   * Np);
+                    Real x = theta * Nt - i;
+                    Real y = phi   * Np - j;
+
+                    Vector<Real> Interp0(INTERP_ORDER);
+                    Vector<Real> Interp1(INTERP_ORDER);
+                    { // Set Interp0, Interp1
+                      auto node = [] (Long i) {
+                        return (Real)i - (INTERP_ORDER-1)/2;
+                      };
+                      for (Long i = 0; i < INTERP_ORDER; i++) {
+                        Real wt_x = 1, wt_y = 1;
+                        for (Long j = 0; j < INTERP_ORDER; j++) {
+                          if (j != i) {
+                            wt_x *= (x - node(j)) / (node(i) - node(j));
+                            wt_y *= (y - node(j)) / (node(i) - node(j));
                           }
+                          Interp0[i] = wt_x;
+                          Interp1[i] = wt_y;
                         }
                       }
+                    }
 
-                      Real f0 = 0;
+                    for (Long k = 0; k < dof; k++) {
+                      Real X0 = 0;
                       for (Long ii = 0; ii < INTERP_ORDER; ii++) {
                         for (Long jj = 0; jj < INTERP_ORDER; jj++) {
                           Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
                           Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
-                          f0 += Interp0[ii] * Interp1[jj] * M[idx_i][idx_j];
+                          X0 += Interp0[ii] * Interp1[jj] * Xf[(k*Nt+idx_i)*Np+idx_j];
                         }
                       }
-
-                      Long elem_idx = tt * Np0 + pp;
+                      Long elem_idx = tt * Mp + pp;
                       Long node_idx = p * ORDER + t;
-                      dXdt[(offset+elem_idx)*COORD_DIM+k][node_idx] = f0;
+                      X[elem_idx*dof+k][node_idx] = X0;
                     }
                   }
                 }
               }
             }
+            return X;
+          };
+
+          Long dof = f.Dim() / S.NElem();
+          SCTL_ASSERT(f.Dim() == S.NElem() * dof);
+          for (Long i = 0; i < S.Nsurf(); i++) {
+            const Long Mt = S.NTor(i);
+            const Long Mp = S.NPol(i);
+            const Long Nelem = Mt * Mp;
+            const Long offset = S.ElemDsp(i);
+            const Long Nt = Mt * ORDER / 5;
+            const Long Np = Mp * ORDER / 5;
+
+            Vector<ElemBasis> f_(Nelem*dof, f.begin() + offset*dof, false);
+            Vector<Real> f_fourier = cheb2grid(f_, Mt, Mp, Nt, Np);
+            f_ = grid2cheb(f_fourier, Nt, Np, Mt, Mp);
           }
+        };
 
-          { // Update S, dt
+        Long iter = 0;
+        Real dt = 0.1;
+        while (1) { // time-step
+          Vector<ElemBasis> dgdnu = compute_gradient(S, pressure, flux_tor, flux_pol)*(-1);
+          //Vector<ElemBasis> dgdnu = compute_pressure_jump(S, pressure, flux_tor, flux_pol)*(-1);
+
+          Vector<ElemBasis> dXdt(dgdnu.Dim()*COORD_DIM);
+          { // Set dXdt
+            dXdt = 0;
+            const Long Nnodes = ElemBasis::Size();
+            Vector<ElemBasis> normal, area_elem;
+            compute_norm_area_elem(S, normal, area_elem);
+            for (Long i = 0; i < S.ElemDsp(S.Nsurf()-1); i++) {
+              for (Long j = 0; j < Nnodes; j++) {
+                dXdt[i*COORD_DIM+0][j] = normal[i*COORD_DIM+0][j] * dgdnu[i][j];
+                dXdt[i*COORD_DIM+1][j] = normal[i*COORD_DIM+1][j] * dgdnu[i][j];
+                dXdt[i*COORD_DIM+2][j] = normal[i*COORD_DIM+2][j] * dgdnu[i][j];
+              }
+            }
+            filter(S, dXdt);
+          }
+          { // Update dt
+            const Long Nelem = S.NElem();
             Stellarator<Real,ORDER> S0 = S, S1 = S, S2 = S;
             for (Long i = 0; i < S.NElem(); i++) {
-              for (Long j = 0; j < Nnodes; j++) {
-                S0.Elem(i, 0)[j] += 0.0 * dt * dXdt[i*COORD_DIM+0][j];
-                S0.Elem(i, 1)[j] += 0.0 * dt * dXdt[i*COORD_DIM+1][j];
-                S0.Elem(i, 2)[j] += 0.0 * dt * dXdt[i*COORD_DIM+2][j];
+              S0.Elem(i, 0) += dXdt[i*COORD_DIM+0] * 0.0 * dt;
+              S0.Elem(i, 1) += dXdt[i*COORD_DIM+1] * 0.0 * dt;
+              S0.Elem(i, 2) += dXdt[i*COORD_DIM+2] * 0.0 * dt;
 
-                S1.Elem(i, 0)[j] += 0.5 * dt * dXdt[i*COORD_DIM+0][j];
-                S1.Elem(i, 1)[j] += 0.5 * dt * dXdt[i*COORD_DIM+1][j];
-                S1.Elem(i, 2)[j] += 0.5 * dt * dXdt[i*COORD_DIM+2][j];
+              S1.Elem(i, 0) += dXdt[i*COORD_DIM+0] * 0.5 * dt;
+              S1.Elem(i, 1) += dXdt[i*COORD_DIM+1] * 0.5 * dt;
+              S1.Elem(i, 2) += dXdt[i*COORD_DIM+2] * 0.5 * dt;
 
-                S2.Elem(i, 0)[j] += 1.0 * dt * dXdt[i*COORD_DIM+0][j];
-                S2.Elem(i, 1)[j] += 1.0 * dt * dXdt[i*COORD_DIM+1][j];
-                S2.Elem(i, 2)[j] += 1.0 * dt * dXdt[i*COORD_DIM+2][j];
-              }
+              S2.Elem(i, 0) += dXdt[i*COORD_DIM+0] * 1.0 * dt;
+              S2.Elem(i, 1) += dXdt[i*COORD_DIM+1] * 1.0 * dt;
+              S2.Elem(i, 2) += dXdt[i*COORD_DIM+2] * 1.0 * dt;
             }
             Real g0 = compute_g(S0, pressure, flux_tor, flux_pol);
             Real g1 = compute_g(S1, pressure, flux_tor, flux_pol);
@@ -3897,30 +3913,42 @@ template <class Real, Integer ORDER=10> class Stellarator {
               std::cout<<g2<<' ';
               std::cout<<dt<<'\n';
             }
-
-            for (Long i = 0; i < S.NElem(); i++) {
-              for (Long j = 0; j < Nnodes; j++) {
-                S.Elem(i, 0)[j] += dt * dXdt[i*COORD_DIM+0][j];
-                S.Elem(i, 1)[j] += dt * dXdt[i*COORD_DIM+1][j];
-                S.Elem(i, 2)[j] += dt * dXdt[i*COORD_DIM+2][j];
-              }
-            }
           }
-          iter++;
 
-          Vector<ElemBasis> pressure_jump = compute_pressure_jump(S, pressure, flux_tor, flux_pol);
           { // Write VTU
             VTUData vtu;
-            vtu.AddElems(S.GetElemList(), dgdnu, ORDER);
+            vtu.AddElems(S.GetElemList(), dgdnu*dt, ORDER);
             vtu.WriteVTK("dgdnu"+std::to_string(iter), comm);
           }
           { // Write VTU
+            VTUData vtu;
+            vtu.AddElems(S.GetElemList(), dXdt*dt, ORDER);
+            vtu.WriteVTK("dXdt"+std::to_string(iter), comm);
+          }
+          { // Write VTU
+            Vector<ElemBasis> pressure_jump = compute_pressure_jump(S, pressure, flux_tor, flux_pol);
             VTUData vtu;
             vtu.AddElems(S.GetElemList(), pressure_jump, ORDER);
             vtu.WriteVTK("pressure_jump"+std::to_string(iter), comm);
           }
-        }
 
+          { // Update S <-- filter(S + dXdt * dt)
+            const Long Nelem = S.NElem();
+            Vector<ElemBasis> X(Nelem*COORD_DIM);
+            for (Long i = 0; i < S.NElem(); i++) {
+              X[i*COORD_DIM+0] = S.Elem(i, 0) + dXdt[i*COORD_DIM+0] * dt * 0.5;
+              X[i*COORD_DIM+1] = S.Elem(i, 1) + dXdt[i*COORD_DIM+1] * dt * 0.5;
+              X[i*COORD_DIM+2] = S.Elem(i, 2) + dXdt[i*COORD_DIM+2] * dt * 0.5;
+            }
+            filter(S, X);
+            for (Long i = 0; i < S.NElem(); i++) {
+              S.Elem(i, 0) = X[i*COORD_DIM+0];
+              S.Elem(i, 1) = X[i*COORD_DIM+1];
+              S.Elem(i, 2) = X[i*COORD_DIM+2];
+            }
+          }
+          iter++;
+        }
         return;
       }
 
@@ -4801,6 +4829,274 @@ template <class Real, Integer ORDER=10> class Stellarator {
     Vector<Long> elem_dsp;
 };
 
+template <class Real, Integer ORDER=10> class MHDEquilib {
+    static constexpr Integer COORD_DIM = 3;
+    static constexpr Integer ELEM_DIM = COORD_DIM-1;
+    using ElemBasis = Basis<Real, ELEM_DIM, ORDER>;
+
+  public:
+    MHDEquilib(const Stellarator<Real,ORDER>& S, const Vector<Real>& pressure, const Vector<Real>& flux_tor, const Vector<Real>& flux_pol) {
+      S_ = S;
+      pressure_ = pressure;
+      flux_tor_ = flux_tor;
+      flux_pol_ = flux_pol;
+      iter = 0;
+    }
+
+    Real operator()(const Eigen::VectorXd& x, Eigen::VectorXd& grad) {
+      const Comm comm = Comm::World();
+      const Long Nelem = S_.NElem();
+      const Long Nnodes = ElemBasis::Size();
+      const Long N = Nelem * COORD_DIM * Nnodes;
+      SCTL_ASSERT(x.rows() == N);
+
+      auto filter = [](const Stellarator<Real,ORDER>& S, Vector<ElemBasis>& f) {
+        auto cheb2grid = [] (const Vector<ElemBasis>& X, Long Mt, Long Mp, Long Nt, Long Np) {
+          const Long dof = X.Dim() / (Mt * Mp);
+          SCTL_ASSERT(X.Dim() == Mt * Mp *dof);
+
+          Vector<Real> Xf(dof*Nt*Np); Xf = 0;
+          const Long Nnodes = ElemBasis::Size();
+          const Matrix<Real>& Mnodes = Basis<Real,1,ORDER>::Nodes();
+          for (Long t = 0; t < Nt; t++) {
+            for (Long p = 0; p < Np; p++) {
+              Real theta = t / (Real)Nt;
+              Real phi   = p / (Real)Np;
+              Long i = (Long)(theta * Mt);
+              Long j = (Long)(phi   * Mp);
+              Real x = theta * Mt - i;
+              Real y = phi   * Mp - j;
+              Long elem_idx = i * Mp + j;
+
+              Vector<Real> Interp0(ORDER);
+              Vector<Real> Interp1(ORDER);
+              { // Set Interp0, Interp1
+                auto node = [&Mnodes] (Long i) {
+                  return Mnodes[0][i];
+                };
+                for (Long i = 0; i < ORDER; i++) {
+                  Real wt_x = 1, wt_y = 1;
+                  for (Long j = 0; j < ORDER; j++) {
+                    if (j != i) {
+                      wt_x *= (x - node(j)) / (node(i) - node(j));
+                      wt_y *= (y - node(j)) / (node(i) - node(j));
+                    }
+                    Interp0[i] = wt_x;
+                    Interp1[i] = wt_y;
+                  }
+                }
+              }
+
+              for (Long ii = 0; ii < ORDER; ii++) {
+                for (Long jj = 0; jj < ORDER; jj++) {
+                  Long node_idx = jj * ORDER + ii;
+                  for (Long k = 0; k < dof; k++) {
+                    Xf[(k*Nt+t)*Np+p] += X[elem_idx*dof+k][node_idx] * Interp0[ii] * Interp1[jj];
+                  }
+                }
+              }
+            }
+          }
+          return Xf;
+        };
+        auto grid2cheb = [] (const Vector<Real>& Xf, Long Nt, Long Np, Long Mt, Long Mp) {
+          Long dof = Xf.Dim() / (Nt*Np);
+          SCTL_ASSERT(Xf.Dim() == dof*Nt*Np);
+          Vector<ElemBasis> X(Mt*Mp*dof);
+          constexpr Integer INTERP_ORDER = 12;
+
+          for (Long tt = 0; tt < Mt; tt++) {
+            for (Long pp = 0; pp < Mp; pp++) {
+              for (Long t = 0; t < ORDER; t++) {
+                for (Long p = 0; p < ORDER; p++) {
+                  Matrix<Real> Mnodes = Basis<Real,1,ORDER>::Nodes();
+                  Real theta = (tt + Mnodes[0][t]) / Mt;
+                  Real phi   = (pp + Mnodes[0][p]) / Mp;
+                  Long i = (Long)(theta * Nt);
+                  Long j = (Long)(phi   * Np);
+                  Real x = theta * Nt - i;
+                  Real y = phi   * Np - j;
+
+                  Vector<Real> Interp0(INTERP_ORDER);
+                  Vector<Real> Interp1(INTERP_ORDER);
+                  { // Set Interp0, Interp1
+                    auto node = [] (Long i) {
+                      return (Real)i - (INTERP_ORDER-1)/2;
+                    };
+                    for (Long i = 0; i < INTERP_ORDER; i++) {
+                      Real wt_x = 1, wt_y = 1;
+                      for (Long j = 0; j < INTERP_ORDER; j++) {
+                        if (j != i) {
+                          wt_x *= (x - node(j)) / (node(i) - node(j));
+                          wt_y *= (y - node(j)) / (node(i) - node(j));
+                        }
+                        Interp0[i] = wt_x;
+                        Interp1[i] = wt_y;
+                      }
+                    }
+                  }
+
+                  for (Long k = 0; k < dof; k++) {
+                    Real X0 = 0;
+                    for (Long ii = 0; ii < INTERP_ORDER; ii++) {
+                      for (Long jj = 0; jj < INTERP_ORDER; jj++) {
+                        Long idx_i = (i + ii-(INTERP_ORDER-1)/2 + Nt) % Nt;
+                        Long idx_j = (j + jj-(INTERP_ORDER-1)/2 + Np) % Np;
+                        X0 += Interp0[ii] * Interp1[jj] * Xf[(k*Nt+idx_i)*Np+idx_j];
+                      }
+                    }
+                    Long elem_idx = tt * Mp + pp;
+                    Long node_idx = p * ORDER + t;
+                    X[elem_idx*dof+k][node_idx] = X0;
+                  }
+                }
+              }
+            }
+          }
+          return X;
+        };
+
+        Long dof = f.Dim() / S.NElem();
+        SCTL_ASSERT(f.Dim() == S.NElem() * dof);
+        for (Long i = 0; i < S.Nsurf(); i++) {
+          const Long Mt = S.NTor(i);
+          const Long Mp = S.NPol(i);
+          const Long Nelem = Mt * Mp;
+          const Long offset = S.ElemDsp(i);
+          const Long Nt = Mt * ORDER / 10;
+          const Long Np = Mp * ORDER / 10;
+
+          Vector<ElemBasis> f_(Nelem*dof, f.begin() + offset*dof, false);
+          Vector<Real> f_fourier = cheb2grid(f_, Mt, Mp, Nt, Np);
+          f_ = grid2cheb(f_fourier, Nt, Np, Mt, Mp);
+        }
+      };
+
+      for (Long i = 0; i < Nelem; i++) { // Set S_
+        for (Long j = 0; j < Nnodes; j++) {
+          S_.Elem(i,0)[j] = x[(i*Nnodes+j)*COORD_DIM+0];
+          S_.Elem(i,1)[j] = x[(i*Nnodes+j)*COORD_DIM+1];
+          S_.Elem(i,2)[j] = x[(i*Nnodes+j)*COORD_DIM+2];
+        }
+      }
+      Real g = Stellarator<Real,ORDER>::compute_g(S_, pressure_, flux_tor_, flux_pol_);
+      Vector<ElemBasis> dgdnu = Stellarator<Real,ORDER>::compute_gradient(S_, pressure_, flux_tor_, flux_pol_);
+      Vector<ElemBasis> dXdt(N);
+      { // Set dXdt
+        dXdt = 0;
+        const Long Nnodes = ElemBasis::Size();
+        Vector<ElemBasis> normal, area_elem;
+        Stellarator<Real,ORDER>::compute_norm_area_elem(S_, normal, area_elem);
+        for (Long i = 0; i < S_.ElemDsp(S_.Nsurf()-1); i++) {
+          for (Long j = 0; j < Nnodes; j++) {
+            dXdt[i*COORD_DIM+0][j] = normal[i*COORD_DIM+0][j] * dgdnu[i][j];
+            dXdt[i*COORD_DIM+1][j] = normal[i*COORD_DIM+1][j] * dgdnu[i][j];
+            dXdt[i*COORD_DIM+2][j] = normal[i*COORD_DIM+2][j] * dgdnu[i][j];
+          }
+        }
+        filter(S_, dXdt);
+      }
+      for (Long i = 0; i < Nelem; i++) { // Set grad
+        for (Long j = 0; j < Nnodes; j++) {
+          grad[(i*Nnodes+j)*COORD_DIM+0] = dXdt[i*COORD_DIM+0][j];
+          grad[(i*Nnodes+j)*COORD_DIM+1] = dXdt[i*COORD_DIM+1][j];
+          grad[(i*Nnodes+j)*COORD_DIM+2] = dXdt[i*COORD_DIM+2][j];
+        }
+      }
+
+      if (1) { // Write VTU
+        VTUData vtu;
+        vtu.AddElems(S_.GetElemList(), dgdnu, ORDER);
+        vtu.WriteVTK("dgdnu"+std::to_string(iter), comm);
+      }
+      std::cout<<"iter = "<<iter<<"    g = "<<g<<'\n';
+
+      iter++;
+      return g;
+    }
+
+    static void ComputeEquilibrium(MHDEquilib& mhd_equilib) {
+      const Long Nelem = mhd_equilib.S_.NElem();
+      const Long Nnodes = ElemBasis::Size();
+      const Long N = Nelem * COORD_DIM * Nnodes;
+
+      LBFGSpp::LBFGSParam<Real> param;
+      param.epsilon = 1e-6;
+      param.max_iterations = 100;
+
+      // Create solver and function object
+      LBFGSpp::LBFGSSolver<Real> solver(param);
+
+      // Initial guess
+      Eigen::VectorXd x = Eigen::VectorXd::Zero(N);
+      for (Long i = 0; i < Nelem; i++) { // Set x
+        for (Long j = 0; j < Nnodes; j++) {
+          x((i*Nnodes+j)*COORD_DIM+0) = mhd_equilib.S_.Elem(i,0)[j];
+          x((i*Nnodes+j)*COORD_DIM+1) = mhd_equilib.S_.Elem(i,1)[j];
+          x((i*Nnodes+j)*COORD_DIM+2) = mhd_equilib.S_.Elem(i,2)[j];
+        }
+      }
+
+      Real fx;
+      Integer niter = solver.minimize(mhd_equilib, x, fx);
+      for (Long i = 0; i < Nelem; i++) { // Set x
+        for (Long j = 0; j < Nnodes; j++) {
+          mhd_equilib.S_.Elem(i,0)[j] = x((i*Nnodes+j)*COORD_DIM+0);
+          mhd_equilib.S_.Elem(i,1)[j] = x((i*Nnodes+j)*COORD_DIM+1);
+          mhd_equilib.S_.Elem(i,2)[j] = x((i*Nnodes+j)*COORD_DIM+2);
+        }
+      }
+
+      std::cout << niter << " iterations" <<'\n';
+      std::cout << "f(x) = " << fx <<'\n';
+    }
+
+    static void test() {
+      constexpr Integer order_singular = 25;
+      constexpr Integer order_direct = 35;
+      Comm comm = Comm::World();
+      Profile::Enable(true);
+
+      Long Nsurf = 2;
+      Stellarator<Real,ORDER> S;
+      Vector<Real> flux_tor(Nsurf), flux_pol(Nsurf), pressure(Nsurf);
+      { // Init S, flux_tor, flux_pol, pressure
+        Vector<Long> NtNp;
+        NtNp.PushBack(50);
+        NtNp.PushBack(8);
+        NtNp.PushBack(50);
+        NtNp.PushBack(8);
+        //for (Long i = 0; i < Nsurf; i++) {
+        //  NtNp.PushBack(30);
+        //  NtNp.PushBack(4);
+        //}
+        S = Stellarator<Real,ORDER>(NtNp);
+        flux_tor = 1;
+        flux_pol = 1;
+        pressure = 0;
+
+        //flux_tor[0] = 1; //0.791881512;
+        //flux_tor[1] = 1;
+        //flux_pol[0] = 0;
+        //flux_pol[1] = 0;
+        //pressure[0] = 0;
+        //pressure[1] = 0;
+      }
+
+      MHDEquilib mhd_equilib(S, pressure, flux_tor, flux_pol);
+      ComputeEquilibrium(mhd_equilib);
+    }
+
+  private:
+    Stellarator<Real,ORDER> S_;
+    Vector<Real> pressure_;
+    Vector<Real> flux_tor_;
+    Vector<Real> flux_pol_;
+    Long iter = 0;
+};
+
+
+
 template <class Real, Integer ORDER=5> class Spheres {
   static constexpr Integer COORD_DIM = 3;
   static constexpr Integer ELEM_DIM = COORD_DIM-1;

src/test.cpp

@@ -60,6 +60,9 @@ void TestMatrix() {
 int main(int argc, char** argv) {
   sctl::Comm::MPI_Init(&argc, &argv);
 
+  sctl::MHDEquilib<double>::test();
+  return 0;
+
   sctl::Stellarator<double>::test();
   return 0;
 
@@ -98,3 +101,4 @@ int main(int argc, char** argv) {
   sctl::Comm::MPI_Finalize();
   return 0;
 }
+