Write new function as per tests.

Author: mvousden

fidimag/common/skyrmion_number.py

@@ -3,9 +3,9 @@
 import numpy as np
 
 
-def skyrmion_number_centre_slice(sim):
+def skyrmion_number_slice(sim, at=None, zIndex=None):
     """
-    Returns the skyrmion number calculated on the x-y plane at the central z
+    Returns the skyrmion number calculated on the x-y plane at a specific z
     co-ordinate.
 
     The current fidimag skyrmion number function finds the skyrmion number for
@@ -16,24 +16,54 @@ def skyrmion_number_centre_slice(sim):
     Arguments:
       sim: LLG object with cuboidal mesh.
 
+    Optional Arguments:
+      If neither of these arguments are specified, this function raises a
+      ValueError.
+
+      at: None, "top", "centre", or "bottom" denoting where to slice in
+          British English. https://en.wikipedia.org/wiki/Strike_It_Lucky
+      zIndex: Integer (not a float), or None denoting the index at which to
+          slice. Overrides "at" if not None. Bound by zIndex~[0, sim.mesh.nz-1]
+
     Returns:
       skyrmionNumber (a float)
     """
 
+    # Deal with inputs.
+    if zIndex is None:
+        if at is "bottom":
+            zIndex = 0
+
+        elif at is "centre":
+            # Find the layer number that corresponds to the centre of the mesh,
+            # preferring the lower layer in a tie.
+            zIndex = sim.mesh.nz / 2.
+            if sim.mesh.nz % 2 == 0:
+                zIndex -= 0.5
+            zIndex -= 0.5
+
+        elif at is "top":
+            zIndex = sim.mesh.nz - 1  # Indeces start at 0.
+
+        else:
+            raise ValueError("[skyrmion_number_slice]: Either "
+                             "'at=bottom|centre|top' or 'zIndex'=[integer] "
+                             "must be passed as arguments.")
+
+    elif zIndex > sim.mesh.nz - 1 or zIndex < 0:
+        raise ValueError("[skyrmion_number_slice]: 'zIndex={}' must be an "
+                         "integer between 0 and sim.mesh.nz - 1, which is "
+                         "'{}' for this simulation object."
+                         .format(zIndex, sim.mesh.nz - 1))
+
     # Find the "length" of a slice in terms of the spin array.
     xyLength = sim.mesh.nx * sim.mesh.ny * 3  # The 3 represents (x,y,z).
 
-    # Find the layer number that corresponds to the centre of the mesh,
-    # preferring the lower layer in a tie.
-    zCentre = sim.mesh.nz / 2.
-    if sim.mesh.nz % 2 == 0:
-        zCentre -= 0.5
-    zCentre -= 0.5
-
-    # Obtain slice of spins cleverly. This also works if the domain is flat.
-    spinSlice = sim.spin[int(xyLength * zCentre):int(xyLength * (zCentre + 1))]
+    # Obtain slice of spins cleverly. This also works if the domain is flat
+    # (i.e. nz=1, zIndex=0)
+    spinSlice = sim.spin[int(xyLength * zIndex):int(xyLength * (zIndex + 1))]
 
-    # Compute the skyrmion number for our spin slice instead.
+    # Compute the skyrmion number for our spin slice.
     if sim._micromagnetic is True:
         return fidimag.extensions.micro_clib.compute_skyrmion_number(\
                spinSlice, sim._skx_number, sim.mesh.nx, sim.mesh.ny,