Merge pull request #72 from Gromwud/main

Perfomance improvements and refactoring

Author: Gromwud

epde/_loop_stats.py

@@ -0,0 +1,106 @@
+"""Lightweight retry/condition-loop instrumentation.
+
+Off by default. Set ``EPDE_LOOP_STATS=1`` to enable; ~30 source-level
+``record(...)`` call sites then accumulate per-loop stats that
+``report()`` formats as a table.
+
+Cost when disabled: a single global-var read per ``record`` call.
+Cost when enabled: a dict lookup + list append per loop exit.
+"""
+from __future__ import annotations
+
+import os
+import sys
+from collections import defaultdict
+from typing import Optional
+
+_ENABLED = os.environ.get('EPDE_LOOP_STATS', '0') == '1'
+
+
+def _new_bucket():
+    return {'entries': 0, 'iters': [], 'hit_cap': 0, 'early_exit': 0, 'caps': set()}
+
+
+_stats = defaultdict(_new_bucket)
+
+
+def enabled() -> bool:
+    return _ENABLED
+
+
+def record(site: str, iters: int, cap: int) -> None:
+    """Record one loop exit.
+
+    ``site`` is a human label like ``"EqRPS.outer"``. ``iters`` is the
+    number of iterations actually executed. ``cap`` is the loop's
+    maximum (use ``sys.maxsize`` for condition-driven loops with no
+    explicit cap).
+    """
+    if not _ENABLED:
+        return
+    b = _stats[site]
+    b['entries'] += 1
+    b['iters'].append(iters)
+    b['caps'].add(cap)
+    if iters >= cap:
+        b['hit_cap'] += 1
+    elif iters <= 1:
+        b['early_exit'] += 1
+
+
+def reset() -> None:
+    _stats.clear()
+
+
+def _stats_for(name: str) -> dict:
+    b = _stats[name]
+    n = b['entries']
+    iters = b['iters']
+    if n == 0:
+        return {'entries': 0}
+    iters_sorted = sorted(iters)
+    median = iters_sorted[n // 2]
+    return {
+        'entries': n,
+        'mean': sum(iters) / n,
+        'median': median,
+        'max': max(iters),
+        'p95': iters_sorted[min(n - 1, int(n * 0.95))],
+        'total_iters': sum(iters),
+        'hit_cap_pct': 100.0 * b['hit_cap'] / n,
+        'early_exit_pct': 100.0 * b['early_exit'] / n,
+        'cap': max(b['caps']) if b['caps'] else 0,
+    }
+
+
+def report(path: Optional[str] = None) -> str:
+    """Format all recorded loops as a table, sorted by total iterations.
+
+    Writes to ``path`` if given AND also returns the string.
+    """
+    sites = sorted(_stats.keys(),
+                   key=lambda s: -sum(_stats[s]['iters']) if _stats[s]['iters'] else 0)
+    lines = []
+    header = (f"{'site':<45} {'entries':>8} {'mean':>7} {'med':>5} "
+              f"{'p95':>5} {'max':>5} {'cap':>6} {'%cap':>6} "
+              f"{'%early':>7} {'totIters':>10}")
+    lines.append(header)
+    lines.append('-' * len(header))
+    if not _ENABLED:
+        lines.append('(EPDE_LOOP_STATS disabled -- set EPDE_LOOP_STATS=1 to record)')
+    for site in sites:
+        s = _stats_for(site)
+        if s['entries'] == 0:
+            continue
+        cap_str = 'inf' if s['cap'] >= sys.maxsize else str(s['cap'])
+        lines.append(
+            f"{site:<45} {s['entries']:>8d} {s['mean']:>7.2f} "
+            f"{s['median']:>5d} {s['p95']:>5d} {s['max']:>5d} "
+            f"{cap_str:>6} {s['hit_cap_pct']:>5.1f}% "
+            f"{s['early_exit_pct']:>6.1f}% {s['total_iters']:>10d}"
+        )
+    text = '\n'.join(lines)
+    if path is not None:
+        with open(path, 'w') as f:
+            f.write(text + '\n')
+    return text

epde/evaluators.py

@@ -28,9 +28,19 @@ def __call__(self, factor, structural: bool = False, grids: list = None,
 
 
 class CustomEvaluator(EvaluatorTemplate):
-    def __init__(self, evaluation_functions_np: Union[Callable, dict] = None, 
+    def __init__(self, evaluation_functions_np: Union[Callable, dict] = None,
                  evaluation_functions_torch: Union[Callable, dict] = None,
-                 eval_fun_params_labels: Union[list, tuple, set] = ['power']):
+                 eval_fun_params_labels: Union[list, tuple, set] = ['power'],
+                 native_vectorized: bool = False):
+        """Wrap one or many evaluation functions for use as a factor evaluator.
+
+        ``native_vectorized=True`` skips the per-element ``np.vectorize``
+        dispatch on the hot path: the func is called ONCE with the full
+        grid arrays. The built-in evaluators in this module all set this
+        flag because their numpy ops (``np.cos``, ``np.sin``, ``np.power``,
+        ``np.full_like``, etc.) vectorize natively. User code passing a
+        non-vectorising callable should leave the default ``False``.
+        """
         self._evaluation_functions_np = evaluation_functions_np
         self._evaluation_functions_torch = evaluation_functions_torch
 
@@ -43,6 +53,7 @@ def __init__(self, evaluation_functions_np: Union[Callable, dict] = None,
             self._single_function_token = True
 
         self.eval_fun_params_labels = eval_fun_params_labels
+        self.native_vectorized = native_vectorized
 
     def __call__(self, factor, structural: bool = False, func_args: List[Union[torch.Tensor, np.ndarray]] = None, 
                  torch_mode: bool = False, **kwargs): # s
@@ -67,23 +78,32 @@ def __call__(self, factor, structural: bool = False, func_args: List[Union[torch
                 if param_descr['name'] == key:
                     eval_fun_kwargs[key] = factor.params[param_idx]
 
-        grid_function = np.vectorize(lambda args: funcs(*args, **eval_fun_kwargs))
-
         if func_args is None:
             new_grid = False
             func_args = factor.grids
         else:
             new_grid = True
-        try:
-            if new_grid:
-                raise AttributeError
-            self.indexes_vect
-        except AttributeError:
-            self.indexes_vect = np.empty_like(func_args[0], dtype=object)
-            for tensor_idx, _ in np.ndenumerate(func_args[0]):
-                self.indexes_vect[tensor_idx] = tuple([subarg[tensor_idx]
-                                                       for subarg in func_args])
-        value = grid_function(self.indexes_vect)
+
+        if self.native_vectorized:
+            # Fast path: call funcs once with the full grid arrays. The
+            # built-in numpy evaluators (trig, sign, grid, inverse,
+            # const, velocity) all return an array of shape
+            # ``func_args[0].shape``. This skips an N-element
+            # ``np.vectorize`` loop that on Wave (65k samples)
+            # dominated evaluator self-time at ~35 s per run.
+            value = funcs(*func_args, **eval_fun_kwargs)
+        else:
+            grid_function = np.vectorize(lambda args: funcs(*args, **eval_fun_kwargs))
+            try:
+                if new_grid:
+                    raise AttributeError
+                self.indexes_vect
+            except AttributeError:
+                self.indexes_vect = np.empty_like(func_args[0], dtype=object)
+                for tensor_idx, _ in np.ndenumerate(func_args[0]):
+                    self.indexes_vect[tensor_idx] = tuple([subarg[tensor_idx]
+                                                           for subarg in func_args])
+            value = grid_function(self.indexes_vect)
         value = value[global_var.grid_cache.g_func != 0]
         value = value.reshape(-1)
         return value
@@ -125,7 +145,9 @@ def simple_function_evaluator(factor, structural: bool = False, grids=None,
 
     else:
         if factor.params[power_param_idx] == 1:
-            value = global_var.tensor_cache.get(factor.cache_label, structural = structural, torch_mode = torch_mode)
+            # Same bucketed key Factor.evaluate uses so trig factors with
+            # within-tolerance freq share a single cached evaluation.
+            value = global_var.tensor_cache.get(factor.structural_label, structural = structural, torch_mode = torch_mode)
             return value
         else:
             value = global_var.tensor_cache.get(factor_params_to_str(factor, set_default_power = True,
@@ -259,30 +281,37 @@ def vhef_grad_15(*grids, **kwargs):
              vhef_grad_10, vhef_grad_11, vhef_grad_12,
              vhef_grad_13, vhef_grad_14, vhef_grad_15]
 
-sign_evaluator = CustomEvaluator(evaluation_functions_np=sign_eval_fun_np, 
-                                evaluation_functions_torch=sign_eval_fun_torch, 
-                                eval_fun_params_labels = ['power', 'dim'])
+sign_evaluator = CustomEvaluator(evaluation_functions_np=sign_eval_fun_np,
+                                evaluation_functions_torch=sign_eval_fun_torch,
+                                eval_fun_params_labels = ['power', 'dim'],
+                                native_vectorized=True)
 
-phased_sine_evaluator = CustomEvaluator(evaluation_functions_np = phased_sine_1d_np, 
+phased_sine_evaluator = CustomEvaluator(evaluation_functions_np = phased_sine_1d_np,
                                         evaluation_functions_torch = phased_sine_1d_torch,
-                                        eval_fun_params_labels = ['power', 'freq', 'phase']) # , use_factors_grids = True
+                                        eval_fun_params_labels = ['power', 'freq', 'phase'],
+                                        native_vectorized=True) # , use_factors_grids = True
 trigonometric_evaluator = CustomEvaluator(evaluation_functions_np = trig_eval_fun_np,
                                           evaluation_functions_torch = trig_eval_fun_torch,
-                                          eval_fun_params_labels=['freq', 'dim', 'power']) # , use_factors_grids = True
+                                          eval_fun_params_labels=['freq', 'dim', 'power'],
+                                          native_vectorized=True) # , use_factors_grids = True
 grid_evaluator = CustomEvaluator(evaluation_functions_np = grid_eval_fun_np,
                                  evaluation_functions_torch = grid_eval_fun_torch,
-                                 eval_fun_params_labels=['dim', 'power']) # , use_factors_grids=True
+                                 eval_fun_params_labels=['dim', 'power'],
+                                 native_vectorized=True) # , use_factors_grids=True
 
 inverse_function_evaluator = CustomEvaluator(evaluation_functions_np = inverse_eval_fun_np,
                                              evaluation_functions_torch = inverse_eval_fun_torch,
-                                             eval_fun_params_labels=['dim', 'power']) # , use_factors_grids=True
+                                             eval_fun_params_labels=['dim', 'power'],
+                                             native_vectorized=True) # , use_factors_grids=True
 
 const_evaluator = CustomEvaluator(evaluation_functions_np = const_eval_fun_np,
-                                  evaluation_functions_torch = const_eval_fun_torch, 
-                                  eval_fun_params_labels = ['power', 'value'])
+                                  evaluation_functions_torch = const_eval_fun_torch,
+                                  eval_fun_params_labels = ['power', 'value'],
+                                  native_vectorized=True)
 const_grad_evaluator = CustomEvaluator(evaluation_functions_np = const_grad_fun_np,
                                        evaluation_functions_torch =  const_grad_fun_np,
-                                       eval_fun_params_labels = ['power', 'value'])
+                                       eval_fun_params_labels = ['power', 'value'],
+                                       native_vectorized=True)
 
 velocity_evaluator = CustomEvaluator(velocity_heating_eval_fun, ['p' + str(idx+1) for idx in range(15)])
 velocity_grad_evaluators = [CustomEvaluator(component, ['p' + str(idx+1) for idx in range(15)])

epde/interface/token_family.py

@@ -396,8 +396,8 @@ def evaluate_all(self, all_vars: List[str]):
                     generated_token.use_grids_cache()
                 generated_token.scaled = False
                 _ = generated_token.evaluate()
-                print(generated_token.cache_label)
-                if generated_token.cache_label not in global_var.tensor_cache.memory_default['numpy'].keys():
+                print(generated_token.structural_label)
+                if generated_token.structural_label not in global_var.tensor_cache.memory_default['numpy'].keys():
                     raise KeyError('Generated token somehow was not stored in cache.')
 
 

epde/operators/common/fitness.py

@@ -186,8 +186,8 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
 
         fitness_value = rl_error
 
-        # if force_out_of_place:
-        #     return fitness_value
+        if force_out_of_place:
+            return fitness_value
 
         objective.aic = None
         objective.aic_calculated = True
@@ -211,8 +211,8 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
                 cv = (std ** 2) / (mu ** 2)
             total_lr = sum(cv) / len(data_shape)
 
-        if force_out_of_place:
-            return fitness_value * total_lr
+        # if force_out_of_place:
+        #     return fitness_value * total_lr
 
         objective.fitness_calculated = True
         objective.fitness_value = fitness_value