Merge pull request #76 from Gromwud/main

Sparsity.py added

Author: Gromwud

epde/_loop_stats.py

@@ -4,13 +4,22 @@
 ``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.
+A second metric class -- wall-clock timers -- is exposed via
+``timer(site)``. Same env-var gate, same ``report()`` output (second
+table). Both classes share the site namespace, so a site may appear in
+both tables (e.g., ``EqRPS.outer`` records iters; ``EqRPS.apply``
+records wall-clock).
+
+Cost when disabled: a single global-var read per ``record`` / ``timer``
+call (the latter returns a shared no-op context-manager singleton).
+Cost when enabled: a dict lookup + list append per loop exit, plus a
+``perf_counter`` pair + accumulate per ``timer`` exit.
 """
 from __future__ import annotations
 
 import os
 import sys
+import time
 from collections import defaultdict
 from typing import Optional
 
@@ -21,7 +30,12 @@ def _new_bucket():
     return {'entries': 0, 'iters': [], 'hit_cap': 0, 'early_exit': 0, 'caps': set()}
 
 
+def _new_timer_bucket():
+    return {'entries': 0, 'total_s': 0.0, 'max_s': 0.0}
+
+
 _stats = defaultdict(_new_bucket)
+_timers = defaultdict(_new_timer_bucket)
 
 
 def enabled() -> bool:
@@ -48,8 +62,78 @@ def record(site: str, iters: int, cap: int) -> None:
         b['early_exit'] += 1
 
 
+class _ActiveTimer:
+    """Per-call active timer; one instance per ``with timer(site):``."""
+    __slots__ = ('_site', '_t0')
+
+    def __init__(self, site: str) -> None:
+        self._site = site
+        self._t0 = 0.0
+
+    def __enter__(self):
+        self._t0 = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        dt = time.perf_counter() - self._t0
+        b = _timers[self._site]
+        b['entries'] += 1
+        b['total_s'] += dt
+        if dt > b['max_s']:
+            b['max_s'] = dt
+        return False
+
+
+class _NoopTimer:
+    """Singleton CM returned when ``EPDE_LOOP_STATS`` is off."""
+    __slots__ = ()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+
+_NOOP_TIMER = _NoopTimer()
+
+
+def timer(site: str):
+    """Return a context manager that accumulates wall-clock under ``site``.
+
+    Disabled fast path: returns a shared singleton no-op CM; the
+    ``with`` block then costs one __enter__/__exit__ method call and
+    nothing else. Enabled path: allocates a small ``_ActiveTimer``
+    per call.
+    """
+    if not _ENABLED:
+        return _NOOP_TIMER
+    return _ActiveTimer(site)
+
+
+def timed(site: str):
+    """Decorator wrapping ``fn`` with ``timer(site)``.
+
+    Avoids re-indenting large ``apply`` bodies when adding probes. When
+    ``EPDE_LOOP_STATS`` is off, cost is one extra function call + the
+    no-op CM enter/exit, all of which are negligible compared to the
+    wrapped operator work.
+    """
+    def deco(fn):
+        def wrapper(*args, **kwargs):
+            with timer(site):
+                return fn(*args, **kwargs)
+        wrapper.__wrapped__ = fn
+        wrapper.__name__ = getattr(fn, '__name__', 'wrapper')
+        wrapper.__qualname__ = getattr(fn, '__qualname__', 'wrapper')
+        wrapper.__doc__ = fn.__doc__
+        return wrapper
+    return deco
+
+
 def reset() -> None:
     _stats.clear()
+    _timers.clear()
 
 
 def _stats_for(name: str) -> dict:
@@ -73,8 +157,17 @@ def _stats_for(name: str) -> dict:
     }
 
 
+def timers_snapshot() -> dict:
+    """Return a copy of the timers dict for external consumers.
+
+    Used by ``profile_loop_stats.py`` to build the cross-system
+    compare table without re-parsing the report text.
+    """
+    return {site: dict(b) for site, b in _timers.items()}
+
+
 def report(path: Optional[str] = None) -> str:
-    """Format all recorded loops as a table, sorted by total iterations.
+    """Format all recorded loops + timers as two tables.
 
     Writes to ``path`` if given AND also returns the string.
     """
@@ -84,6 +177,7 @@ def report(path: Optional[str] = None) -> str:
     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('LOOPS')
     lines.append(header)
     lines.append('-' * len(header))
     if not _ENABLED:
@@ -99,6 +193,34 @@ def report(path: Optional[str] = None) -> str:
             f"{cap_str:>6} {s['hit_cap_pct']:>5.1f}% "
             f"{s['early_exit_pct']:>6.1f}% {s['total_iters']:>10d}"
         )
+
+    lines.append('')
+    timer_sites = sorted(_timers.keys(),
+                         key=lambda s: -_timers[s]['total_s'])
+    timer_header = (f"{'site':<45} {'entries':>8} {'total_s':>10} "
+                    f"{'mean_ms':>10} {'max_ms':>10} {'share%':>7}")
+    lines.append('TIMERS')
+    lines.append(timer_header)
+    lines.append('-' * len(timer_header))
+    if not _ENABLED:
+        lines.append('(EPDE_LOOP_STATS disabled -- set EPDE_LOOP_STATS=1 to record)')
+    if timer_sites:
+        total_max = max(_timers[s]['total_s'] for s in timer_sites)
+    else:
+        total_max = 0.0
+    for site in timer_sites:
+        b = _timers[site]
+        n = b['entries']
+        if n == 0:
+            continue
+        mean_ms = 1000.0 * b['total_s'] / n
+        max_ms = 1000.0 * b['max_s']
+        share = (100.0 * b['total_s'] / total_max) if total_max > 0 else 0.0
+        lines.append(
+            f"{site:<45} {n:>8d} {b['total_s']:>10.2f} "
+            f"{mean_ms:>10.2f} {max_ms:>10.2f} {share:>6.1f}%"
+        )
+
     text = '\n'.join(lines)
     if path is not None:
         with open(path, 'w') as f:

epde/globals.py

@@ -22,6 +22,58 @@
 from epde.preprocessing.smoothers import NN
 
 
+# Gram-construction configuration, read by VWSRSparsity.apply,
+# PhysicsInformedLasso.fit, EqRightPartSelector._precompute_super_gram,
+# and L2LRFitness.apply. ``mode='vcoef'`` (default) uses the
+# varying-coefficient stability estimator (``VaryingCoefSetup``);
+# ``mode='axis'`` is the legacy axis-aligned sliding-window backup
+# (``GramSetup`` reduced by the var/mu^2 CV in
+# ``PhysicsInformedLasso.get_cv``).
+gram_mode: str = 'vcoef'
+
+# Per-rep seed for additive Gaussian noise applied at ``cfg.load_data()``;
+# rewritten each rep so every rep sees an independent noise realization.
+noise_seed = None
+
+# ``gram_mode='vcoef'`` varying-coefficient stability config (see
+# ``epde.operators.common.stability.VaryingCoefSetup``). ``vc_modes_cache`` resolves the
+# per-axis basis resolution once per ``(grid_shape, main_var)`` from the
+# Taylor microscale and reuses it for every candidate so the basis is
+# identical across individuals; cleared on ``set_gram_config``. ``vc_k_max``
+# caps modes per axis; ``vc_freq_coef`` scales the frequency ridge that
+# suppresses noise leakage into the non-constant energy.
+vc_modes_cache: dict = {}
+vc_k_max: int = 6
+vc_freq_coef: float = 1.0
+
+# When True, ``VaryingCoefSetup._solve_gammas`` solves the mode block
+# PER-FEATURE (block-diagonal in feature index) instead of jointly: cross-
+# feature mode collinearity is dropped so a true constant-coefficient term's
+# region-variation B (=nc_deb/C) is not inflated by collinear grid-modulated
+# cousins (``x*u_xx``/``sin*u_xx`` sharing ``u_xx``'s mode energy, which pushed
+# the weak true term's L1 threshold above its signal -> the ac t0/3 collapse).
+# Extends the existing Frisch-Waugh constant-block decoupling to the modes.
+# Default True; set False for the legacy joint mode solve.
+vc_mode_decouple: bool = True
+
+
+def set_gram_config(mode: str = 'vcoef'):
+    """Override the global Gram-construction mode before ``build_search``.
+
+    Used by ``projects/thesis/thesis_runner.py`` / ``profile_loop_stats.py``
+    to switch between the varying-coefficient default (``'vcoef'``) and the
+    axis-aligned sliding-window backup (``'axis'``) via a single CLI flag.
+    """
+    global gram_mode
+    if mode not in ('axis', 'vcoef'):
+        raise ValueError(
+            f'gram_mode must be "axis" or "vcoef"; got {mode!r}')
+    gram_mode = mode
+    # Stale per-axis basis resolution from a prior CLI/config must not bleed
+    # into a new invocation -- the source data or grid_shape may have changed.
+    vc_modes_cache.clear()
+
+
 def init_caches(set_grids: bool = False, device = 'cpu'):
     """
     Initialization global variables for keeping input data, values of grid and useful tensors such as evaluated terms

epde/interface/interface.py

@@ -22,6 +22,8 @@
 
 import epde.globals as global_var
 
+from epde import _loop_stats
+
 from epde.optimizers.builder import StrategyBuilder
 from epde.optimizers.builder import OptimizationPatternDirector
 
@@ -754,8 +756,9 @@ def saved_derivaties(self):
             print('Trying to get derivatives before their calculation. Call EPDESearch.create_pool() to calculate derivatives')
             return None
 
+    @_loop_stats.timed('EpdeSearch.fit')
     def fit(self, data: Union[np.ndarray, list, tuple] = None, equation_terms_max_number=6,
-            equation_factors_max_number=1, variable_names=['u',], eq_sparsity_interval=(1e-4, 2.5), 
+            equation_factors_max_number=1, variable_names=['u',], eq_sparsity_interval=(1e-4, 2.5),
             derivs=None, max_deriv_order=1, additional_tokens = None, data_fun_pow: int = 1, deriv_fun_pow: int = 1,
             optimizer: Union[SimpleOptimizer, MOEADDOptimizer] = None, pool: TFPool = None,
             population: List[SoEq] = None, data_nn = None, ann_epochs_max = 1e5,