Refactor SimpleDemag so that it is a subclass of Energy

Author: rpep

fidimag/micro/simple_demag.py

@@ -1,7 +1,8 @@
 import numpy as np
+from .energy import Energy
 
 
-class SimpleDemag(object):
+class SimpleDemag(Energy):
     """
     Demagnetising field for thin films in the i-direction.
     Hj = Hk = 0 and Hi = - Mi.
@@ -12,29 +13,12 @@ def __init__(self, Nx=0, Ny=0.5, Nz=0.5, name='SimpleDemag'):
         field_strength is Ms by default
 
         """
-        
+
         self.Nx = Nx
         self.Ny = Ny
         self.Nz = Nz
-        
         self.name = name
 
-    def setup(self, mesh, spin, Ms):
-        self.mesh = mesh
-        self.spin = spin
-        self.n = mesh.n
-
-        self.Ms = Ms
-        self.Ms_long = np.zeros(3 * mesh.n)
-
-        self.Ms_long.shape = (3, -1)
-        for i in range(mesh.n):
-            self.Ms_long[:, i] = Ms[i]
-
-        self.Ms_long.shape = (-1,)
-        self.field = np.zeros(3 * self.n)
-        #self.field[:] = helper.init_vector(self.H0, self.mesh)
-
 
     def compute_field(self, t=0):
         self.spin.shape = (3,-1)
@@ -49,9 +33,7 @@ def compute_field(self, t=0):
     def compute_energy(self):
 
         mu_0 = 4*np.pi*1e-7
-
-        sf = 0.5* self.field * self.spin * self.Ms_long * mu_0
-
-        energy = -np.sum(sf)
+        sf = -0.5 * self.field * self.spin * mu_0
+        energy = np.sum(sf.reshape(-1, 3), axis=1) * self.Ms
 
         return energy * self.mesh.cellsize