perf: Use gfortran -O3 and more numba for the Python

src/pydrex/diagnostics.py

@@ -15,10 +15,12 @@
     grain axis and the j-th external axis (in the global Eulerian frame).
 
 """
+
 import functools as ft
 from multiprocessing import Pool
 
 import numpy as np
+import numba as nb
 import scipy.linalg as la
 
 from pydrex import stats as _stats
@@ -207,7 +209,7 @@ def bingham_average(orientations, axis="a"):
     return mean_vector / la.norm(mean_vector)
 
 
-def finite_strain(deformation_gradient, **kwargs):
+def finite_strain(deformation_gradient, driver="ev"):
     """Extract measures of finite strain from the deformation gradient.
 
     Extracts the largest principal strain value and the vector defining the axis of
@@ -218,7 +220,7 @@ def finite_strain(deformation_gradient, **kwargs):
     # Get eigenvalues and eigenvectors of the left stretch (Cauchy-Green) tensor.
     B_λ, B_v = la.eigh(
         deformation_gradient @ deformation_gradient.transpose(),
-        driver=kwargs.pop("driver", "ev"),
+        driver=driver,
     )
     # Stretch ratio is sqrt(λ) - 1, the (-1) is so that undeformed strain is 0 not 1.
     return np.sqrt(B_λ[-1]) - 1, B_v[:, -1]
@@ -265,7 +267,11 @@ def symmetry(orientations, axis="a"):
 
 
 def misorientation_indices(
-    orientation_stack, system: _geo.LatticeSystem, bins=None, ncpus=None, pool=None,
+    orientation_stack,
+    system: _geo.LatticeSystem,
+    bins=None,
+    ncpus=None,
+    pool=None,
 ):
     """Calculate M-indices for a series of polycrystal textures.
 
@@ -312,7 +318,7 @@ def misorientation_index(orientations, system: _geo.LatticeSystem, bins=None):
     See `_geo.LatticeSystem` for supported systems.
 
     .. warning::
-        This method must be able to allocate an array of shape 
+        This method must be able to allocate an array of shape
         $ \frac{N!}{2(N-2)!}× M^{2} $
         for N the length of `orientations` and M the number of symmetry operations for
         the given `system`.
@@ -354,6 +360,7 @@ def coaxial_index(orientations, axis1="b", axis2="a"):
     return 0.5 * (2 - (P1 / (G1 + P1)) - (G2 / (G2 + P2)))
 
 
+@nb.njit(fastmath=True)
 def smallest_angle(vector, axis, plane=None):
     """Get smallest angle between a unit `vector` and a bidirectional `axis`.
 
@@ -363,13 +370,19 @@ def smallest_angle(vector, axis, plane=None):
 
     """
     if plane is not None:
-        _plane = np.asarray(plane)
-        _vector = np.asarray(vector) - _plane * np.dot(vector, _plane)
+        _vector = vector - plane * np.dot(vector, plane)
     else:
-        _vector = np.asarray(vector)
+        _vector = vector
     angle = np.rad2deg(
         np.arccos(
-            np.clip(np.dot(_vector, axis) / (la.norm(_vector) * la.norm(axis)), -1, 1)
+            np.clip(
+                np.asarray(  # https://github.com/numba/numba/issues/3469
+                    np.dot(_vector, axis)
+                    / (np.linalg.norm(_vector) * np.linalg.norm(axis))
+                ),
+                -1,
+                1,
+            )
         )
     )
     if angle > 90:

tools/perf_compare.sh

@@ -37,7 +37,7 @@ run_f90() {
     cp drex_forward_simpleshear.f90 "$OUTDIR"
     >/dev/null pushd "$OUTDIR"
     sed -i -E "s/size3 = [0-9]+/size3 = $SIZE3/" drex_forward_simpleshear.f90
-    gfortran drex_forward_simpleshear.f90 -o run_drex
+    gfortran -O3 drex_forward_simpleshear.f90 -o run_drex
     perf stat -r $N_RUNS -o drex.stat ./run_drex >/dev/null
     >/dev/null popd
 }
@@ -81,6 +81,7 @@ shift $(( $OPTIND - 1 ))
 mkdir -p "$OUTDIR"
 cat << EOF > "$OUTDIR/pydrex_forward_shear.py"
 import numpy as np
+import numba as nb
 
 from pydrex import core as _core
 from pydrex import diagnostics as _diagnostics
@@ -90,7 +91,7 @@ from pydrex import utils as _utils
 from pydrex import velocity as _velocity
 
 
-shear_direction = [0, 1, 0]
+shear_direction = np.array([0.0, 1.0, 0.0])
 strain_rate = 1e-4
 _, get_velocity_gradient = _velocity.simple_shear_2d("Y", "X", strain_rate)
 timestamps = np.linspace(0, 1e4, 201)
@@ -110,12 +111,17 @@ deformation_gradient = np.eye(3)
 θ_fse = np.empty_like(timestamps)
 θ_fse[0] = 45
 
+
+@nb.njit
+def get_position(t):
+    return np.zeros(3)
+
 for t, time in enumerate(timestamps[:-1], start=1):
     deformation_gradient = mineral.update_orientations(
         params,
         deformation_gradient,
         get_velocity_gradient,
-        pathline=(time, timestamps[t], lambda t: np.zeros(3)),
+        pathline=(time, timestamps[t], get_position),
     )
     _, fse_v = _diagnostics.finite_strain(deformation_gradient)
     θ_fse[t] = _diagnostics.smallest_angle(fse_v, shear_direction)