add numba, x10 speedup

Author: pavel-kirienko

jafit/ja.py

@@ -6,12 +6,20 @@
 """
 
 from __future__ import annotations
+from typing import Any, Callable
 import itertools
 from logging import getLogger
 import dataclasses
 import numpy as np
 import numpy.typing as npt
 
+try:
+    from numba import jit
+except ImportError:
+
+    def jit(**kwargs: Any) -> Callable[[Callable], Callable]:
+        return lambda f: f
+
 
 mu_0 = 1.2566370614359173e-6  # Vacuum permeability [henry/meter]
 
@@ -125,6 +133,7 @@ def solve(
 
     Returns sample points for the H, M, and B fields in the order of their appearance.
     """
+    c_r, M_s, a, k_p, alpha = coef.c_r, coef.M_s, coef.a, coef.k_p, coef.alpha
     hmb_fragments: list[list[tuple[float, float, float]]] = []
 
     # noinspection PyPep8Naming
@@ -136,9 +145,9 @@ def sweep(H: float, M: float, sign: int) -> list[tuple[float, float, float]]:
             # Integrate using Heun's method (instead of Euler's) for better stability.
             dH = sign * dH_abs
             try:
-                dM_dH_1 = _dM_dH(coef, H=H, M=M, direction=sign)
+                dM_dH_1 = _dM_dH(c_r=c_r, M_s=M_s, a=a, k_p=k_p, alpha=alpha, H=H, M=M, direction=sign)
                 M_1 = M + dM_dH_1 * dH
-                dM_dH_2 = _dM_dH(coef, H=H + dH, M=M_1, direction=sign)
+                dM_dH_2 = _dM_dH(c_r=c_r, M_s=M_s, a=a, k_p=k_p, alpha=alpha, H=H + dH, M=M_1, direction=sign)
                 M_2 = M + 0.5 * (dM_dH_1 + dM_dH_2) * dH
             except (FloatingPointError, ZeroDivisionError) as ex:
                 if idx < 10:
@@ -193,40 +202,42 @@ def sweep(H: float, M: float, sign: int) -> list[tuple[float, float, float]]:
 
 
 # noinspection PyPep8Naming
-def _dM_dH(coef: Coef, *, H: float, M: float, direction: int) -> float:
+@jit(nogil=True, nopython=True)
+def _dM_dH(*, c_r: float, M_s: float, a: float, k_p: float, alpha: float, H: float, M: float, direction: int) -> float:
+    # noinspection PyTypeChecker
     """
     Evaluates the magnetic susceptibility derivative at the given point of the M(H) curve.
     The result is sensitive to the sign of the H change; the direction is defined as sign(dH).
     This implements the model described in "Jiles–Atherton Magnetic Hysteresis Parameters Identification", Pop et al.
 
-    >>> fun = lambda H, M, d: _dM_dH(COEF_COMSOL_JA_MATERIAL, H=H, M=M, direction=d)
+    >>> fun = lambda H, M, d: _dM_dH(*dataclasses.asdict(COEF_COMSOL_JA_MATERIAL), H=H, M=M, direction=d)
     >>> assert np.isclose(fun(0, 0, +1), fun(0, 0, -1))
     >>> assert np.isclose(fun(+1, 0, +1), fun(-1, 0, -1))
     >>> assert np.isclose(fun(-1, 0, +1), fun(+1, 0, -1))
     >>> assert np.isclose(fun(-1, 0.8e6, +1), fun(+1, -0.8e6, -1))
     >>> assert np.isclose(fun(+1, 0.8e6, +1), fun(-1, -0.8e6, -1))
     """
-    if direction not in (-1, +1):
-        raise ValueError(f"Invalid direction: {direction}")
-
-    H_e = H + coef.alpha * M
-    M_an = coef.M_s * _langevin(H_e / coef.a)
-    dM_an_dH_e = coef.M_s / coef.a * _dL_dx(H_e / coef.a)
-    dM_irr_dH = (M_an - M) / (coef.k_p * direction * (1 - coef.c_r) - coef.alpha * (M_an - M))
-    return (coef.c_r * dM_an_dH_e + (1 - coef.c_r) * dM_irr_dH) / (1 - coef.alpha * coef.c_r)
+    assert direction in (-1, +1)
+    H_e = H + alpha * M
+    M_an = M_s * _langevin(H_e / a)
+    dM_an_dH_e = M_s / a * _dL_dx(H_e / a)
+    dM_irr_dH = (M_an - M) / (k_p * direction * (1 - c_r) - alpha * (M_an - M))
+    return (c_r * dM_an_dH_e + (1 - c_r) * dM_irr_dH) / (1 - alpha * c_r)  # type: ignore
 
 
+@jit(nogil=True, nopython=True)
 def _langevin(x: float) -> float:
     """
     L(x) = coth(x) - 1/x
     For tensors, the function is applied element-wise.
     """
     if np.abs(x) < _EPSILON:  # For very small |x|, use the series expansion ~ x/3
         return x / 3.0
-    return float(1.0 / np.tanh(x) - 1.0 / x)
+    return 1.0 / np.tanh(x) - 1.0 / x  # type: ignore
 
 
 # noinspection PyPep8Naming
+@jit(nogil=True, nopython=True)
 def _dL_dx(x: float) -> float:
     """
     Derivative of Langevin L(x) = coth(x) - 1/x.
@@ -236,7 +247,7 @@ def _dL_dx(x: float) -> float:
         return 1.0 / 3.0
     # exact expression: -csch^2(x) + 1/x^2
     # csch^2(x) = 1 / sinh^2(x)
-    return float(-1.0 / (np.sinh(x) ** 2) + 1.0 / (x**2))
+    return -1.0 / (np.sinh(x) ** 2) + 1.0 / (x**2)  # type: ignore
 
 
 _logger = getLogger(__name__)

jafit/jafit.py

@@ -179,7 +179,7 @@ def visualize(
 ) -> None:
     import matplotlib.pyplot as plt
 
-    fig, (ax_m, ax_b) = plt.subplots(2, 1, figsize=(12, 10), sharex="all")
+    fig, (ax_m, ax_b) = plt.subplots(2, 1, figsize=(12, 10), sharex="all")  # type: ignore
     try:
         # Plot the curves predicted by the JA model
         for i, fragment in enumerate(sol.H_M_B_segments, start=1):
@@ -261,6 +261,7 @@ def main() -> None:
     logging.getLogger("matplotlib").setLevel(logging.ERROR)
     logging.getLogger("numpy").setLevel(logging.WARNING)
     logging.getLogger("scipy").setLevel(logging.WARNING)
+    logging.getLogger("numba").setLevel(logging.WARNING)
 
     np.seterr(all="raise")
 

pyproject.toml

@@ -28,7 +28,7 @@ show_error_context = true
 mypy_path = []
 
 [[tool.mypy.overrides]]
-module = ["scipy", "scipy.*", "matplotlib"]
+module = ["scipy", "scipy.*", "matplotlib", "numba"]
 ignore_missing_imports = true
 
 # --------------------------------------------------  BLACK  --------------------------------------------------