Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 125 additions & 3 deletions epde/_loop_stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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.
"""
Expand All @@ -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:
Expand All @@ -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:
Expand Down
52 changes: 52 additions & 0 deletions epde/globals.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion epde/interface/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading