Merge pull request #116 from computationalmodelling/helper-cython

Refactor Helper functions and Zeeman for better performance

Author: davidcortesortuno

.gitignore

@@ -5,6 +5,7 @@
 *.swp
 *.pdf
 *.tmp
+*.so
 /local/
 
 # ignore automatically generated cython files

doc/core_eqs.rst

@@ -18,7 +18,6 @@ described at a semi-classical level.
 
   Interactions between magnetic moments are specified using the Heisenberg
   formalism.  
-|
 
 * Micromagnetics
 

doc/index.rst

@@ -12,7 +12,6 @@ The code for Fidimag is available under an open source license on `GitHub
 <https://github.com/computationalmodelling/fidimag>`_.
 
 Contents:
-
 .. toctree::
    :maxdepth: 2
    :caption: Installation Instructions
@@ -31,6 +30,7 @@ Contents:
    ipynb/isolated_skyrmion
    ipynb/spin-polarised-current-driven-skyrmion
    ipynb/spin-waves-in-periodic-system
+   ipynb/FMR-stdprob
 
 .. toctree::
    :maxdepth: 2
@@ -46,4 +46,3 @@ Contents:
 
    nebm
 
-.. _GitHub

fidimag/common/helper.py

@@ -8,71 +8,7 @@
 from mpl_toolkits.mplot3d import Axes3D
 from matplotlib.colors import colorConverter
 from matplotlib.collections import PolyCollection, LineCollection
-
-
-def normalise(a):
-    """
-    normalise the given array a
-    """
-    a.shape = (-1, 3)
-    b = np.sqrt(a[:, 0] ** 2 + a[:, 1] ** 2 + a[:, 2] ** 2)
-    ids = (b == 0)
-    b[ids] = 1.0
-    a[:, 0] /= b
-    a[:, 1] /= b
-    a[:, 2] /= b
-    a.shape = (-1,)
-
-
-def init_vector(m0, mesh, norm=False, *args):
-    n = mesh.n
-    field = np.zeros((n, 3))
-
-    if isinstance(m0, list) or isinstance(m0, tuple):
-        field[:, :] = m0
-        field = np.reshape(field, 3 * n, order='C')
-
-    elif hasattr(m0, '__call__'):
-        # Check only once that the function returns appropriately...
-        v = m0(mesh.coordinates[0], *args)
-        if len(v) != 3:
-            raise Exception(
-                'The length of the value in init_vector method must be 3.')
-        for i in range(n):
-            field[i, :] = m0(mesh.coordinates[i], *args)
-        field = np.reshape(field, 3 * n, order='C')
-
-    elif isinstance(m0, np.ndarray):
-        if m0.shape == (3, ):
-            field[:] = m0  # broadcasting
-        else:
-            field.shape = (-1)
-            field[:] = m0  # overwriting the whole thing
-    field.shape = (-1,)
-    if norm:
-        normalise(field)
-
-    return field
-
-
-def init_scalar(value, mesh, *args):
-    n = mesh.n
-    mesh_v = np.zeros(n)
-
-    if isinstance(value, (int, float)):
-        mesh_v[:] = value
-    
-    elif hasattr(value, '__call__'):
-        for i in range(n):
-            mesh_v[i] = value(mesh.coordinates[i], *args)
-
-    elif isinstance(value, np.ndarray):
-        if value.shape == mesh_v.shape:
-            mesh_v[:] = value[:]
-        else:
-            raise ValueError("Array size must match the mesh size")
-
-    return mesh_v
+from fidimag.extensions.common_clib import normalise, init_scalar, init_vector, init_vector_func_fast
 
 
 def extract_data(mesh, npys, pos, comp='x'):
@@ -128,15 +64,15 @@ def plot_m(mesh, npy, comp='x', based=None):
         data = data - based
 
     data.shape = (-1, 3)
-    m = data[:,cmpi]
+    m = data[:, cmpi]
 
     nx = mesh.nx
     ny = mesh.ny
     nz = mesh.nz
 
     m.shape = (nz, ny, nx)
 
-    m2 = m[0,:,:]
+    m2 = m[0, :, :]
 
     fig = plt.figure()
     # norm=color.Normalize(-1,1)
@@ -207,7 +143,7 @@ def plot_energy_3d(name, key_steps=50, filename=None):
     if each_n_step < 1:
         each_n_step = 1
 
-    cc = lambda arg: colorConverter.to_rgba(arg, alpha=0.6)
+    def cc(arg): return colorConverter.to_rgba(arg, alpha=0.6)
     colors = [cc('r'), cc('g'), cc('b'), cc('y')]
     facecolors = []
     line_data = []
@@ -241,5 +177,6 @@ def plot_energy_3d(name, key_steps=50, filename=None):
 
 
 def compute_RxRy(mesh, spin, nx_start=0, nx_stop=-1, ny_start=0, ny_stop=-1):
-    res = clib.compute_RxRy(spin, mesh.nx, mesh.ny, mesh.nz, nx_start, nx_stop, ny_start, ny_stop)
+    res = clib.compute_RxRy(spin, mesh.nx, mesh.ny,
+                            mesh.nz, nx_start, nx_stop, ny_start, ny_stop)
     return res

fidimag/common/lib/common_clib.pyx

@@ -1,3 +1,5 @@
+cimport numpy as np
+import numpy as np
 
 # -----------------------------------------------------------------------------
 
@@ -33,7 +35,7 @@ cdef extern from "common_clib.h":
                          double *mxH, double *mxmxH, double *mxmxH_last, double *tau,
                          int* pins, int n)
 
-    void sd_compute_step (double *spin, double *spin_last, double *magnetisation, 
+    void sd_compute_step (double *spin, double *spin_last, double *magnetisation,
                           double *field,
                           double *mxH, double *mxmxH, double *mxmxH_last, double *tau,
                           int *pins, int n, int counter, double tmin, double tmax)
@@ -128,8 +130,103 @@ def compute_sd_step(double [:] spin,
                     int [:] pins,
                     n, counter, tmin, tmax):
 
-    sd_compute_step(&spin[0], &spin_last[0], &magnetisation[0], 
+    sd_compute_step(&spin[0], &spin_last[0], &magnetisation[0],
                     &field[0], &mxH[0],
                     &mxmxH[0], &mxmxH_last[0], &tau[0], &pins[0],
                     n, counter, tmin, tmax
                     )
+
+def normalise(a):
+    """
+    normalise the given array a
+    """
+    a.shape = (-1, 3)
+    b = np.sqrt(a[:, 0] ** 2 + a[:, 1] ** 2 + a[:, 2] ** 2)
+    ids = (b == 0)
+    b[ids] = 1.0
+    a[:, 0] /= b
+    a[:, 1] /= b
+    a[:, 2] /= b
+    a.shape = (-1,)
+
+def init_scalar(value, mesh, *args):
+
+    n = mesh.n
+
+    mesh_v = np.zeros(n)
+
+    if isinstance(value, (int, float)):
+        mesh_v[:] = value
+    elif hasattr(value, '__call__'):
+        for i in range(n):
+            mesh_v[i] = value(mesh.coordinates[i], *args)
+
+    elif isinstance(value, np.ndarray):
+        if value.shape == mesh_v.shape:
+            mesh_v[:] = value[:]
+        else:
+            raise ValueError("Array size must match the mesh size")
+
+    return mesh_v
+
+def init_vector(m0, mesh, norm=False, *args):
+
+    n = mesh.n
+
+    spin = np.zeros((n, 3))
+
+    if isinstance(m0, list) or isinstance(m0, tuple):
+        spin[:, :] = m0
+        spin = np.reshape(spin, 3 * n, order='C')
+
+    elif hasattr(m0, '__call__'):
+        v = m0(mesh.coordinates[0], *args)
+        if len(v) != 3:
+            raise Exception(
+                'The length of the value in init_vector method must be 3.')
+        for i in range(n):
+            spin[i, :] = m0(mesh.coordinates[i], *args)
+        spin = np.reshape(spin, 3 * n, order='C')
+
+    elif isinstance(m0, np.ndarray):
+        if m0.shape == (3, ):
+            spin[:] = m0  # broadcasting
+        else:
+            spin.shape = (-1)
+            spin[:] = m0  # overwriting the whole thing
+
+    spin.shape = (-1,)
+
+    if norm:
+        normalise(spin)
+
+    return spin
+
+def init_vector_func_fast(m0, mesh, double[:] field, norm=False, *args):
+    """
+    An unsafe method of setting the field. Depends on
+    the setter code being memory safe.
+
+    m0 must be a Python function that takes the mesh
+    and field as arguments. Within that, the user
+    must handle evaluating the function at different
+    coordinate points. It needs to be able to handle
+    the spatial dependence itself, and write the
+    field valuse into the field array. This can
+    be written with Cython which will give much
+    better performance. For example:
+
+    from libc.math cimport sin
+
+    def fast_sin_init(mesh, double[:] field, *params):
+        t, axis, Bmax, fc = params
+        for i in range(mesh.n):
+            field[3*i+0] = Bmax * axis[0] * sin(fc*t)
+            field[3*i+1] = Bmax * axis[1] * sin(fc*t)
+            field[3*i+2] = Bmax * axis[2] * sin(fc*t)
+
+    """  
+    m0(mesh, field, *args)
+    if norm:
+        normalise(field)
+    return field

fidimag/micro/__init__.py

@@ -7,4 +7,3 @@
 from .dmi import DMI
 from .baryakhtar import LLBar, LLBarFull
 from .simple_demag import SimpleDemag
-

fidimag/micro/zeeman.py

@@ -1,7 +1,7 @@
 import numpy as np
-
 from fidimag.common.constant import mu_0
 import fidimag.common.helper as helper
+import inspect
 
 
 class Zeeman(object):
@@ -115,34 +115,19 @@ class TimeZeeman(Zeeman):
 
     """
     The time dependent external field, also can vary with space
-
-    The function time_fun must be a function which takes two arguments:
-
-    def time_fun(pos, t):
-        x, y, z = pos
-        # compute Bx, By, Bz as a function of x y, z and t.
-        Bx = ...
-        By = ...
-        Bz = ...
-        return (Bx, By, Bz)
-
-
     """
 
-    def __init__(self, H0, time_fun, name='TimeZeeman'):
-
+    def __init__(self, H0, time_fun, extra_args=[], name='TimeZeeman'):
         self.H0 = H0
         self.time_fun = time_fun
         self.name = name
         self.jac = True
+        self.extra_args = extra_args
 
     def setup(self, mesh, spin, Ms):
         super(TimeZeeman, self).setup(mesh, spin, Ms)
         self.H_init = self.field.copy()
 
     def compute_field(self, t=0, spin=None):
-        self.field[:] = helper.init_vector(self.time_fun,
-                                           self.mesh,
-                                           False,
-                                           t)
+        self.field[:] = self.H_init[:] * self.time_fun(t, *self.extra_args)
         return self.field

fidimag/user/README.md

@@ -0,0 +1,32 @@
+# User Extensions
+
+The user extensions directory is here to allow for 
+the user to be able to straightforwardly run compiled
+code within Fidimag for performance reasons. We
+consider this an advanced feature and do not recommend
+trying this unless you have experience writing and
+building C/Cython programs.
+
+Some of the energy classes perform callbacks to 
+user-supplied functions. Performance for this 
+is generally poor, as there is an overhead to 
+calling Python functions repeatedly. Hence,
+we place this folder here to allow you to expose
+functions written in Cython/C conveniently.
+
+An example has been supplied. We suggest copying
+the folder and modifying each of the files in 
+this. Please note that we have automated the
+building of the extensions, but you can only
+have a single Cython .pyx file per directory, 
+because a single Cython module is created in
+each folder. The module that is created will have 
+the name of this file.
+
+You do not explicitly need to write an __init__.py file;
+your extension will be importable immediately from
+fidimag.extensions.user.$FOLDERNAME
+
+The __init__.py file lets you do the slightly shorter:
+from fidimag.user.$FOLDERNAME import *
+

fidimag/user/example/__init__.py

@@ -0,0 +1 @@
+from fidimag.extensions.user.example import *

fidimag/user/example/example.pyx

@@ -0,0 +1,16 @@
+from libc.math cimport cos, sin
+
+def fast_sin_init(mesh, double[:] field, *params):
+    t, axis, Bmax, fc = params
+    for i in range(mesh.n):
+        field[3*i+0] = Bmax * axis[0] * sin(fc*t)
+        field[3*i+1] = Bmax * axis[1] * sin(fc*t)
+        field[3*i+2] = Bmax * axis[2] * sin(fc*t)
+
+def TimeZeemanFast_test_time_fun(mesh, double[:] field, *params):
+    cdef int i
+    t, frequency = params
+    for i in range(mesh.n):
+        field[3*i+0] = 0
+        field[3*i+1] = 0
+        field[3*i+2] = 10 * cos(frequency * t)