Merge pull request #1732 from drewejohnson/1713-dep-results-index

Easier access and indexing to depletion time points

openmc/deplete/results_list.py

@@ -1,8 +1,12 @@
+import numbers
+import bisect
+import math
+
 import h5py
 import numpy as np
 
 from .results import Results, VERSION_RESULTS
-from openmc.checkvalue import check_filetype_version, check_value
+from openmc.checkvalue import check_filetype_version, check_value, check_type
 
 
 __all__ = ["ResultsList"]
@@ -192,3 +196,103 @@ class ResultsList(list):
         for ix, res in enumerate(items):
             times[ix] = res.proc_time
         return times
+
+    def get_times(self, time_units="d") -> np.ndarray:
+        """Return the points in time that define the depletion schedule
+
+
+        .. versionadded:: 0.12.1
+
+        Parameters
+        ----------
+        time_units : {"s", "d", "h", "min"}, optional
+            Return the vector in these units. Default is to
+            convert to days
+
+        Returns
+        -------
+        numpy.ndarray
+            1-D vector of time points
+
+        """
+        check_type("time_units", time_units, str)
+
+        times = np.fromiter(
+            (r.time[0] for r in self),
+            dtype=self[0].time.dtype,
+            count=len(self),
+        )
+
+        if time_units == "d":
+            times /= (60 * 60 * 24)
+        elif time_units == "h":
+            times /= (60 * 60)
+        elif time_units == "min":
+            times /= 60
+        elif time_units != "s":
+            raise ValueError(
+                'Unable to set "time_units" to {} since it is not '
+                'in ("s", "d", "min", "h")'.format(time_units)
+            )
+        return times
+
+    def get_step_where(
+        self, time, time_units="d", atol=1e-6, rtol=1e-3
+    ) -> int:
+        """Return the index closest to a given point in time
+
+        In the event ``time`` lies exactly between two points, the
+        lower index will be returned. It is possible that the index
+        will be at most one past the point in time requested, but only
+        according to tolerances requested.
+
+        Passing ``atol=math.inf`` and ``rtol=math.inf`` will return
+        the closest index to the requested point.
+
+
+        .. versionadded:: 0.12.1
+
+        Parameters
+        ----------
+        time : float
+            Desired point in time
+        time_units : {"s", "d", "min", "h"}, optional
+            Units on ``time``. Default: days
+        atol : float, optional
+            Absolute tolerance (in ``time_units``) if ``time`` is not
+            found.
+        rtol : float, optional
+            Relative tolerance if ``time`` is not found.
+
+        Returns
+        -------
+        int
+
+        """
+        check_type("time", time, numbers.Real)
+        check_type("atol", atol, numbers.Real)
+        check_type("rtol", rtol, numbers.Real)
+
+        times = self.get_times(time_units)
+
+        if times[0] < time < times[-1]:
+            ix = bisect.bisect_left(times, time)
+            if ix == times.size:
+                ix -= 1
+            # Bisection will place us either directly on the point
+            # or one-past the first value less than time
+            elif time - times[ix - 1] <= times[ix] - time:
+                ix -= 1
+        elif times[0] >= time:
+            ix = 0
+        elif time >= times[-1]:
+            ix = times.size - 1
+
+        if math.isclose(time, times[ix], rel_tol=rtol, abs_tol=atol):
+            return ix
+
+        raise ValueError(
+            "A value of {} {} was not found given absolute and "
+            "relative tolerances {} and {}.".format(
+                time, time_units, atol, rtol)
+        )

tests/unit_tests/test_deplete_resultslist.py

@@ -1,6 +1,7 @@
 """Tests the ResultsList class"""
 
 from pathlib import Path
+from math import inf
 
 import numpy as np
 import pytest
@@ -65,3 +66,66 @@ def test_get_eigenvalue(res):
     np.testing.assert_allclose(t, t_ref)
     np.testing.assert_allclose(k[:, 0], k_ref)
     np.testing.assert_allclose(k[:, 1], u_ref)
+
+
+@pytest.mark.parametrize("unit", ("s", "d", "min", "h"))
+def test_get_steps(unit):
+    # Make a ResultsList full of near-empty Result instances
+    # Just fill out a time schedule
+    results = openmc.deplete.ResultsList()
+    # Time in units of unit
+    times = np.linspace(0, 100, num=5)
+    if unit == "d":
+        conversion_to_seconds = 60 * 60 * 24
+    elif unit == "h":
+        conversion_to_seconds = 60 * 60
+    elif unit == "min":
+        conversion_to_seconds = 60
+    else:
+        conversion_to_seconds = 1
+
+    for ix in range(times.size):
+        res = openmc.deplete.Results()
+        res.time = times[ix:ix + 1] * conversion_to_seconds
+        results.append(res)
+
+    for expected, value in enumerate(times):
+        actual = results.get_step_where(
+            value, time_units=unit, atol=0, rtol=0)
+        assert actual == expected, (value, results[actual].time[0])
+
+    with pytest.raises(ValueError):
+        # Emulate a result file with a non-zero initial point in time
+        # as in starting from a restart
+        results.get_step_where(times[0] - 1, time_units=unit, atol=0, rtol=0)
+
+    with pytest.raises(ValueError):
+        results.get_step_where(times[-1] + 1, time_units=unit, atol=0, rtol=0)
+
+    # Grab intermediate points with a small offset
+    delta = (times[1] - times[0])
+    offset = delta * 0.1
+    for expected, value in enumerate(times[1:-1], start=1):
+        # Shoot a little low and a little high
+        for mult in (1, -1):
+            target = value + mult * offset
+            # Compare using absolute and relative tolerances
+            actual = results.get_step_where(
+                target, time_units=unit, atol=offset * 2, rtol=inf)
+            assert actual == expected, (
+                target, times[actual], times[expected], offset)
+
+            actual = results.get_step_where(
+                target, time_units=unit, atol=inf, rtol=offset / value)
+            assert actual == expected, (
+                target, times[actual], times[expected], offset)
+        # Check that the lower index is returned for the exact mid-point
+        target = value + delta * 0.5
+        actual = results.get_step_where(
+            target, time_units=unit, atol=delta, rtol=delta / value)
+        assert actual == expected
+
+    # Shoot way over with no tolerance -> just give closest value
+    actual = results.get_step_where(
+        times[-1] * 100, time_units=unit, atol=inf, rtol=inf)
+    assert actual == times.size - 1