Add working lazy FMM and generated code up to high order

rpep · rpep · commit bfc5d9423898 · 2019-07-17T10:36:21.000+01:00
diff --git a/fidimag/atomistic/demag.py b/fidimag/atomistic/demag.py
@@ -1,6 +1,12 @@
 import fidimag.extensions.dipolar as clib
 import numpy as np
 from .energy import Energy
+import numpy as np
+import fidimag
+from fidimag.atomistic.energy import Energy
+import fidimag.extensions.fmm as fmm
+import time
+
 
 
 class Demag(Energy):
@@ -94,3 +100,39 @@ def compute_energy(self):
         self.energy /= self.scale
 
         return energy / self.scale
+
+
+class DemagFMM(Energy): 
+    def __init__(self, order, ncrit, theta, name="DemagFMM"):
+        self.name = name
+        assert order > 0, "Order must be 1 or higher"
+        self.order = order
+        assert ncrit >= 2, "ncrit must be greater than 1."
+        self.ncrit = ncrit
+        assert theta >= 0.0, "theta must be >= 0.0"
+        self.theta = theta
+
+    def setup(self, mesh, spin, mu_s, mu_s_inv):
+        super(DemagFMM, self).setup(mesh, spin, mu_s, mu_s_inv)
+        self.n = mesh.n
+        print(mesh.coordinates)
+        self.m_temp = spin.copy()
+        self.m_temp[0::3] *= self.mu_s
+        self.m_temp[1::3] *= self.mu_s
+        self.m_temp[2::3] *= self.mu_s
+        self.fmm = fmm.FMM(self.n, self.ncrit, self.theta,
+                           self.order,
+                           mesh.coordinates * mesh.unit_length,
+                           self.m_temp)
+
+    def compute_field(self, t=0, spin=None):
+        self.m_temp[:] = spin if spin is not None else self.spin
+        self.m_temp[0::3] *= self.mu_s
+        self.m_temp[1::3] *= self.mu_s
+        self.m_temp[2::3] *= self.mu_s
+
+        self.field[:] = 0.0
+        #self.fmm.set(self.m_temp)
+        self.fmm.compute_field(self.theta, self.field)
+        self.field *= 1e-7
+        return self.field
diff --git a/fidimag/atomistic/fmmlib/calculate.cpp b/fidimag/atomistic/fmmlib/calculate.cpp
@@ -5,6 +5,7 @@
 #include <iostream>
 #include <stack>
 #include <cmath>
+#include <algorithm>
 #include<cstdio>
 
 void P2P(double x, double y, double z, double mux, double muy, double muz, double *F) {
@@ -30,9 +31,10 @@ void P2P_noatomic(double x, double y, double z, double mux, double muy, double m
   F[2] += (3*mu_dot_r * z / R5 - muz / R3);
 }
 
+
 void evaluate_P2M(std::vector<Particle> &particles, std::vector<Cell> &cells,
 		          size_t cell, size_t ncrit, size_t exporder) {
-  std::cout << "Nparticles = " << particles.size() << std::endl;
+  // std::cout << "Nparticles = " << particles.size() << std::endl;
   double *M = new double[Nterms(exporder+1)]();
   //#pragma omp for
   for(size_t c = 0; c < cells.size(); c++) {
@@ -67,14 +69,14 @@ void evaluate_M2M(std::vector<Particle> &particles, std::vector<Cell> &cells,
   #pragma omp for
   for (size_t c = cells.size() - 1; c > 0; c--) {
     size_t p = cells[c].parent;
-    std::cout << "M2M: " << c << " to " << p << std::endl;
+    // std::cout << "M2M: " << c << " to " << p << std::endl;
     double dx = cells[p].x - cells[c].x;
     double dy = cells[p].y - cells[c].y;
     double dz = cells[p].z - cells[c].z;
     M2M(dx, dy, dz, cells[c].M, cells[p].M, exporder);
   }
 
-  std::cout << "evalm2m_cpp: " << cells[0].M[0] << "," << cells[0].M[1] << "," << cells[0].M[2] << "," << cells[0].M[3] << std::endl;
+  // std::cout << "evalm2m_cpp: " << cells[0].M[0] << "," << cells[0].M[1] << "," << cells[0].M[2] << "," << cells[0].M[3] << std::endl;
 }
 
 
@@ -210,12 +212,12 @@ void evaluate_M2L_lazy(std::vector<Cell> &cells,
 }
 
 void evaluate_P2P_lazy(std::vector<Cell> &cells, std::vector<Particle> &particles,
-                       std::vector<std::pair<size_t, size_t>> &P2P_list, std::vector<double> &F) {
+                       std::vector<std::pair<size_t, size_t>> &P2P_list, double *F) {
    #pragma omp for
    for(size_t i = 0; i < P2P_list.size(); i++) {
        size_t A = P2P_list[i].first;
        size_t B = P2P_list[i].second;
-       P2P_Cells(A, B, cells, particles, F.data());
+       P2P_Cells(A, B, cells, particles, F);
    }
 }
 
@@ -274,3 +276,54 @@ void evaluate_direct(std::vector<Particle> &particles, double *F, size_t n) {
       }
   }
 }
+
+
+void evaluate_approx(std::vector<Particle> &particles, std::vector<Cell> &cells,
+                     size_t ncrit, double theta, size_t order, double *F) {
+    evaluate_P2M(particles, cells, 0, ncrit, order);
+    evaluate_M2M(particles, cells, order);
+    interact_dehnen(0, 0, cells, particles, theta, order, ncrit, F);
+    evaluate_L2L(cells, order);
+    evaluate_L2P(particles, cells, F, ncrit, order);
+}
+
+void evaluate_approx_lazy(std::vector<Particle> &particles, std::vector<Cell> &cells,
+                          size_t ncrit, size_t order, double *F,
+                          std::vector<std::pair<size_t, size_t>> &M2L_list,
+                          std::vector<std::pair<size_t, size_t>> &P2P_list) {
+    #pragma omp parallel
+    evaluate_P2M(particles, cells, 0, ncrit, order);
+
+    evaluate_M2M(particles, cells, order);
+    #pragma omp barrier
+    #pragma omp parallel
+    {
+      evaluate_M2L_lazy(cells,M2L_list,order);
+      evaluate_P2P_lazy(cells, particles, P2P_list, F);
+      #pragma omp barrier
+      evaluate_L2L(cells, order);
+      #pragma omp barrier
+      evaluate_L2P(particles, cells, F, ncrit, order);
+    }
+
+}
+
+
+void build_interaction_lists(std::vector<std::pair<size_t, size_t>> &M2L_list,
+                             std::vector<std::pair<size_t, size_t>> &P2P_list,
+                             std::vector<Cell> &cells,
+                             std::vector<Particle> &particles,
+                             double theta,
+                             size_t order,
+                             size_t ncrit
+    ) {
+
+    std::sort(M2L_list.begin(), M2L_list.end(),
+           [](std::pair<size_t, size_t> &left, std::pair<size_t, size_t> &right) {
+                return left.first < right.first;
+               }
+           );
+
+    std::cout << "M2L_list size = " << M2L_list.size() << std::endl;
+    std::cout << "P2P_list size = " << P2P_list.size() << std::endl;
+}
diff --git a/fidimag/atomistic/fmmlib/calculate.hpp b/fidimag/atomistic/fmmlib/calculate.hpp
@@ -16,6 +16,14 @@
 
 void P2P(double x, double y, double z, double mux, double muy, double muz, double *F);
 
+void evaluate_approx(std::vector<Particle> &particles, std::vector<Cell> &cells,
+                     size_t ncrit, double theta, size_t order, double *F);
+
+void evaluate_approx_lazy(std::vector<Particle> &particles, std::vector<Cell> &cells,
+                          size_t ncrit, size_t order, double *F,
+                          std::vector<std::pair<size_t, size_t>> &M2L_list,
+                          std::vector<std::pair<size_t, size_t>> &P2P_list);
+
 void evaluate_P2M(std::vector<Particle> &particles, std::vector<Cell> &cells,
 		  size_t cell, size_t ncrit, size_t exporder);
 
@@ -42,8 +50,6 @@ void interact_dehnen_lazy(const size_t A, const size_t B, const std::vector<Cell
 void P2P_Cells(size_t A, size_t B, std::vector<Cell> &cells,
 	 			       std::vector<Particle> &particles, double *F);
 
-void evaluate_P2P_lazy(std::vector<Cell> &cells,
-                      std::vector<std::pair<size_t, size_t>> &P2P_list);
 
 void evaluate_M2L_lazy(std::vector<Cell> &cells,
                      std::vector<std::pair<size_t, size_t>> &M2L_list,
@@ -53,4 +59,13 @@ void evaluate_M2L_lazy(std::vector<Cell> &cells,
                     std::vector<std::pair<size_t, size_t>> &M2L_list, size_t order);
 
 void evaluate_P2P_lazy(std::vector<Cell> &cells, std::vector<Particle> &particles,
-                    std::vector<std::pair<size_t, size_t>> &P2P_list, std::vector<double> &F);
+                       std::vector<std::pair<size_t, size_t>> &P2P_list, double *F);
+
+void build_interaction_lists(std::vector<std::pair<size_t, size_t>> &M2L_list,
+                             std::vector<std::pair<size_t, size_t>> &P2P_list,
+                             std::vector<Cell> &cells,
+                             std::vector<Particle> &particles,
+                             double theta,
+                             size_t order,
+                             size_t ncrit
+    );
diff --git a/fidimag/atomistic/fmmlib/fmm.pyx b/fidimag/atomistic/fmmlib/fmm.pyx
@@ -3,7 +3,7 @@
 from fidimag.atomistic.energy import Energy
 from libcpp.vector cimport vector
 from libcpp.utility cimport pair
-
+from libcpp.algorithm cimport sort
 cimport numpy as np
 import numpy as np
 
@@ -31,6 +31,14 @@ cdef extern from "tree.hpp":
 
 
 cdef extern from "calculate.hpp":
+    void evaluate_approx(vector[Particle] particles, vector[Cell] cells,
+                     size_t ncrit, double theta, size_t order, double *F)
+
+    void evaluate_approx_lazy(vector[Particle] particles, vector[Cell] cells,
+                          size_t ncrit, size_t order, double *F,
+                          vector[pair[size_t, size_t]] M2L_list,
+                          vector[pair[size_t, size_t]] P2P_list)
+
     void evaluate_P2M(vector[Particle] particles,
                       vector[Cell] cells,
                       size_t cell,
@@ -62,8 +70,23 @@ cdef extern from "calculate.hpp":
                          double *F,
                          size_t n)
 
+    void interact_dehnen_lazy(size_t A, size_t B,
+                              vector[Cell] cells,
+                              vector[Particle] particles,
+                              double theta, size_t order,
+                              const size_t ncrit,
+                              vector[pair[size_t, size_t]] M2L_list,
+                              vector[pair[size_t, size_t]] P2P_list)
 
 
+    void build_interaction_lists(vector[pair[size_t, size_t]] M2L_list,
+                                 vector[pair[size_t, size_t]]P2P_list,
+                                 vector[Cell] cells,
+                                 vector[Particle] particles,
+                                 double theta,
+                                 size_t order,
+                                 size_t ncrit
+    )
 
 
 cdef class FMM:
@@ -75,18 +98,23 @@ cdef class FMM:
     cdef Cell root
     cdef vector[double] M
     cdef vector[double] L
+    cdef vector[pair[size_t, size_t]] M2L_list
+    cdef vector[pair[size_t, size_t]] P2P_list
     cdef size_t Msize
     cdef size_t Lsize
     cdef public double [:, :] coords
+    cdef double theta
 
-    def __cinit__(self, size_t n, size_t ncrit, size_t order, double [:, :] coords, double [:] mu):
+    def __cinit__(self, size_t n, size_t ncrit, double theta, size_t order, double [:, :] coords, double [:] mu):
+        if order > 11:
+            raise ValueError("Order needs to be < 12")
         # self.particles = vector[Particle]
-
+        self.theta = theta
         # Don't remove this line, or the memory goes out of scope!
         self.coords = coords
         self.ncrit = ncrit
         self.order = order
-        print('FMM Order = {}'.format(order))
+        # print('FMM Order = {}'.format(order))
         xs = np.asarray(self.coords[:, 0])
         ys = np.asarray(self.coords[:, 1])
         zs = np.asarray(self.coords[:, 2])
@@ -108,16 +136,7 @@ cdef class FMM:
                               &mu[3*i])
             self.particles.push_back(self.p)
         self.cells = build_tree(self.particles, self.root, ncrit, order)
-        print(f"Cython: cells.size() {self.cells.size()}")
-
-
-        print(f"r[0] = {self.particles[0].r[0]}, {self.particles[0].r[1]}, {self.particles[0].r[2]}")
-
-        print(f"mu[0] = {self.particles[0].mu[0]}, {self.particles[0].mu[1]}, {self.particles[0].mu[2]}")
-
-        # Need to allocate memory for the M and L arrays.
-        #vector[double] M(cells.size() * (Nterms(order) - Nterms(0)), 0.0)
-        #vector[double] L(cells.size() * Nterms(order - 1), 0.0)
+        print(f"DemagFMM tree built with {self.cells.size()} cells")
 
         self.Msize = Nterms(order) - Nterms(0)
         self.Lsize = Nterms(order - 1)
@@ -128,6 +147,12 @@ cdef class FMM:
             self.cells[i].M = &self.M[i*self.Msize]
             self.cells[i].L = &self.L[i*self.Lsize]
 
+        # vector[pair[size_t, size_t]] M2L_list
+        # vector[pair[size_t, size_t]] P2P_list
+        print("Setting up interaction list")
+        # interact_dehnen_lazy(0, 0, self.cells, self.particles, theta, order, ncrit, self.M2L_list, self.P2P_list)
+        build_interaction_lists(self.M2L_list, self.P2P_list, self.cells, self.particles, self.theta, order, ncrit)
+        print("Done")
 
     def P2M(self):
         evaluate_P2M(self.particles, self.cells, 0, self.ncrit, self.order)
@@ -137,24 +162,31 @@ cdef class FMM:
 
 
     def compute_field(self, double theta, double [:] F):
-        F[:] = 0
-        print('compute field')
-        print("P2M starting")
-        evaluate_P2M(self.particles, self.cells, 0, self.ncrit, self.order)
-        print("M2M starting")
-        evaluate_M2M(self.particles, self.cells, self.order)
-
-        print(f"mu[0] = {self.particles[0].mu[0]}, {self.particles[0].mu[1]}, {self.particles[0].mu[2]}")
-
-        print(np.max(self.M))
-
-        print("interact_dehnen starting")
-        interact_dehnen(0, 0, self.cells, self.particles, theta, self.order, self.ncrit, &F[0])
-        print("L2L starting")
-        evaluate_L2L(self.cells, self.order)
-        print("L2P starting")
-        evaluate_L2P(self.particles, self.cells, &F[0], self.ncrit, self.order)
-        print(F[0])
+        #print("Computing field...")
+        for i in range(self.Msize * self.cells.size()):
+            self.M[i] = 0.0
+        for i in range(self.Lsize * self.cells.size()):
+            self.L[i] = 0.0
+        for i in range(3*self.n):
+            F[i] = 0.0
+        #print('compute field')
+        #print("P2M starting")
+        # evaluate_P2M(self.particles, self.cells, 0, self.ncrit, self.order)
+        # #print("M2M starting")
+        # evaluate_M2M(self.particles, self.cells, self.order)
+
+        # #print(f"mu[0] = {self.particles[0].mu[0]}, {self.particles[0].mu[1]}, {self.particles[0].mu[2]}")
+
+        # #print(np.max(self.M))
+
+        # #print("interact_dehnen starting")
+        # interact_dehnen(0, 0, self.cells, self.particles, theta, self.order, self.ncrit, &F[0])
+        # #print("L2L starting")
+        # evaluate_L2L(self.cells, self.order)
+        # #print("L2P starting")
+        # evaluate_L2P(self.particles, self.cells, &F[0], self.ncrit, self.order)
+        # evaluate_approx(self.particles, self.cells, self.ncrit, theta, self.order, &F[0])
+        evaluate_approx_lazy(self.particles, self.cells, self.ncrit, self.order, &F[0], self.M2L_list, self.P2P_list)
 
     def compute_field_exact(self, double [:] F_exact):
         evaluate_direct(self.particles, &F_exact[0], self.n)
diff --git a/fidimag/atomistic/fmmlib/generate.py b/fidimag/atomistic/fmmlib/generate.py
@@ -1,6 +1,6 @@
 import fmmgen
 
-order = 5
+order = 13
 
 fmmgen.generate_code(order, "operators", CSE=True, cython_wrapper=False, potential=False, field=True)
 
diff --git a/fidimag/atomistic/fmmlib/operators.cpp b/fidimag/atomistic/fmmlib/operators.cpp
diff --git a/fidimag/atomistic/fmmlib/operators.h b/fidimag/atomistic/fmmlib/operators.h
