Merge pull request #1213 from paulromano/fast-float-endf

Speed up reading floating-point numbers from ENDF

openmc/data/_endf.pyx

@@ -0,0 +1,8 @@
+# cython: c_string_type=str, c_string_encoding=ascii
+
+cdef extern from "endf.c":
+    double cfloat_endf(const char* buffer, int n)
+
+def float_endf(s):
+    cdef const char* c_string = s
+    return cfloat_endf(c_string, len(s))

openmc/data/endf.c

@@ -0,0 +1,39 @@
+#include <stdlib.h>
+
+double cfloat_endf(const char* buffer, int n)
+{
+  char arr[12]; // 11 characters plus a null terminator
+  int j = 0; // current position in arr
+  int found_significand = 0;
+  int found_exponent = 0;
+  for (int i = 0; i < n; ++i) {
+    // Skip whitespace characters
+    char c = buffer[i];
+    if (c == ' ') continue;
+
+    if (found_significand) {
+      if (!found_exponent) {
+        if (c == '+' || c == '-') {
+          // In the case that we encounter +/- and we haven't yet encountered
+          // e/E, we manually add it
+          arr[j++] = 'e';
+          found_exponent = 1;
+
+        } else if (c == 'e' || c == 'E' || c == 'd' || c == 'D') {
+          arr[j++] = 'e';
+          found_exponent = 1;
+          continue;
+        }
+      }
+    } else if (c == '.' || (c >= '0' && c <= '9')) {
+      found_significand = 1;
+    }
+
+    // Copy character
+    arr[j++] = c;
+  }
+
+  // Done copying. Add null terminator and convert to double
+  arr[j] = '\0';
+  return atof(arr);
+}

openmc/data/endf.py

@@ -20,6 +20,11 @@ from numpy.polynomial.polynomial import Polynomial
 from .data import ATOMIC_SYMBOL, gnd_name
 from .function import Tabulated1D, INTERPOLATION_SCHEME
 from openmc.stats.univariate import Uniform, Tabular, Legendre
+try:
+    from ._endf import float_endf
+    _CYTHON = True
+except ImportError:
+    _CYTHON = False
 
 
 _LIBRARY = {0: 'ENDF/B', 1: 'ENDF/A', 2: 'JEFF', 3: 'EFF',
@@ -69,7 +74,7 @@ SUM_RULES = {1: [2, 3],
 ENDF_FLOAT_RE = re.compile(r'([\s\-\+]?\d*\.\d+)([\+\-]) ?(\d+)')
 
 
-def float_endf(s):
+def py_float_endf(s):
     """Convert string of floating point number in ENDF to float.
 
     The ENDF-6 format uses an 'e-less' floating point number format,
@@ -92,6 +97,10 @@ def float_endf(s):
     return float(ENDF_FLOAT_RE.sub(r'\1e\2\3', s))
 
 
+if not _CYTHON:
+    float_endf = py_float_endf
+
+
 def int_endf(s):
     """Convert string of integer number in ENDF to int.
 

openmc/stats/univariate.py

@@ -312,7 +312,7 @@ class Normal(Univariate):
     r"""Normally distributed sampling.
 
     The Normal Distribution is characterized by two parameters
-    :math:`\mu` and :math:`\sigma` and has density function 
+    :math:`\mu` and :math:`\sigma` and has density function
     :math:`p(X) dX = 1/(\sqrt{2\pi}\sigma) e^{-(X-\mu)^2/(2\sigma^2)}`
 
     Parameters
@@ -325,9 +325,9 @@ class Normal(Univariate):
     Attributes
     ----------
     mean_value : float
-        Mean of the Normal distribution 
+        Mean of the Normal distribution
     std_dev : float
-        Standard deviation of the Normal distribution 
+        Standard deviation of the Normal distribution
     """
 
     def __init__(self, mean_value, std_dev):
@@ -380,17 +380,17 @@ class Normal(Univariate):
 class Muir(Univariate):
     """Muir energy spectrum.
 
-    The Muir energy spectrum is a Gaussian spectrum, but for 
+    The Muir energy spectrum is a Gaussian spectrum, but for
     convenience reasons allows the user 3 parameters to define
     the distribution, e0 the mean energy of particles, the mass
-    of reactants m_rat, and the ion temperature kt. 
+    of reactants m_rat, and the ion temperature kt.
 
     Parameters
     ----------
     e0 : float
         Mean of the Muir distribution in units of eV
     m_rat : float
-        Ratio of the sum of the masses of the reaction inputs to an 
+        Ratio of the sum of the masses of the reaction inputs to an
         AMU
     kt : float
          Ion temperature for the Muir distribution in units of eV
@@ -400,7 +400,7 @@ class Muir(Univariate):
     e0 : float
         Mean of the Muir distribution in units of eV
     m_rat : float
-        Ratio of the sum of the masses of the reaction inputs to an 
+        Ratio of the sum of the masses of the reaction inputs to an
         AMU
     kt : float
          Ion temperature for the Muir distribution in units of eV
@@ -582,26 +582,27 @@ class Legendre(Univariate):
 
     def __init__(self, coefficients):
         self.coefficients = coefficients
+        self._legendre_poly = None
 
     def __call__(self, x):
-        return self._legendre_polynomial(x)
+        # Create Legendre polynomial if we haven't yet
+        if self._legendre_poly is None:
+            l = np.arange(len(self._coefficients))
+            coeffs = (2.*l + 1.)/2. * self._coefficients
+            self._legendre_poly = np.polynomial.Legendre(coeffs)
+
+        return self._legendre_poly(x)
 
     def __len__(self):
-        return len(self._legendre_polynomial.coef)
+        return len(self._coefficients)
 
     @property
     def coefficients(self):
-        poly = self._legendre_polynomial
-        l = np.arange(poly.degree() + 1)
-        return 2./(2.*l + 1.) * poly.coef
+        return self._coefficients
 
     @coefficients.setter
     def coefficients(self, coefficients):
-        cv.check_type('Legendre expansion coefficients', coefficients,
-                      Iterable, Real)
-        l = np.arange(len(coefficients))
-        coeffs = (2.*l + 1.)/2. * np.array(coefficients)
-        self._legendre_polynomial = np.polynomial.Legendre(coeffs)
+        self._coefficients = np.asarray(coefficients)
 
     def to_xml_element(self, element_name):
         raise NotImplementedError

setup.py

@@ -68,10 +68,10 @@ kwargs = {
     },
 }
 
-# If Cython is present, add resonance reconstruction capability
+# If Cython is present, add resonance reconstruction and fast float_endf
 if have_cython:
     kwargs.update({
-        'ext_modules': cythonize('openmc/data/reconstruct.pyx'),
+        'ext_modules': cythonize('openmc/data/*.pyx'),
         'include_dirs': [np.get_include()]
     })
 

tests/unit_tests/test_endf.py

@@ -9,6 +9,19 @@ def test_float_endf():
     assert endf.float_endf('6.022-23') == approx(6.022e-23)
     assert endf.float_endf(' +1.01+ 2') == approx(101.0)
     assert endf.float_endf(' -1.01- 2') == approx(-0.0101)
+    assert endf.float_endf('+ 2 . 3+ 1') == approx(23.0)
+    assert endf.float_endf('-7 .8 -1') == approx(-0.78)
+    assert endf.float_endf('3.14e0') == approx(3.14)
+    assert endf.float_endf('3.14E0') == approx(3.14)
+    assert endf.float_endf('3.14e-1') == approx(0.314)
+    assert endf.float_endf('3.14d0') == approx(3.14)
+    assert endf.float_endf('3.14D0') == approx(3.14)
+    assert endf.float_endf('3.14d-1') == approx(0.314)
+    assert endf.float_endf('1+2') == approx(100.0)
+    assert endf.float_endf('-1+2') == approx(-100.0)
+    assert endf.float_endf('1.+2') == approx(100.0)
+    assert endf.float_endf('-1.+2') == approx(-100.0)
+    assert endf.float_endf('        ') == 0.0
 
 
 def test_int_endf():