Merge pull request #162 from seismic-anisotropy/perf

Add simple f90 performance comparison script (single core)

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:

src/pydrex/minerals.py

@@ -6,6 +6,7 @@
     leading to an overall anisotropic orientation distribution
 
 """
+
 import io
 from dataclasses import dataclass, field
 from zipfile import ZipFile
@@ -20,6 +21,7 @@
 from pydrex import io as _io
 from pydrex import logger as _log
 from pydrex import tensors as _tensors
+from pydrex import utils as _utils
 
 OLIVINE_STIFFNESS = np.array(
     [
@@ -334,22 +336,6 @@ def update_orientations(
 
         # ===== Set up callables for the ODE solver and internal processing =====
 
-        def extract_vars(y):
-            # TODO: Check if we can avoid .copy() here.
-            deformation_gradient = y[:9].copy().reshape((3, 3))
-            orientations = (
-                y[9 : self.n_grains * 9 + 9]
-                .copy()
-                .reshape((self.n_grains, 3, 3))
-                .clip(-1, 1)
-            )
-            # TODO: Check if we can avoid .copy() here.
-            fractions = (
-                y[self.n_grains * 9 + 9 : self.n_grains * 10 + 9].copy().clip(0, None)
-            )
-            fractions /= fractions.sum()
-            return deformation_gradient, orientations, fractions
-
         def eval_rhs(t, y):
             """Evaluate right hand side of the D-Rex PDE."""
             # assert not np.any(np.isnan(y)), y[np.isnan(y)].shape
@@ -371,7 +357,9 @@ def eval_rhs(t, y):
 
             strain_rate = (velocity_gradient + velocity_gradient.transpose()) / 2
             strain_rate_max = np.abs(la.eigvalsh(strain_rate)).max()
-            deformation_gradient, orientations, fractions = extract_vars(y)
+            deformation_gradient, orientations, fractions = _utils.extract_vars(
+                y, self.n_grains
+            )
             # Uses nondimensional values of strain rate and velocity gradient.
             orientations_diff, fractions_diff = _core.derivatives(
                 phase=self.phase,
@@ -395,26 +383,6 @@ def eval_rhs(t, y):
                 )
             )
 
-        def apply_gbs(orientations, fractions, config):
-            """Apply grain boundary sliding for small grains."""
-            mask = fractions < (config["gbs_threshold"] / self.n_grains)
-            # _log.debug(
-            #     "grain boundary sliding activity (volume percentage): %s",
-            #     len(np.nonzero(mask)) / len(fractions),
-            # )
-            # No rotation: carry over previous orientations.
-            orientations[mask, :, :] = self.orientations[-1][mask, :, :]
-            fractions[mask] = config["gbs_threshold"] / self.n_grains
-            fractions /= fractions.sum()
-            # _log.debug(
-            #     "grain volume fractions: median=%e, min=%e, max=%e, sum=%e",
-            #     np.median(fractions),
-            #     np.min(fractions),
-            #     np.max(fractions),
-            #     np.sum(fractions),
-            # )
-            return orientations, fractions
-
         def perform_step(solver):
             """Perform SciPy solver step and appropriate processing."""
             message = solver.step()
@@ -424,8 +392,16 @@ def perform_step(solver):
             #     "%s step_size=%e", solver.__class__.__qualname__, solver.step_size
             # )
 
-            deformation_gradient, orientations, fractions = extract_vars(solver.y)
-            orientations, fractions = apply_gbs(orientations, fractions, config)
+            deformation_gradient, orientations, fractions = _utils.extract_vars(
+                solver.y, self.n_grains
+            )
+            orientations, fractions = _utils.apply_gbs(
+                orientations,
+                fractions,
+                config["gbs_threshold"],
+                self.orientations[-1],
+                self.n_grains,
+            )
             solver.y[9:] = np.hstack((orientations.flatten(), fractions))
 
         # ===== Initialise and run the solver using the above callables =====
@@ -469,7 +445,9 @@ def perform_step(solver):
             perform_step(solver)
 
         # Extract final values for this simulation step, append to storage.
-        deformation_gradient, orientations, fractions = extract_vars(solver.y.squeeze())
+        deformation_gradient, orientations, fractions = _utils.extract_vars(
+            solver.y.squeeze(), self.n_grains
+        )
         self.orientations.append(orientations)
         self.fractions.append(fractions)
         return deformation_gradient

src/pydrex/utils.py

@@ -1,4 +1,5 @@
 """> PyDRex: Miscellaneous utility methods."""
+
 from datetime import datetime
 import subprocess
 import os
@@ -27,17 +28,44 @@ def strain_increment(dt, velocity_gradient):
     )
 
 
+@nb.njit
+def apply_gbs(orientations, fractions, gbs_threshold, orientations_prev, n_grains):
+    """Apply grain boundary sliding for small grains."""
+    mask = fractions < (gbs_threshold / n_grains)
+    # _log.debug(
+    #     "grain boundary sliding activity (volume percentage): %s",
+    #     len(np.nonzero(mask)) / len(fractions),
+    # )
+    # No rotation: carry over previous orientations.
+    orientations[mask, :, :] = orientations_prev[mask, :, :]
+    fractions[mask] = gbs_threshold / n_grains
+    fractions /= fractions.sum()
+    # _log.debug(
+    #     "grain volume fractions: median=%e, min=%e, max=%e, sum=%e",
+    #     np.median(fractions),
+    #     np.min(fractions),
+    #     np.max(fractions),
+    #     np.sum(fractions),
+    # )
+    return orientations, fractions
+
+
+@nb.njit
+def extract_vars(y, n_grains):
+    """Extract deformation gradient, orientation matrices and grain sizes from y."""
+    deformation_gradient = y[:9].reshape((3, 3))
+    orientations = y[9 : n_grains * 9 + 9].reshape((n_grains, 3, 3)).clip(-1, 1)
+    fractions = y[n_grains * 9 + 9 : n_grains * 10 + 9].clip(0, None)
+    fractions /= fractions.sum()
+    return deformation_gradient, orientations, fractions
+
+
 def remove_nans(a):
     """Remove NaN values from array."""
     a = np.asarray(a)
     return a[~np.isnan(a)]
 
 
-def readable_timestamp(timestamp, tformat="%H:%M:%S"):
-    """Convert timestamp in fractional seconds to human readable format."""
-    return datetime.fromtimestamp(timestamp).strftime(tformat)
-
-
 def default_ncpus():
     """Get a safe default number of CPUs available for multiprocessing.
 
@@ -134,11 +162,11 @@ def redraw_legend(ax, fig=None, remove_all=True, **kwargs):
         and `fig.set_layout_engine("none")` before saving the figure.
 
     """
-    handler_map={
+    handler_map = {
         PathCollection: HandlerPathCollection(
             update_func=_remove_legend_symbol_transparency
         ),
-        Line2D: HandlerLine2D(update_func=_remove_legend_symbol_transparency)
+        Line2D: HandlerLine2D(update_func=_remove_legend_symbol_transparency),
     }
     if fig is None:
         legend = ax.get_legend()

tests/test_corner_flow_2d.py

@@ -160,10 +160,7 @@ def test_prescribed(self, outdir, seed, ncpus):
                         )
                         bingham_vectors[idx] = direction_mean
 
-                    _log.debug(
-                        "Total walltime: %s",
-                        _utils.readable_timestamp(perf_counter() - _begin),
-                    )
+                    _log.debug("Total walltime: %s", perf_counter() - _begin)
 
                     if outdir is not None:
                         labels.append(rf"$z_{{f}}$ = {z_exit/1e3:.1f} km")

tests/test_simple_shear_2d.py

@@ -1,4 +1,5 @@
 """> PyDRex: 2D simple shear tests."""
+
 import contextlib as cl
 import functools as ft
 from multiprocessing import Pool
@@ -303,10 +304,7 @@ def test_dvdx_ensemble(
                 if outdir is not None:
                     labels.append(f"$M^∗$ = {gbm_mobility}")
 
-            _log.info(
-                "elapsed CPU time: %s",
-                _utils.readable_timestamp(np.abs(process_time() - clock_start)),
-            )
+            _log.info("elapsed CPU time: %s", np.abs(process_time() - clock_start))
 
         # Take ensemble means and optionally plot figure.
         strains = timestamps * strain_rate
@@ -428,10 +426,7 @@ def test_dvdx_GBM(self, outdir, seeds_nearX45, ncpus):
                 if outdir is not None:
                     labels.append(f"$M^∗$ = {params['gbm_mobility']}")
 
-            _log.info(
-                "elapsed CPU time: %s",
-                _utils.readable_timestamp(np.abs(process_time() - clock_start)),
-            )
+            _log.info("elapsed CPU time: %s", np.abs(process_time() - clock_start))
 
             # Take ensemble means and optionally plot figure.
             _log.info("postprocessing results...")
@@ -459,7 +454,7 @@ def test_dvdx_GBM(self, outdir, seeds_nearX45, ncpus):
                 np.savez(
                     f"{out_basepath}_angles.npz",
                     angles=result_angles,
-                    err=result_angles_err
+                    err=result_angles_err,
                 )
                 fig, ax, colors = _vis.alignment(
                     None,
@@ -597,10 +592,7 @@ def test_GBM_calibration(self, outdir, seeds, ncpus):
                 if outdir is not None:
                     labels.append(f"$M^∗$ = {params['gbm_mobility']}")
 
-            _log.info(
-                "elapsed CPU time: %s",
-                _utils.readable_timestamp(np.abs(process_time() - clock_start)),
-            )
+            _log.info("elapsed CPU time: %s", np.abs(process_time() - clock_start))
 
             # Take ensemble means and optionally plot figure.
             _log.info("postprocessing results...")

tests/test_simple_shear_3d.py

@@ -1,4 +1,5 @@
 """> PyDRex: Simple shear 3D tests."""
+
 import contextlib as cl
 import numpy as np
 from multiprocessing import Pool
@@ -13,7 +14,6 @@
 from pydrex import velocity as _velocity
 from pydrex import stats as _stats
 from pydrex import diagnostics as _diagnostics
-from pydrex import utils as _utils
 from pydrex import io as _io
 
 # Subdirectory of `outdir` used to store outputs from these tests.
@@ -206,10 +206,7 @@ def test_direction_change(
                         ]
                     )
 
-            _log.info(
-                "elapsed CPU time: %s",
-                _utils.readable_timestamp(np.abs(process_time() - clock_start)),
-            )
+            _log.info("elapsed CPU time: %s", np.abs(process_time() - clock_start))
             _log.info("calculating ensemble averages...")
             olA_from_proj_XZ_mean = olA_from_proj_XZ.mean(axis=0)
             olA_from_proj_XZ_err = olA_from_proj_XZ.std(axis=0)