.

include/sctl/boundary_quadrature.hpp

@@ -4936,53 +4936,86 @@ template <class Real, Integer ORDER=10> class Stellarator {
 };
 
 template <class Real, Integer ORDER=10> class MHDEquilib {
+    static constexpr Integer fourier_upsample = 2;
     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;
-    }
+    static Vector<Real> 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);
 
-    Real operator()(const Eigen::VectorXd& x, Eigen::VectorXd& grad) {
-      const Comm comm = Comm::World();
-      const Long Nelem = S_.NElem();
+      Vector<Real> Xf(dof*Nt*Np); Xf = 0;
       const Long Nnodes = ElemBasis::Size();
-      const Long N = Nelem * COORD_DIM * Nnodes;
-      SCTL_ASSERT(x.rows() == N);
-
-      auto filter = [](const Stellarator<Real,ORDER>& S, const Comm& comm, Vector<ElemBasis>& f, Real sigma) {
-        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);
+      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;
+              }
+            }
+          }
 
-          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);
+          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;
+    }
+    static Vector<ElemBasis> 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 = [&Mnodes] (Long i) {
-                  return Mnodes[0][i];
+                auto node = [] (Long i) {
+                  return (Real)i - (INTERP_ORDER-1)/2;
                 };
-                for (Long i = 0; i < ORDER; i++) {
+                for (Long i = 0; i < INTERP_ORDER; i++) {
                   Real wt_x = 1, wt_y = 1;
-                  for (Long j = 0; j < ORDER; j++) {
+                  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));
@@ -4993,172 +5026,214 @@ template <class Real, Integer ORDER=10> class MHDEquilib {
                 }
               }
 
-              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 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 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;
-                      }
-                    }
-                  }
+        }
+      }
+      return X;
+    }
+    static void fourier_filter(sctl::Vector<Real>& X, long Nt_, long Np_, Real sigma, const Comm& comm) {
+      long dof = X.Dim() / (Nt_ * Np_);
+      SCTL_ASSERT(X.Dim() == dof * Nt_ * Np_);
+
+      sctl::FFT<Real> fft_r2c, fft_c2r;
+      sctl::StaticArray<sctl::Long, 2> fft_dim = {Nt_, Np_};
+      fft_r2c.Setup(sctl::FFT_Type::R2C, 1, sctl::Vector<sctl::Long>(2, fft_dim, false), omp_get_max_threads());
+      fft_c2r.Setup(sctl::FFT_Type::C2R, 1, sctl::Vector<sctl::Long>(2, fft_dim, false), omp_get_max_threads());
+
+      long Nt = Nt_;
+      long Np = fft_r2c.Dim(1) / (Nt * 2);
+      SCTL_ASSERT(fft_r2c.Dim(1) == Nt * Np * 2);
+
+      auto filter_fn = [](Real x2, Real sigma) {return exp(-x2/(2*sigma*sigma));};
+      //auto filter_fn = [](Real x2, Real sigma) {return (x2<sigma*sigma?1.0:0.0);};
+
+      sctl::Vector<Real> normal, gradX;
+      biest::SurfaceOp<Real> op(comm, Nt_, Np_);
+      sctl::Vector<Real> coeff(fft_r2c.Dim(1));
+      for (long k = 0; k < dof; k++) {
+        sctl::Vector<Real> X_(Nt_*Np_, X.begin() + k*Nt_*Np_, false);
+        fft_r2c.Execute(X_, coeff);
+        for (long t = 0; t < Nt; t++) {
+          for (long p = 0; p < Np; p++) {
+            Real tt = (t - (t > Nt / 2 ? Nt : 0)) / (Real)(Nt / 2);
+            Real pp = p / (Real)Np;
+            Real f = filter_fn(tt*tt+pp*pp, sigma);
+            coeff[(t * Np + p) * 2 + 0] *= f;
+            coeff[(t * Np + p) * 2 + 1] *= f;
+          }
+        }
+        fft_c2r.Execute(coeff, X_);
+      }
+    };
 
-                  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;
-        };
-        auto fourier_filter = [](sctl::Vector<Real>& X, long Nt_, long Np_, Real sigma, const Comm& comm) {
-          long dof = X.Dim() / (Nt_ * Np_);
-          SCTL_ASSERT(X.Dim() == dof * Nt_ * Np_);
-
-          sctl::FFT<Real> fft_r2c, fft_c2r;
-          sctl::StaticArray<sctl::Long, 2> fft_dim = {Nt_, Np_};
-          fft_r2c.Setup(sctl::FFT_Type::R2C, 1, sctl::Vector<sctl::Long>(2, fft_dim, false), omp_get_max_threads());
-          fft_c2r.Setup(sctl::FFT_Type::C2R, 1, sctl::Vector<sctl::Long>(2, fft_dim, false), omp_get_max_threads());
-
-          long Nt = Nt_;
-          long Np = fft_r2c.Dim(1) / (Nt * 2);
-          SCTL_ASSERT(fft_r2c.Dim(1) == Nt * Np * 2);
-
-          auto filter_fn = [](Real x2, Real sigma) {return exp(-x2/(2*sigma*sigma));};
-          //auto filter_fn = [](Real x2, Real sigma) {return (x2<sigma*sigma?1.0:0.0);};
-
-          sctl::Vector<Real> normal, gradX;
-          biest::SurfaceOp<Real> op(comm, Nt_, Np_);
-          sctl::Vector<Real> coeff(fft_r2c.Dim(1));
-          for (long k = 0; k < dof; k++) {
-            sctl::Vector<Real> X_(Nt_*Np_, X.begin() + k*Nt_*Np_, false);
-            fft_r2c.Execute(X_, coeff);
-            for (long t = 0; t < Nt; t++) {
-              for (long p = 0; p < Np; p++) {
-                Real tt = (t - (t > Nt / 2 ? Nt : 0)) / (Real)(Nt / 2);
-                Real pp = p / (Real)Np;
-                Real f = filter_fn(tt*tt+pp*pp, sigma);
-                coeff[(t * Np + p) * 2 + 0] *= f;
-                coeff[(t * Np + p) * 2 + 1] *= f;
-              }
-            }
-            fft_c2r.Execute(coeff, X_);
-          }
-        };
+    static void filter(const Stellarator<Real,ORDER>& S, const Comm& comm, Vector<ElemBasis>& f, Real sigma) {
+      Long dof = f.Dim() / S.NElem();
+      SCTL_ASSERT(f.Dim() == S.NElem() * dof);
+      for (Long i = 0; i < S.Nsurf()-1; 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 * 4;
+        const Long Np = Mp * ORDER * 4;
+
+        Vector<ElemBasis> f_(Nelem*dof, f.begin() + offset*dof, false);
+        Vector<Real> f_fourier = cheb2grid(f_, Mt, Mp, Nt, Np);
+        fourier_filter(f_fourier, Nt, Np, 0.25 * sigma, comm);
+        f_ = grid2cheb(f_fourier, Nt, Np, Mt, Mp);
+      }
+    };
 
-        Long dof = f.Dim() / S.NElem();
-        SCTL_ASSERT(f.Dim() == S.NElem() * dof);
-        for (Long i = 0; i < S.Nsurf()-1; 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 * 4;
-          const Long Np = Mp * ORDER * 4;
-
-          Vector<ElemBasis> f_(Nelem*dof, f.begin() + offset*dof, false);
-          Vector<Real> f_fourier = cheb2grid(f_, Mt, Mp, Nt, Np);
-          fourier_filter(f_fourier, Nt, Np, 0.25 * sigma, comm);
-          f_ = grid2cheb(f_fourier, Nt, Np, Mt, Mp);
-        }
-      };
+    static Vector<Real> cheb2grid(const Stellarator<Real,ORDER>& S, const Vector<ElemBasis>& f) {
+      const Long Nnodes = ElemBasis::Size();
+      const Long Nelem = S.NElem();
+      const Long dof = f.Dim() / Nelem;
+      SCTL_ASSERT(Nelem * dof == f.Dim());
+
+      Vector<Real> f_fourier(dof * Nelem * (ORDER*ORDER*fourier_upsample*fourier_upsample));
+      for (Long i = 0; i < S.Nsurf(); i++) {
+        const Long Mt = S.NTor(i);
+        const Long Mp = S.NPol(i);
+        const Long offset = S.ElemDsp(i);
+        const Long Nt = Mt * ORDER * fourier_upsample;
+        const Long Np = Mp * ORDER * fourier_upsample;
+
+        const Vector<ElemBasis> f_(Mt*Mp*dof, (Iterator<ElemBasis>)f.begin() + offset*dof, false);
+        Vector<Real> f_fourier_(dof*Nt*Np, f_fourier.begin() + dof*offset * (ORDER*ORDER*fourier_upsample*fourier_upsample), false);
+        f_fourier_ = cheb2grid(f_, Mt, Mp, Nt, Np);
+        SCTL_ASSERT(f_fourier_.Dim() == dof*Nt*Np);
+      }
+      return f_fourier;
+    }
+    static Vector<ElemBasis> grid2cheb(const Stellarator<Real,ORDER>& S, const Vector<Real>& f_fourier) {
+      const Long Nnodes = ElemBasis::Size();
+      const Long Nelem = S.NElem();
+      const Long dof = f_fourier.Dim() / (Nelem * (ORDER*ORDER*fourier_upsample*fourier_upsample));
+      SCTL_ASSERT(dof * Nelem * (ORDER*ORDER*fourier_upsample*fourier_upsample) == f_fourier.Dim());
+
+      Vector<ElemBasis> f(Nelem * dof);
+      for (Long i = 0; i < S.Nsurf(); i++) {
+        const Long Mt = S.NTor(i);
+        const Long Mp = S.NPol(i);
+        const Long offset = S.ElemDsp(i);
+        const Long Nt = Mt * ORDER * fourier_upsample;
+        const Long Np = Mp * ORDER * fourier_upsample;
+
+        Vector<ElemBasis> f_(Mt*Mp*dof, f.begin() + offset*dof, false);
+        const Vector<Real> f_fourier_(dof*Nt*Np, (Iterator<Real>)f_fourier.begin() + dof*offset * (ORDER*ORDER*fourier_upsample*fourier_upsample), false);
+        f_ = grid2cheb(f_fourier_, Nt, Np, Mt, Mp);
+        SCTL_ASSERT(f_.Dim() == Mt*Mp*dof);
+      }
+      return f;
+    }
+
+    template <class Real, class GradOp> static Long GradientDescent(GradOp& grad_op, Eigen::VectorXd& x, Real& fx, Long max_iter, Real tol) {
+      Real dt = 0.1;
+      for (Long iter = 0; iter < max_iter; iter++) {
+        Eigen::VectorXd grad(x.size());
+        fx = grad_op(x, grad);
+        { // Update dt
+          Eigen::VectorXd grad_(x.size());
+          Eigen::VectorXd x1 = x - grad * dt * 0.5;
+          Eigen::VectorXd x2 = x - grad * dt * 1.0;
+          Real fx1 = grad_op(x1, grad_);
+          Real fx2 = grad_op(x2, grad_);
+
+          { // Calculate optimal step size dt
+            Real a = 2*fx - 4*fx1 + 2*fx2;
+            Real b =-3*fx + 4*fx1 -   fx2;
+            Real c = fx;
+            Real s = -b/(2*a);
+            dt *= s;
+
+            Real fx_ = a*s*s + b*s + c;
+            std::cout<<"g = "<<fx_<<' ';
+            std::cout<<fx<<' ';
+            std::cout<<fx1<<' ';
+            std::cout<<fx2<<' ';
+            std::cout<<dt<<'\n';
+          }
+        }
+        x = x - grad * dt;
+        if (fx < tol) return iter;
+      }
+      return max_iter;
+    }
+
+  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;
+
+      Vector<Real> X_fourier(x.size());
+      for (Long i = 0; i < x.size(); i++) { // Set X_fourier
+        X_fourier[i] = x(i);
+      }
+      Vector<ElemBasis> X = grid2cheb(S_, X_fourier);
 
       Real g;
       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];
-        }
-      }
-      Stellarator<Real,ORDER> SS; //////////////////////////
-      { // Update S <-- filter(S)
-        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);
-          X[i*COORD_DIM+1] = S_.Elem(i, 1);
-          X[i*COORD_DIM+2] = S_.Elem(i, 2);
-        }
-        SS = S_;
-        filter(S_, comm, X, 0.1);
-        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];
+          S_.Elem(i,0)[j] = X[i*COORD_DIM+0][j];
+          S_.Elem(i,1)[j] = X[i*COORD_DIM+1][j];
+          S_.Elem(i,2)[j] = X[i*COORD_DIM+2][j];
         }
       }
       Vector<ElemBasis> dgdnu = Stellarator<Real,ORDER>::compute_gradient(S_, pressure_, flux_tor_, flux_pol_, &g);
-      Vector<ElemBasis> dXdt(Nelem*COORD_DIM);
-      { // 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];
+      //Vector<ElemBasis> dgdnu = Stellarator<Real,ORDER>::compute_pressure_jump(S_, pressure_, flux_tor_, flux_pol_, &g);
+
+      Vector<Real> dXdt_fourier;
+      { // Set grad
+        //filter(S_, comm, dgdnu, 0.1);
+        //Vector<Real> dgdnu_fourier = cheb2grid(S_, dgdnu);
+
+        { // deprecate
+          Vector<ElemBasis> dXdt(Nelem*COORD_DIM);
+          { // 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 dXdt
+            filter(S_, comm, dXdt, 0.1);
+          }
+          dXdt_fourier = cheb2grid(S_, dXdt);
+          SCTL_ASSERT(grad.size() == dXdt_fourier.Dim());
+          for (Long i = 0; i < grad.size(); i++) { // Set grad
+            grad(i) = dXdt_fourier[i];
           }
-        }
-        //filter(S_, comm, dXdt, 0.1);
-      }
-      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];
         }
       }
 
@@ -5169,14 +5244,10 @@ template <class Real, Integer ORDER=10> class MHDEquilib {
       }
       if (1) { // Write VTU
         VTUData vtu;
+        Vector<ElemBasis> dXdt = grid2cheb(S_, dXdt_fourier);
         vtu.AddElems(S_.GetElemList(), dXdt, ORDER);
         vtu.WriteVTK("dXdt"+std::to_string(iter), comm);
       }
-      if (1) { // Write VTU
-        VTUData vtu;
-        vtu.AddElems(SS.GetElemList(), dgdnu, ORDER);
-        vtu.WriteVTK("S"+std::to_string(iter), comm);
-      }
       std::cout<<"iter = "<<iter<<"    g = "<<g<<'\n';
 
       iter++;
@@ -5184,39 +5255,39 @@ template <class Real, Integer ORDER=10> class MHDEquilib {
     }
 
     static void ComputeEquilibrium(MHDEquilib& mhd_equilib) {
+      Comm comm = Comm::World();
       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-8;
-      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];
+      Eigen::VectorXd x;
+      { // Set x
+        Vector<ElemBasis> X(Nelem * COORD_DIM);
+        for (Long i = 0; i < Nelem; i++) { // Set x
+          X[i*COORD_DIM+0] = mhd_equilib.S_.Elem(i,0);
+          X[i*COORD_DIM+1] = mhd_equilib.S_.Elem(i,1);
+          X[i*COORD_DIM+2] = mhd_equilib.S_.Elem(i,2);
+        }
+        Vector<Real> X_fourier = cheb2grid(mhd_equilib.S_, X);
+        x.resize(X_fourier.Dim());
+        for (Long i = 0; i < X_fourier.Dim(); i++) {
+          x(i) = X_fourier[i];
         }
       }
 
       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);
-        }
+      if (0) {
+        LBFGSpp::LBFGSParam<Real> param;
+        param.max_iterations = 100;
+        param.epsilon = 1e-8;
+        LBFGSpp::LBFGSSolver<Real> solver(param);
+        Integer niter = solver.minimize(mhd_equilib, x, fx);
+      } else {
+        Integer niter = GradientDescent(mhd_equilib, x, fx, 100, 1e-8);
+      }
+      { // Set x
+        // TODO
       }
-
-      std::cout << niter << " iterations" <<'\n';
-      std::cout << "f(x) = " << fx <<'\n';
     }
 
     static void test() {