added function to handle dicts of values

openmc/data/grid.py

@@ -31,6 +31,7 @@ def linearize(x, f, tolerance=0.001):
     # Initialize stack
     x_stack = [x[0]]
     y_stack = [f(x[0])]
+    print(y_stack)
 
     for i in range(x.shape[0] - 1):
         x_stack.insert(0, x[i + 1])
@@ -112,3 +113,70 @@ def thin(x, y, tolerance=0.001):
         y_out[i_remove] = np.nan
 
     return x_out[np.isfinite(x_out)], y_out[np.isfinite(y_out)]
+
+def linearizeIter(x, f, tolerance=0.001, unified=True):
+    """Return a tabulated representation of multiple functions of one
+    variable.
+
+    Parameters
+    ----------
+    x : Iterable of float
+        Initial x values at which the function should be evaluated
+    f : Callable
+        Function of a single variable that returns a dictionary
+    tolerance : float
+        Tolerance on the interpolation error
+    unified : boolean
+        Flag to indicate usage of a unified grid for all functions
+        if True, or independent grids if False
+
+    Returns
+    -------
+    numpy.ndarray
+        Tabulated values of the independent variable
+    dictionary of numpy.ndarray's
+        Tabulated values of the dependent variable
+
+    """
+    if unified==True:
+        # Initialize dictionary of output
+        y_dict = f(x[0])
+    
+        for item in y_dict:
+            #Initialize output
+            x_out = []
+            y_out = []
+            print(str(item))
+            #Initialize stacks
+            x_stack = [x[0]]
+            print(y_dict)
+            y_stack = [y_dict[item]]
+            for i in range(x.shape[0] - 1):
+                print(x_stack)
+                x_stack.insert(0, x[i + 1])
+                print(x_stack)
+                y_stack.insert(0, f(x[i + 1])[item])
+        
+                while True:
+                    x_high, x_low = x_stack[-2:]
+                    y_high, y_low = y_stack[-2:]
+                    x_mid = 0.5*(x_low + x_high)
+                    y_mid = f(x_mid)[item]
+        
+                    y_interp = y_low + (y_high - y_low)/(x_high - x_low)*(x_mid - x_low)
+                    error = abs((y_interp - y_mid)/y_mid)
+                    if error > tolerance:
+                        x_stack.insert(-1, x_mid)
+                        y_stack.insert(-1, y_mid)
+                    else:
+                        x_out.append(x_stack.pop())
+                        y_out.append(y_stack.pop())
+                        if len(x_stack) == 1:
+                            break
+        
+            x_out.append(x_stack.pop())
+            y_out.append(y_stack.pop())
+            x=np.array(x_out) #Use x_out for initial x values in next item
+
+        y_dict_out = f(np.array(x_out))
+        return np.array(x_out), y_dict_out