diff --git a/epde/_loop_stats.py b/epde/_loop_stats.py index 2c05bd56..4556f4c2 100644 --- a/epde/_loop_stats.py +++ b/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: diff --git a/epde/globals.py b/epde/globals.py index 2978c2ee..dcdb6d3b 100644 --- a/epde/globals.py +++ b/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 diff --git a/epde/interface/interface.py b/epde/interface/interface.py index 9291cd41..52d13868 100644 --- a/epde/interface/interface.py +++ b/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, diff --git a/epde/operators/common/fitness.py b/epde/operators/common/fitness.py index 6fa9e524..bf13ee8b 100644 --- a/epde/operators/common/fitness.py +++ b/epde/operators/common/fitness.py @@ -23,7 +23,18 @@ from sklearn.linear_model import LinearRegression, Ridge from scipy.optimize import minimize from epde.supplementary import minmax_normalize -from epde.supplementary import calculate_weights +from epde.operators.common.stability import (calculate_weights, vc_stability_total_lr) +from epde import _loop_stats + + +def _gram_dispatch_kwargs(): + """Return ``(gram_cls, gram_kwargs)`` for ``calculate_weights``. The + axis backup path uses ``GramSetup`` (``calculate_weights``'s default), + so this is always ``(None, None)``; the ``vcoef`` default does not + route through ``calculate_weights`` (it scores via ``VaryingCoefSetup`` + directly in ``PhysicsInformedLasso.fit``). + """ + return (None, None) LOSS_NAN_VAL = 1e7 @@ -54,6 +65,7 @@ class L2Fitness(CompoundOperator): key = 'DiscrepancyBasedFitness' + @_loop_stats.timed('L2Fitness.apply') def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool = False): """ Calculate the fitness function values. The result is not returned, but stored in the equation.fitness_value attribute. @@ -142,20 +154,34 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool = _, sw_target, sw_features = objective.evaluate(normalize=True, return_val=False) if sw_features is None: total_lr = 1.0 + elif global_var.gram_mode == 'vcoef': + if getattr(objective, '_cached_vc_score', None) is not None: + total_lr = float(np.sum(objective._cached_vc_score)) + else: + # Non-zero terms only; exclude the intercept when it is zeroed + # (weights_final[-1] == 0) -- same policy as the axis path. + _, t_nz, f_nz = objective.evaluate(normalize=False, return_val=False) + total_lr = (1.0 if f_nz is None else vc_stability_total_lr( + f_nz, t_nz, self.g_fun_vals, data_shape, + main_var=objective.main_var_to_explain, + fit_intercept=objective.weights_final[-1] != 0)) else: if hasattr(objective, '_cached_sw_weights') and objective._cached_sw_weights is not None: sw_weights = objective._cached_sw_weights else: + _gc, _gk = _gram_dispatch_kwargs() sw_weights = calculate_weights( sw_features, sw_target, self.g_fun_vals, data_shape, objective.weights_final[-1] != 0, + gram_cls=_gc, gram_kwargs=_gk, ) sw_arr = np.array(sw_weights) - std = sw_arr.std(axis=0, ddof=1) mu = sw_arr.mean(axis=0) + std = sw_arr.std(axis=0, ddof=1) with np.errstate(divide='ignore', invalid='ignore'): cv = (std ** 2) / (mu ** 2) - total_lr = sum(cv) / len(data_shape) + cv[mu == 0] = 0.0 + total_lr = sum(np.nan_to_num(cv)) / len(data_shape) except Exception: total_lr = 1.0 objective.stability_calculated = True @@ -168,6 +194,7 @@ def use_default_tags(self): class L2LRFitness(CompoundOperator): key = 'DiscrepancyBasedFitnessWithCV' + @_loop_stats.timed('L2LRFitness.apply') def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool = False): """ Calculate the fitness function values. The result is not returned, but stored in the equation.fitness_value attribute. @@ -186,7 +213,7 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool = if force_out_of_place: self.suboperators['sparsity'].apply(objective, subop_args['sparsity']) - if all(objective.weights_internal == 0): + if all(objective.weights_internal[:-1] == 0): return None self.suboperators['coeff_calc'].apply(objective, subop_args['coeff_calc']) @@ -203,22 +230,11 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool = if features is None: discr = target - target.mean() else: - # ``features`` width depends on the ``normalize`` flag passed to - # ``evaluate`` above: ``normalize=True`` returns all N-1 - # non-target columns; ``normalize=False`` filters to only the - # nonzero-weight columns. ``weights_final[:-1]`` matches the - # latter shape (nonzero count); ``weights_internal`` matches the - # former (full N-1, with zeros). Pick whichever lines up with - # the actual feature matrix -- same pattern as L2Fitness.apply. - n_cols = features.shape[1] if features.ndim > 1 else 1 - mask = objective.weights_internal != 0 - if n_cols == len(mask): - discr_feats = np.dot(features, objective.weights_internal) - elif n_cols == int(mask.sum()): + if objective.weights_internal[-1]: discr_feats = np.dot(features, objective.weights_final[:-1]) + discr_feats = discr_feats + objective.weights_final[-1] else: - discr_feats = np.zeros(features.shape[0]) - discr_feats = discr_feats + objective.weights_final[-1] + discr_feats = np.dot(features, objective.weights_final) discr = target - discr_feats rl_error = np.sum(np.abs(discr)) / np.sum(np.abs(target)) @@ -235,23 +251,31 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool = objective.aic_calculated = True data_shape = global_var.grid_cache.inner_shape - if features is None: - # Degenerate candidate (all features pruned by sparsity). - # Nothing to fit sliding-window weights on -- skip the CV - # calculation and report unit stability so downstream callers - # still get a finite value. - total_lr = 1.0 + if global_var.gram_mode == 'vcoef': + if getattr(objective, '_cached_vc_score', None) is not None: + total_lr = float(np.sum(objective._cached_vc_score)) + else: + total_lr = vc_stability_total_lr( + features, target, self.g_fun_vals, data_shape, + main_var=objective.main_var_to_explain, + fit_intercept=objective.weights_internal[-1] != 0) else: if hasattr(objective, '_cached_sw_weights') and objective._cached_sw_weights is not None: weights = objective._cached_sw_weights else: - weights = calculate_weights(features, target, self.g_fun_vals, data_shape, objective.weights_final[-1] != 0) + _gc, _gk = _gram_dispatch_kwargs() + weights = calculate_weights( + features, target, self.g_fun_vals, data_shape, + objective.weights_internal[-1] != 0, + gram_cls=_gc, gram_kwargs=_gk, + ) weights_arr = np.array(weights) - std = weights_arr.std(axis=0, ddof=1) mu = weights_arr.mean(axis=0) + std = weights_arr.std(axis=0, ddof=1) with np.errstate(divide='ignore', invalid='ignore'): cv = (std ** 2) / (mu ** 2) - total_lr = sum(cv) / len(data_shape) + cv[mu == 0] = 0.0 + total_lr = sum(np.nan_to_num(cv)) / len(data_shape) # if force_out_of_place: # return fitness_value * total_lr @@ -430,19 +454,33 @@ def apply(self, objective: SoEq, arguments: dict, force_out_of_place: bool = Fal # Calculate r-loss data_shape = global_var.grid_cache.inner_shape _, target, features = eq.evaluate(normalize=True, return_val=False) - if hasattr(eq, '_cached_sw_weights') and eq._cached_sw_weights is not None: - weights = eq._cached_sw_weights + if global_var.gram_mode == 'vcoef': + if getattr(eq, '_cached_vc_score', None) is not None: + total_lr = float(np.sum(eq._cached_vc_score)) + else: + total_lr = vc_stability_total_lr( + features, target, self.g_fun_vals, data_shape, + main_var=objective.main_var_to_explain, + fit_intercept=objective.weights_internal[-1] != 0) else: - weights = calculate_weights(features, target, self.g_fun_vals, data_shape) - weights_arr = np.array(weights) - std = weights_arr.std(axis=0, ddof=1) - mu = weights_arr.mean(axis=0) + if hasattr(eq, '_cached_sw_weights') and eq._cached_sw_weights is not None: + weights = eq._cached_sw_weights + else: + _gc, _gk = _gram_dispatch_kwargs() + weights = calculate_weights( + features, target, self.g_fun_vals, data_shape, + gram_cls=_gc, gram_kwargs=_gk, + ) + weights_arr = np.array(weights) + mu = weights_arr.mean(axis=0) + std = weights_arr.std(axis=0, ddof=1) - # Safe division - with np.errstate(divide='ignore', invalid='ignore'): - cv = (std ** 2) / (mu ** 2) + # Safe division + with np.errstate(divide='ignore', invalid='ignore'): + cv = (std ** 2) / (mu ** 2) + cv[mu == 0] = 0.0 - total_lr = sum(cv) / len(data_shape) + total_lr = sum(np.nan_to_num(cv)) / len(data_shape) eq.fitness_calculated = True eq.fitness_value = lp @@ -570,12 +608,30 @@ def _compute_stability_for_equation(self, eq: Equation): _, target, features = eq.evaluate(normalize=False, return_val=False) data_shape = global_var.grid_cache.inner_shape self.get_g_fun_vals() - weights = calculate_weights(features, target, self.g_fun_vals, data_shape) - weights_arr = np.array(weights) - std = weights_arr.std(axis=0, ddof=1) - mu = weights_arr.mean(axis=0) - cv = (std ** 2) / (mu ** 2) - total_lr = np.sum(cv) / len(data_shape) + if global_var.gram_mode == 'vcoef': + if getattr(eq, '_cached_vc_score', None) is not None: + total_lr = float(np.sum(eq._cached_vc_score)) + else: + total_lr = vc_stability_total_lr( + features, target, self.g_fun_vals, data_shape, + main_var=objective.main_var_to_explain, + fit_intercept=objective.weights_internal[-1] != 0) + else: + if hasattr(eq, '_cached_sw_weights') and eq._cached_sw_weights is not None: + weights = eq._cached_sw_weights + else: + _gc, _gk = _gram_dispatch_kwargs() + weights = calculate_weights( + features, target, self.g_fun_vals, data_shape, + gram_cls=_gc, gram_kwargs=_gk, + ) + weights_arr = np.array(weights) + mu = weights_arr.mean(axis=0) + std = weights_arr.std(axis=0, ddof=1) + with np.errstate(divide='ignore', invalid='ignore'): + cv = (std ** 2) / (mu ** 2) + cv[mu == 0] = 0.0 + total_lr = np.sum(np.nan_to_num(cv)) / len(data_shape) eq.coefficients_stability = total_lr eq.stability_calculated = True diff --git a/epde/operators/common/right_part_selection.py b/epde/operators/common/right_part_selection.py index 363fd408..365d7f8f 100644 --- a/epde/operators/common/right_part_selection.py +++ b/epde/operators/common/right_part_selection.py @@ -15,7 +15,7 @@ from epde.operators.utils.template import CompoundOperator from epde.decorators import HistoryExtender from epde.structure.main_structures import Term, Equation -from epde.supplementary import GramSetup, retry_until_unique +from epde.operators.common.stability import (GramSetup, VaryingCoefSetup) from epde import _loop_stats class EqRightPartSelector(CompoundOperator): @@ -44,17 +44,8 @@ class EqRightPartSelector(CompoundOperator): ''' key = 'FitnessCheckingRightPartSelector' - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - # Negative cache: structure hashes that produced ``inf`` fitness on - # every eligible target_idx. Scope is one operator-tree lifetime, - # which matches one ``EpdeSearch.fit()`` call -- the strategy - # builder in ``epde/optimizers/moeadd/strategy.py`` constructs a - # fresh ``EqRightPartSelector()`` per build. Repeat encounters - # skip the term-sweep via internal ``objective.randomize()``. - self._bad_structures: set = set() - @staticmethod + @_loop_stats.timed('EqRPS.gram_super') def _precompute_super_gram(objective: Equation) -> None: """Build a per-equation super-Gram over all structure terms and attach it to ``objective._gram_super`` for the upcoming term-sweep. @@ -79,8 +70,13 @@ def _precompute_super_gram(objective: Equation) -> None: objective._gram_super = None _loop_stats.record('EqRPS.gram_super_skip', 1, 1) return - objective._gram_super = GramSetup.precompute_super( - Z, sample_weights, grid_shape) + if global_var.gram_mode == 'vcoef': + objective._gram_super = VaryingCoefSetup.precompute_super( + Z, sample_weights, grid_shape, + main_var=objective.main_var_to_explain) + else: # 'axis' backup + objective._gram_super = GramSetup.precompute_super( + Z, sample_weights, grid_shape) _loop_stats.record('EqRPS.gram_super_built', 1, 1) except Exception: # Defensive: any unexpected failure (shape mismatch, missing @@ -89,6 +85,7 @@ def _precompute_super_gram(objective: Equation) -> None: objective._gram_super = None _loop_stats.record('EqRPS.gram_super_skip', 1, 1) + @_loop_stats.timed('EqRPS.apply') @HistoryExtender('\n -> The equation structure was detected: ', 'a') def apply(self, objective : Equation, arguments : dict): """Select a right-part term for ``objective`` in-place. @@ -144,49 +141,41 @@ def apply(self, objective : Equation, arguments : dict): objective.restore_property(mandatory_family=False, deriv=True) _loop_stats.record('EqRPS.inner_derivative', inner_attempts, inner_max_iter) - # Negative cache: this structure already zeroed out for every - # target_idx on a prior call (sparsity is deterministic on - # the same input). Skip the term-sweep, reroll, retry. - if objective.terms_labels in self._bad_structures: - _loop_stats.record('EqRPS.bad_structure_skip', 1, 1) - objective.randomize() - continue - # Tier 3: precompute the super-Gram over all terms ONCE per # outer iter so the term-sweep below derives per-target # GramSetup views via pure slicing instead of rebuilding the # windowed XTWX matmul for every candidate. self._precompute_super_gram(objective) - for target_idx, target_term in enumerate(objective.structure): - if not objective.structure[target_idx].contains_deriv(objective.main_var_to_explain): - continue - objective.target_idx = target_idx - fitness = self.suboperators['fitness_calculation'].apply(objective, arguments = subop_args['fitness_calculation'], force_out_of_place = True) - if fitness is not None and fitness < min_fitness: - min_fitness = fitness - min_idx = target_idx - weights_internal = objective.weights_internal - weights_final = objective.weights_final - sw_weights = objective._cached_sw_weights - - objective.weights_internal_evald = False - objective.weights_final_evald = False + with _loop_stats.timer('EqRPS.term_sweep'): + for target_idx, target_term in enumerate(objective.structure): + if not objective.structure[target_idx].contains_deriv(objective.main_var_to_explain): + continue + objective.target_idx = target_idx + fitness = self.suboperators['fitness_calculation'].apply(objective, arguments = subop_args['fitness_calculation'], force_out_of_place = True) + if fitness is not None and fitness < min_fitness: + min_fitness = fitness + min_idx = target_idx + weights_internal = objective.weights_internal + weights_final = objective.weights_final + sw_weights = objective._cached_sw_weights + vc_score = objective._cached_vc_score + + objective.weights_internal_evald = False + objective.weights_final_evald = False if np.isinf(min_fitness): # Every eligible target produced inf fitness for the # post-restore structure -- reroll this single equation - # locally (cheap) and continue the outer loop. The - # negative cache below remembers the structure shape so - # future outer iters / future calls skip the term-sweep. + # locally (cheap) and continue the outer loop. _loop_stats.record('EqRPS.inf_fitness_regen', 1, 1) - self._bad_structures.add(objective.terms_labels) objective.randomize() continue objective.weights_internal = weights_internal objective.weights_final = weights_final objective._cached_sw_weights = sw_weights + objective._cached_vc_score = vc_score objective.weights_internal_evald = True objective.weights_final_evald = True objective.target_idx = min_idx @@ -204,6 +193,13 @@ def apply(self, objective : Equation, arguments : dict): objective._gram_super = None objective.right_part_selected = True objective.remove_zero_terms() + # Hard invariant: no duplicate terms may leave RPS. simplify and + # scrub both regenerate-then-drop, so a surviving duplicate is a + # logic error to surface HERE -- not one generation later at the + # crossover/mutation assert (the crash site that "lies"). + _final_sigs = {term.factors_labels for term in objective.structure} + assert len(_final_sigs) == len(objective.structure), \ + 'EqRightPartSelector.apply: duplicate terms survived RPS.' def simplify_equation(self, objective: Equation): # Get nonzero terms @@ -215,6 +211,59 @@ def simplify_equation(self, objective: Equation): if len(equation_terms) <= 1: return False + + # Degree reduction: when a SINGLE non-target term remains, the + # equation is ``coef * f = g`` with f, g products of powered + # factors. If every factor power in BOTH f and g shares a common + # divisor p >= 2, the whole equation is a p-th power -- take the + # p-th root (divide every power by p; the coefficient is recomputed + # downstream). E.g. ``c*(u_xx)^2 = (u_tt)^2`` -> ``sqrt(c)*u_xx = u_tt``. + # Keeps the lowest-degree equivalent form so it is not penalised / + # mistaken for a distinct higher-order structure. + if len(nonzero_terms) == 2: + powers, integral = [], True + for term in nonzero_terms: + for factor in term.structure: + for i in factor.params_description: + if factor.params_description[i]["name"] == "power": + p = factor.params[i] + if float(p) != int(p) or int(p) < 1: + integral = False + powers.append(int(p)) + if integral and powers: + root = int(np.gcd.reduce(np.array(powers, dtype=int))) + if root >= 2: + # p-th root: divide every factor power by the gcd, + # collapsing the equation to its lowest equivalent + # degree (c*(u_xx)^2=(u_tt)^2 -> sqrt(c)*u_xx=u_tt). + for term in nonzero_terms: + for factor in term.structure: + for i in factor.params_description: + if factor.params_description[i]["name"] == "power": + factor.set_param(int(factor.params[i]) // root, idx=i) + term.reset_saved_state() + # The reduction can collapse a survivor onto a zero-weight + # candidate already in the structure (u^2 -> u when a u + # term exists). Such a colliding copy is ALWAYS a + # zero-weight non-survivor -- two genuinely nonzero terms + # cannot collide via a p-th root unless they were already + # equal, which the entry assert forbids -- so drop the + # redundant copies outright (drop-immediately). The reduced + # low-degree form survives on the kept terms; no revert. + keep_ids = {id(t) for t in nonzero_terms} + kept_labels = {t.factors_labels for t in nonzero_terms} + redundant = [t for t in objective.structure + if id(t) not in keep_ids + and t.factors_labels in kept_labels] + for t in redundant: + _regen_or_drop_term( + objective, t, max_iter=0, + stats_name='simplify_equation.degree_reduction') + try: + objective.reset_state(reset_right_part=False) + except TypeError: + objective.reset_state() + return True common_factors = list(frozenset.intersection(*equation_terms)) if not common_factors: return False @@ -244,27 +293,20 @@ def simplify_equation(self, objective: Equation): term.structure = [factor for factor in term.structure if factor not in factors_simplified] term.reset_saved_state() - # If term's order became zero -- replace term. - # Cap retries so a constrained token pool can't - # deadlock the optimizer (same hazard fixed in - # ``enforce_rps_uniqueness``). - def _replacement_acceptable(): - empty = len(term.structure) == 0 - not_meaningful = not term.contains_meaningful() - signatures = {t.factors_labels for t in objective.structure} - duplicate = len(signatures) != len(objective.structure) - return not (empty or not_meaningful or duplicate) - - # Cap-hit policy is silently-accept-whatever-state: the - # term may still be empty/non-meaningful/duplicate when - # the cap fires, and the outer RPS loop is expected to - # detect that on its next pass. - retry_until_unique( - predicate=_replacement_acceptable, - mutate=lambda: term.randomize(), - max_iter=max_iter, - stats_name='simplify_equation.replace_term', - ) + # If the term's order became zero (or it now duplicates + # another term), regenerate it; if the pool can't yield a + # unique, meaningful replacement within the cap, DROP it. + # A duplicate must never ride out of RPS -- see the exit + # assert in ``apply``. + status = _regen_or_drop_term( + objective, term, max_iter=max_iter, + stats_name='simplify_equation.replace_term') + if status in ('target', 'floor'): + # Offending term is the RPS target, or dropping would + # degenerate the equation -> decline this + # simplification and let the outer RPS loop reset and + # re-select. + return False # Structure changed: invalidate stale fitness / # weights / AIC caches while leaving RPS to the @@ -358,6 +400,83 @@ def use_default_tags(self): self._tags = {'equation right part selection', 'gene level', 'contains suboperators', 'inplace'} +def _regen_or_drop_term(equation: Equation, term, *, max_iter: int = 100, + min_terms: int = 2, + stats_name: str = 'simplify_equation.regen_or_drop') -> str: + """Make ``equation.structure`` unique w.r.t. ``term`` by regenerating + ``term`` up to ``max_iter`` times; if it is still empty / non-meaningful + / a duplicate, DROP it from the structure. + + This is the simplify/scrub cap-hit policy -- *regenerate-n-then-drop* -- + deliberately distinct from the *keep-or-revert* ``retry_until_unique`` + policy used by ``Equation.__init__`` and the mutation operators. + ``max_iter == 0`` means "drop immediately if unacceptable" (no + regeneration) -- used by the degree-reduction branch and the scrub + duplicate gate. + + The acceptability predicate ranges over the FULL structure, so a + duplicate against a zero-weight candidate elsewhere is caught too. The + RPS target is never dropped: if ``term`` is the target AND a duplicate, + the OTHER member of its duplicate group is dropped instead; if the + target is merely empty/non-meaningful, ``'target'`` is returned for the + caller to handle. Refuses to drop below ``min_terms`` (returns + ``'floor'``). On a drop, ``target_idx`` is reindexed exactly as in + ``Equation.remove_zero_terms``. + + Returns one of ``'ok'`` (already acceptable), ``'regenerated'``, + ``'dropped'``, ``'target'``, ``'floor'``. + """ + idx = next((j for j, t in enumerate(equation.structure) if t is term), None) + if idx is None: + return 'ok' # already dropped earlier in this pass + + def _acceptable(): + if len(term.structure) == 0 or not term.contains_meaningful(): + return False + signatures = {t.factors_labels for t in equation.structure} + return len(signatures) == len(equation.structure) + + cap = max_iter if max_iter > 0 else 1 + if _acceptable(): + _loop_stats.record(stats_name, 1, cap) + return 'ok' + + attempts = 0 + for _ in range(max_iter): + attempts += 1 + term.randomize() + term.reset_saved_state() + if _acceptable(): + _loop_stats.record(stats_name, attempts, cap) + equation._invalidate_label_cache() + return 'regenerated' + _loop_stats.record(stats_name, max(attempts, 1), cap) + + # Exhausted (or max_iter == 0): drop the offending term, if legal. + tgt = getattr(equation, 'target_idx', None) + drop_idx = idx + if tgt is not None and idx == tgt: + # Can't drop the RPS target. If it duplicates another term, drop + # that other (non-target) member; if it is merely empty/non- + # meaningful, leave it for the caller to resolve. + my_label = term.factors_labels + other = next((j for j, t in enumerate(equation.structure) + if j != idx and t.factors_labels == my_label), None) + if other is None: + equation._invalidate_label_cache() + return 'target' + drop_idx = other + if len(equation.structure) <= min_terms: + equation._invalidate_label_cache() + return 'floor' + equation.structure = [t for j, t in enumerate(equation.structure) + if j != drop_idx] + if tgt is not None and drop_idx < tgt: + equation.target_idx -= 1 + equation._invalidate_label_cache() + return 'dropped' + + def _scrub_conflicting_terms(equation: Equation, fixed_rps, *, max_iter: int = 2000, skip_idx=None) -> bool: """Replace any term in ``equation.structure`` whose factor signature is a @@ -387,12 +506,13 @@ def _scrub_conflicting_terms(equation: Equation, fixed_rps, *, max_iter: int = 2 def _conflicts(t): return any(rs.issubset(t.factors_labels) for rs in fixed_rps) + # Snapshot the conflicting terms by identity: the randomize loop below + # never resizes ``structure``, but the duplicate-drop pass afterwards + # does, so index-based iteration would be unsafe. + conflicting = [term for idx, term in enumerate(equation.structure) + if idx != skip_idx and _conflicts(term)] changed = False - for idx, term in enumerate(equation.structure): - if idx == skip_idx: - continue - if not _conflicts(term): - continue + for term in conflicting: attempts = 0 for _ in range(max_iter): attempts += 1 @@ -406,6 +526,15 @@ def _conflicts(t): changed = True if changed: + # Cap-hit may leave a scrubbed term as a DUPLICATE (regenerate + # exhausted). A conflicting-but-unique term is tolerated -- the + # bidirectional outer loop re-selects -- but a duplicate must not + # ride out of RPS, so drop it (the n regenerate attempts were + # already spent in the loop above). The skip_idx term is excluded + # from ``conflicting`` and is never dropped. + for term in conflicting: + _regen_or_drop_term(equation, term, max_iter=0, + stats_name='scrub_conflicting_terms.drop') try: equation.reset_state(reset_right_part=False) except TypeError: @@ -436,6 +565,7 @@ class SoEqRightPartSelector(CompoundOperator): """ key = 'SoEqRightPartSelector' + @_loop_stats.timed('SoEqRPS.apply') def apply(self, objective, arguments: dict): """Run per-equation RPS forward + bidirectional passes in-place. diff --git a/epde/operators/common/sparsity.py b/epde/operators/common/sparsity.py index 961c8fbc..644e3432 100644 --- a/epde/operators/common/sparsity.py +++ b/epde/operators/common/sparsity.py @@ -16,7 +16,8 @@ from sklearn.base import BaseEstimator, RegressorMixin # import seaborn as sns import matplotlib.pyplot as plt -from epde.supplementary import calculate_weights, GramSetup +from epde.operators.common.stability import (calculate_weights, GramSetup, + VaryingCoefSetup) from epde import _loop_stats @@ -233,10 +234,15 @@ class PhysicsInformedLasso(BaseEstimator, RegressorMixin): - Aggressive: Instant elimination of features that hit zero during optimization. """ - def __init__(self, max_iter=1000, tol=1e-4, grid_shape=None): + def __init__(self, max_iter=1000, tol=1e-4, grid_shape=None, + main_var: str = None): self.max_iter = max_iter self.tol = tol self.grid_shape = grid_shape + # Threaded through to ``VaryingCoefSetup`` so the basis-mode + # resolver picks the equation's own primary variable when + # multi-var systems use different scales per equation. + self.main_var = main_var self.coef_ = None self.full_coef_ = None # Includes the intercept @@ -244,26 +250,26 @@ def _soft_threshold(self, x, lambda_): return np.sign(x) * np.maximum(np.abs(x) - lambda_, 0.0) def get_cv(self, weights): - """Calculates Squared Coefficient of Variation (std^2 / mean^2).""" - weights_arr = np.array(weights) - std = weights_arr.std(axis=0, ddof=1) - mu = weights_arr.mean(axis=0) - + """Per-feature CV-stability metric for the axis backup path: + ``(std / mean)^2 = var / mu^2`` across the sliding windows. + + The squared coefficient of variation of each feature's per-window + weight. It blows up (large CV) for features whose fitted coefficient + is unstable or near-zero-mean across horizons, so the + ``active_thresholds = cv * max_corr`` step in + :meth:`PhysicsInformedLasso.fit` prunes them first. The default + ``gram_mode='vcoef'`` path does not call this -- it scores via + ``VaryingCoefSetup.score`` instead. + """ + weights_arr = np.asarray(weights) with np.errstate(divide='ignore', invalid='ignore'): - cv = std ** 2 / mu ** 2 - # cv = std ** 2 - + std = weights_arr.std(axis=0, ddof=1) + mu = weights_arr.mean(axis=0) + cv = (std ** 2) / (mu ** 2) + cv[mu == 0] = 0.0 return np.nan_to_num(cv) - # def get_cv(self, weights): - # weights_arr = np.asarray(weights) - # q1, q3 = np.percentile(weights_arr, [25, 75], axis=0) - # spread = (q3 - q1) / 1.349 # IQR/1.349 ≈ σ for Gaussian - # center = np.median(weights_arr, axis=0) - # with np.errstate(divide='ignore', invalid='ignore'): - # cv = spread ** 2 / (center ** 2 + spread ** 2) - # return np.nan_to_num(cv) - + @_loop_stats.timed('PhysicsInformedLasso.fit') def fit(self, X, y, sample_weights=None, gram_setup=None): n_samples, n_features = X.shape @@ -290,7 +296,17 @@ def fit(self, X, y, sample_weights=None, gram_setup=None): # super-Gram, reuse it -- saves the windowed matmul that # otherwise repeats for every candidate target_idx in one sweep. if gram_setup is None: - gram_setup = GramSetup(X, y, sample_weights, self.grid_shape) + if global_var.gram_mode == 'vcoef': + gram_setup = VaryingCoefSetup( + X, y, sample_weights, self.grid_shape, + main_var=self.main_var) + else: # 'axis' backup + gram_setup = GramSetup(X, y, sample_weights, self.grid_shape) + + # Varying-coefficient mode returns a per-feature stability score + # directly (no per-window weight stack), so the in-fit CV-threshold + # path branches on it below. + is_vcoef = getattr(gram_setup, 'is_vcoef', False) outer_iteration = 0 max_outer_iters = total_features # Max possible eliminations @@ -308,7 +324,9 @@ def fit(self, X, y, sample_weights=None, gram_setup=None): # 2. Calculate physical priors ONLY for the active library -- # slice the precomputed full Gram by the current active mask. - weights = gram_setup.solve(active_mask) + # ``vcoef`` yields the per-feature score directly; the axis + # path returns a per-window weight stack reduced by ``get_cv``. + weights = None if is_vcoef else gram_setup.solve(active_mask) # Slice data for the CD run X_active = X_aug[:, active_mask] @@ -318,16 +336,39 @@ def fit(self, X, y, sample_weights=None, gram_setup=None): # so threshold scale tracks the current problem as features drop. max_corr = np.max(np.abs(X_T_y[active_mask])) - # 3. CV performs as adaptive alpha - active_cv = self.get_cv(weights) + # 3. CV performs as adaptive alpha. In vcoef mode this is each + # term's stability score (Var(gamma_0) + NC_deb)/gamma_0^2 -- + # significance plus debiased region-variation -- which prunes weak / + # zero / unstable / spuriously-varying terms. + active_cv = (gram_setup.score(active_mask) if is_vcoef + else self.get_cv(weights)) + # Tackle the most physically unstable feature first so unstable # terms get shrunk to zero before they pollute the residual. - cv_order = np.argsort(active_cv)[::-1] active_thresholds = active_cv * max_corr + # active_thresholds = active_cv # active_thresholds = active_cv * norm_sq_active - # Initialize coefficients - active_coef = weights.mean(axis=0) + cv_order = np.argsort(active_cv)[::-1] + # cv_order = np.argsort(active_thresholds)[::-1] + + # Initialize coefficients from a single global weighted-OLS on + # the full dataset rather than the mean over per-window OLS + # coefficients. The per-window mean is biased toward zero on + # heterogeneous data (e.g. KdV solitons, where most windows see + # ~0 signal so the mean is shrunk by ``M_signal / M``); the + # global OLS is unbiased. + sw_active = (sample_weights if sample_weights is not None + else np.ones(n_samples)) + try: + XTWX_full = X_active.T @ (sw_active[:, None] * X_active) + XTWy_full = X_active.T @ (sw_active * y) + active_coef = np.linalg.solve(XTWX_full, XTWy_full) + except np.linalg.LinAlgError: + # ``vcoef`` has no per-window stack to average; fall back to + # a zero start (CD recovers it) instead of weights.mean. + active_coef = (np.zeros(int(active_mask.sum())) if is_vcoef + else weights.mean(axis=0)) residual = y - (X_active @ active_coef) @@ -407,6 +448,27 @@ def fit(self, X, y, sample_weights=None, gram_setup=None): _loop_stats.record('PhysicsInformedLasso.RFE_outer', outer_iters_executed, max_outer_iters) self.cached_weights_ = weights + # Per-active-term stability scores on the converged mask, summed as the + # stability objective in fitness. ``None`` for the axis backup path. + self.cached_vc_score_ = (gram_setup.score(active_mask) + if is_vcoef else None) + + # Relaxed-LASSO refit: replace the surviving CD-output + # coefficients with a single global weighted-OLS on + # ``X[:, active_mask]``. Sparsity decisions (which features + # survived) are preserved; only the magnitudes become unbiased + # global estimates. + if np.any(active_mask): + sw_active = (sample_weights if sample_weights is not None + else np.ones(n_samples)) + X_final = X_aug[:, active_mask] + try: + XTWX_final = X_final.T @ (sw_active[:, None] * X_final) + XTWy_final = X_final.T @ (sw_active * y) + refit = np.linalg.solve(XTWX_final, XTWy_final) + self.full_coef_[active_mask] = refit + except np.linalg.LinAlgError: + pass # singular -> keep CD result # Map back to standard sklearn attributes self.coef_ = self.full_coef_[:-1] @@ -442,7 +504,8 @@ class LASSOSparsity(CompoundOperator): """ key = 'LASSOBasedSparsity' - + + @_loop_stats.timed('LASSOSparsity.apply') def apply(self, objective : Equation, arguments : dict): """ Apply the operator, to fit the LASSO regression to the equation object to detect the @@ -538,10 +601,13 @@ class VWSRSparsity(CompoundOperator): """ key = 'VWSRBasedSparsity' + @_loop_stats.timed('VWSRSparsity.apply') def apply(self, objective : Equation, arguments : dict): self_args, subop_args = self.parse_suboperator_args(arguments = arguments) - estimator = PhysicsInformedLasso(grid_shape=global_var.grid_cache.inner_shape) + estimator = PhysicsInformedLasso( + grid_shape=global_var.grid_cache.inner_shape, + main_var=objective.main_var_to_explain) self.g_fun_vals = global_var.grid_cache.g_func[global_var.grid_cache.g_func_mask] @@ -557,16 +623,20 @@ def apply(self, objective : Equation, arguments : dict): target = Z[:, t] feature_indexes = [i for i in range(Z.shape[1]) if i != t] features = Z[:, feature_indexes] - gram_setup = GramSetup.from_full(gram_super, t) + if gram_super.get('mode') == 'vcoef': + gram_setup = VaryingCoefSetup.from_full(gram_super, t) + else: + gram_setup = GramSetup.from_full(gram_super, t) else: _, target, features = objective.evaluate(normalize=True, return_val=False) gram_setup = None estimator.fit(features, target, self.g_fun_vals, gram_setup=gram_setup) - objective.weights_internal = estimator.coef_ + objective.weights_internal = np.array([*estimator.coef_, estimator.intercept_]) objective.weights_internal_evald = True - objective.weights_final = np.append([weight for weight in estimator.coef_ if weight != 0], estimator.intercept_) + objective.weights_final = np.array([weight for weight in objective.weights_internal if weight != 0]) objective.weights_final_evald = True objective._cached_sw_weights = estimator.cached_weights_ + objective._cached_vc_score = estimator.cached_vc_score_ # See LASSOSparsity.apply: _eval_cache survives a weights update; # only structural resets via ``Equation.reset_state`` should wipe it. diff --git a/epde/operators/common/stability.py b/epde/operators/common/stability.py new file mode 100644 index 00000000..010ea71f --- /dev/null +++ b/epde/operators/common/stability.py @@ -0,0 +1,911 @@ +"""Stability-estimation subsystem for the EPDE structural search. + +Two coefficient-stability estimators used by the sparsity / fitness / +right-part operators to score and prune candidate library terms: + +* ``GramSetup`` -- axis-aligned sliding-window CV (the ``gram_mode='axis'`` + backup), with ``calculate_weights`` as a single-shot wrapper. +* ``VaryingCoefSetup`` -- the default ``gram_mode='vcoef'`` varying-coefficient + stability, summed by ``vc_stability_total_lr``; basis resolution resolved + by ``resolve_vc_modes_from_input`` / ``taylor_microscale``. + +Relocated verbatim from ``epde.supplementary``. ``import epde.globals`` is +kept lazy (inside the bodies that use it) to avoid a circular import: +globals.py imports utilities from supplementary at module top level. +""" +import numpy as np +from numpy.lib.stride_tricks import sliding_window_view + +from epde import _loop_stats + + +# Default behaviour for GramSetup's sliding-window CV: ``True`` treats +# each axis as periodic so windows near the boundary wrap around (the +# last few windows span data[end-k:end] + data[0:window_size-k]). +# ``False`` is the legacy linear semantics where the number of windows +# along an axis is ``N - window_size + 1`` and no wrap occurs. +_DEFAULT_CIRCULAR_CV = True + + +def _windowed_take(arr: np.ndarray, dim: int, window_size: int, + num_horizons: int, step_size: int, + circular: bool) -> np.ndarray: + """Return ``sliding_window_view`` along ``dim``, optionally with + circular padding so windows near the boundary wrap to the start. + + Caller computes ``num_horizons`` (the number of valid start positions + along ``dim``) and ``step_size`` (subsampling stride). Under circular + mode, ``num_horizons == arr.shape[dim]`` and the input is padded by + ``window_size - 1`` samples copied from the start at the end via + ``np.pad(..., mode='wrap')``. Under linear mode, ``num_horizons == + arr.shape[dim] - window_size + 1`` and no padding is applied. + """ + if circular: + pad = [(0, 0)] * arr.ndim + pad[dim] = (0, window_size - 1) + arr = np.pad(arr, pad, mode='wrap') + windows = sliding_window_view(arr, window_shape=window_size, axis=dim) + return windows.take(indices=range(0, num_horizons, step_size), axis=dim) + + +def _cholesky_solve_batched(A, b): + """Solve ``A @ x = b`` batched over the leading axis using Cholesky. + + ``A`` is assumed symmetric positive-definite (shape ``(batch, n, n)``); + ``b`` is the RHS ``(batch, n, 1)``. Returns ``(x, L)`` where ``x`` is + the solution and ``L`` is the lower-triangular factor (so the caller + can reuse it for iterative refinement). If Cholesky fails on any batch + entry, returns ``(None, None)`` to signal "use the lstsq fallback". + + numpy doesn't ship a batched triangular solver, so the two triangular + solves go through ``np.linalg.solve`` -- still SPD-stable and ~1.5x + faster than feeding the full ``A`` to ``np.linalg.solve``. + """ + try: + L = np.linalg.cholesky(A) + except np.linalg.LinAlgError: + return None, None + try: + z = np.linalg.solve(L, b) + x = np.linalg.solve(L.transpose(0, 2, 1), z) + except np.linalg.LinAlgError: + return None, L + return x, L + + +def _per_batch_lstsq(A, b): + """Per-batch SVD-based least-squares solve. Used as the safety net + when Cholesky reports the equilibrated batch is non-SPD. Returns + weights of shape ``(batch, n, 1)`` matching the input RHS layout so + the caller can compose with subsequent matrix products without + reshaping. + """ + batch_size = A.shape[0] + n = A.shape[1] + out = np.empty((batch_size, n, 1)) + for i in range(batch_size): + sol, *_ = np.linalg.lstsq(A[i], b[i, :, 0], rcond=None) + out[i, :, 0] = sol + return out + + +class GramSetup: + """Precomputed batched normal-equation matrices for fast active-mask + solves. Splits :func:`calculate_weights` into a setup phase (compute + ``X^T diag(w) X`` and ``X^T diag(w) y`` per window-batch per dimension, + using the FULL augmented feature matrix) and a solve phase (slice each + full Gram matrix by an active-feature mask and solve). The setup is + mask-independent; only the solve depends on which columns are active. + + Used by :class:`PhysicsInformedLasso.fit`, whose outer RFE loop calls + ``calculate_weights`` per shrinking column subset. With this split the + expensive ``X^T diag(w) X`` matmul runs ONCE per fit and each outer + iter only pays the cost of an (active × active) solve. The math is + exact: a sub-block of a Gram matrix equals the Gram of the + corresponding sub-columns. + """ + + def __init__(self, X, y, sample_weights, grid_shape, + circular_cv: bool = _DEFAULT_CIRCULAR_CV): + n_samples = X.shape[0] + # Always augment X with the intercept column so callers can toggle + # ``fit_intercept`` via the active mask's last bit rather than + # re-running setup. + X_aug = np.hstack([X, np.ones((n_samples, 1))]) + n_features_aug = X_aug.shape[1] + + X_grid = X_aug.reshape(*grid_shape, n_features_aug) + y_grid = y.reshape(*grid_shape) + sample_weights_grid = sample_weights.reshape(*grid_shape) + + self.n_features_aug = n_features_aug + self.grid_shape = grid_shape + self._per_dim = [] + + for dim in range(len(grid_shape)): + window_size = grid_shape[dim] // 2 + # Circular: every position along the axis is a valid window + # start (the input is virtually periodic); linear: only the + # first ``window_size + 1`` positions yield a full window. + num_horizons = grid_shape[dim] if circular_cv else window_size + 1 + step_size = max(1, num_horizons // 30) + + X_windows = _windowed_take(X_grid, dim, window_size, + num_horizons, step_size, circular_cv) + y_windows = _windowed_take(y_grid, dim, window_size, + num_horizons, step_size, circular_cv) + w_windows = _windowed_take(sample_weights_grid, dim, window_size, + num_horizons, step_size, circular_cv) + + X_windows = np.moveaxis(X_windows, dim, 0) + y_windows = np.moveaxis(y_windows, dim, 0) + w_windows = np.moveaxis(w_windows, dim, 0) + X_windows = np.moveaxis(X_windows, -2, -1) + + batch_size = X_windows.shape[0] + X_batch = X_windows.reshape(batch_size, -1, n_features_aug) + y_batch = y_windows.reshape(batch_size, -1) + weights_batch = w_windows.reshape(batch_size, -1, 1) + + XTW = X_batch.transpose(0, 2, 1) * weights_batch.transpose(0, 2, 1) + XTWX_full = XTW @ X_batch + XTWy_full = XTW @ y_batch[..., None] + + # Per-batch column scales for equilibration in :meth:`solve`. + # ``diag`` is the per-feature L2 norm squared (weighted) of the + # underlying X columns; ``sqrt`` brings it back to a column- + # norm scale. The ``1e-30`` floor is a degenerate-column guard + # (well below any meaningful data scale) so ``1/scale`` stays + # finite for near-zero columns. + diag = np.diagonal(XTWX_full, axis1=1, axis2=2) + scales = np.sqrt(np.maximum(np.abs(diag), 1e-30)) + + self._per_dim.append((XTWX_full, XTWy_full, scales)) + + def solve(self, active_mask=None, ridge_rel=None, ridge_floor=None): + """Solve the normal equations for the active-feature subset across + every window-batch in every spatial dimension. ``active_mask`` is a + length-``n_features_aug`` boolean array; pass ``None`` for the full + set (equivalent to the legacy ``fit_intercept=True`` path). Returns + weights of shape ``(total_windows_across_dims, active_count)``. + + Stability strategy (preserves the Gram-sub-block precompute trick): + + 1. **Column equilibration**: rescale columns by + ``1/sqrt(diag(XTWX))`` so the equilibrated Gram has unit + diagonals and a much smaller effective condition number than + the raw ``XTWX`` (which carries the squared condition number + of the underlying ``sqrt(W) X``). + 2. **Cholesky on the equilibrated SPD batch** (with batched LU + fallback if scipy's batched triangular solve isn't available + on this numpy). Cholesky has tighter backward error than LU + and is ~2x faster on SPD inputs. + 3. **One step of iterative refinement** on the original (un- + equilibrated) system, recovering 6-8 decimal digits that + normal-equation conditioning costs. + 4. **Per-batch lstsq safety net** for any window-batch where + Cholesky fails (non-SPD after equilibration -- rare). + + ``ridge_rel`` / ``ridge_floor`` are kept as no-op kwargs for + backward compatibility with callers from the previous adaptive- + ridge era; the equilibrated solve does not need a per-feature + ridge, only a tiny flat ``1e-10`` on the unit-diagonal matrix. + """ + if active_mask is None: + active_mask = np.ones(self.n_features_aug, dtype=bool) + active_size = int(active_mask.sum()) + + all_weights = [] + for XTWX_full, XTWy_full, scales_full in self._per_dim: + # Two-step boolean slice. Boolean indexing copies, so the + # result is a fresh array we can modify in place without + # corrupting the cached full Gram. + XTWX_a = XTWX_full[:, active_mask, :][:, :, active_mask] + XTWy_a = XTWy_full[:, active_mask, :] + s_a = scales_full[:, active_mask] # (batch, k) + inv_s = 1.0 / s_a # (batch, k) + + # Equilibrate: A = D^-1 XTWX D^-1, b = D^-1 XTWy. After this + # the diagonal of A is 1 by construction; the off-diagonals + # are the correlation coefficients between the underlying + # columns of sqrt(W) X. + A = XTWX_a * inv_s[:, :, None] * inv_s[:, None, :] + b = XTWy_a * inv_s[:, :, None] + + # Tiny flat ridge on the equilibrated diagonal (now ~1 by + # construction) to keep Cholesky well-defined when columns + # are exactly collinear. + idx = np.arange(active_size) + A[:, idx, idx] += 1e-10 + + batch_size = A.shape[0] + w_norm, L = _cholesky_solve_batched(A, b) + if w_norm is None: + # Cholesky failed somewhere in the batch; per-entry + # lstsq safety net on the equilibrated system. + w_norm = _per_batch_lstsq(A, b) + + # Iterative refinement on the ORIGINAL system to claw back + # digits lost to normal-equation condition squaring. + # w0 = D^-1 w_norm is the candidate solution in original + # coordinates; the residual r = XTWy - XTWX @ w0 measures + # how much it misses the original equation; the correction + # dw_norm solves the same equilibrated system on D^-1 r and + # is unscaled back to dw. + w0 = w_norm * inv_s[:, :, None] + r = XTWy_a - XTWX_a @ w0 + r_norm = r * inv_s[:, :, None] + if L is not None: + try: + z = np.linalg.solve(L, r_norm) + dw_norm = np.linalg.solve(L.transpose(0, 2, 1), z) + except np.linalg.LinAlgError: + dw_norm = _per_batch_lstsq(A, r_norm) + else: + dw_norm = _per_batch_lstsq(A, r_norm) + w = w0 + dw_norm * inv_s[:, :, None] + + all_weights.append(w.squeeze(-1)) + return np.vstack(all_weights) + + @classmethod + def precompute_super(cls, Z, sample_weights, grid_shape, + circular_cv: bool = _DEFAULT_CIRCULAR_CV): + """Build a per-dim super-Gram over ``Z_aug = column_stack(Z, ones)`` + once. Returns an opaque dict consumed by :meth:`from_full` to + derive per-target GramSetup views via pure slicing. + + Used by EqRightPartSelector's term-sweep: instead of rebuilding + the windowed XTWX matrix for each candidate target column (which + does the same reshape/window/matmul on the SAME underlying Z), + precompute once over the full Z and slice out target-specific + sub-blocks. The math is exact -- (Z[:, ~t])^T W Z[:, ~t] is the + sub-block of (Z_aug)^T W Z_aug at rows/cols ``~t U intercept``. + + ``Z`` is shape (n_samples, n_terms); ``sample_weights`` is the + flat per-sample weight vector; ``grid_shape`` is the same shape + ``GramSetup.__init__`` consumes. ``circular_cv`` mirrors the + ``__init__`` flag -- callers that flow through both paths + (PhysicsInformedLasso single-call + EqRPS super-Gram sweep) must + keep these values aligned for the per-target views to remain + sub-blocks of the super-Gram (the math requires identical window + sets across the two passes). + """ + n_samples, n_terms = Z.shape + Z_aug = np.hstack([Z, np.ones((n_samples, 1))]) + n_features_aug = n_terms + 1 + + Z_grid = Z_aug.reshape(*grid_shape, n_features_aug) + sw_grid = sample_weights.reshape(*grid_shape) + per_dim_super = [] + + for dim in range(len(grid_shape)): + window_size = grid_shape[dim] // 2 + num_horizons = grid_shape[dim] if circular_cv else window_size + 1 + step_size = max(1, num_horizons // 30) + + Z_windows = _windowed_take(Z_grid, dim, window_size, + num_horizons, step_size, circular_cv) + w_windows = _windowed_take(sw_grid, dim, window_size, + num_horizons, step_size, circular_cv) + + Z_windows = np.moveaxis(Z_windows, dim, 0) + w_windows = np.moveaxis(w_windows, dim, 0) + Z_windows = np.moveaxis(Z_windows, -2, -1) + + batch_size = Z_windows.shape[0] + Z_batch = Z_windows.reshape(batch_size, -1, n_features_aug) + weights_batch = w_windows.reshape(batch_size, -1, 1) + + ZTW = Z_batch.transpose(0, 2, 1) * weights_batch.transpose(0, 2, 1) + XTWX_super = ZTW @ Z_batch + + diag = np.diagonal(XTWX_super, axis1=1, axis2=2) + scales_super = np.sqrt(np.maximum(np.abs(diag), 1e-30)) + + per_dim_super.append((XTWX_super, scales_super)) + + return { + 'per_dim_super': per_dim_super, + 'n_features_aug': n_features_aug, + 'grid_shape': grid_shape, + 'n_terms': n_terms, + # Cached so downstream VWSRSparsity can derive per-target + # ``target`` / ``features`` by slicing instead of re-calling + # objective.evaluate(normalize=True) -- which would force + # another vstack + transpose of the same term evaluations + # for every candidate target_idx in the sweep. + 'Z': Z, + } + + @classmethod + def from_full(cls, super_data, target_idx_in_terms): + """Construct a per-target GramSetup view from precomputed + super-Gram data via slicing. + + ``super_data`` is the dict returned by :meth:`precompute_super`. + ``target_idx_in_terms`` is the column index within Z (the terms + portion) that the caller wants as the regression target; the + intercept column stays in the feature set automatically. + + The returned object has the same ``n_features_aug``, ``grid_shape`` + and ``_per_dim`` shape contract as a regular ``GramSetup`` so + downstream code (``PhysicsInformedLasso.fit`` -> ``solve``) is + unchanged. + """ + per_dim_super = super_data['per_dim_super'] + n_features_aug_super = super_data['n_features_aug'] + grid_shape = super_data['grid_shape'] + + if not (0 <= target_idx_in_terms < n_features_aug_super - 1): + raise IndexError( + f'target_idx_in_terms={target_idx_in_terms} out of range ' + f'[0, {n_features_aug_super - 1}) for super-Gram with ' + f'{n_features_aug_super - 1} terms (+1 intercept).' + ) + + active = np.ones(n_features_aug_super, dtype=bool) + active[target_idx_in_terms] = False + + instance = cls.__new__(cls) + instance.n_features_aug = int(active.sum()) # n_terms (incl. intercept) + instance.grid_shape = grid_shape + instance._per_dim = [] + for XTWX_super, scales_super in per_dim_super: + XTWX_target = XTWX_super[:, active, :][:, :, active] + # XTWy = (Z_aug[:, ~t])^T W Z[:, t] = column t of XTWX_super + # at rows in ``active``. + XTWy_target = XTWX_super[:, active, + target_idx_in_terms:target_idx_in_terms + 1] + scales_target = scales_super[:, active] + instance._per_dim.append((XTWX_target, XTWy_target, scales_target)) + return instance + + +def taylor_microscale(field: np.ndarray, grid_shape, axis: int, + eps: float = 1e-30, + deriv_field: np.ndarray = None) -> float: + """Return the Taylor microscale of ``field`` along ``axis``. + + ``lambda^2 = / <(du/dx)^2>``, with unit grid spacing -- the + result has units of grid points so it sets the locality scale + directly. Constant or near-constant fields produce ``inf``. + + When ``deriv_field`` is supplied (the already-computed + ``d(field)/d(axis)``), use it directly instead of ``np.gradient``; + cache fast path that also gives a higher-quality estimate on noisy + data when the EPDE pool stored ANN-smoothed derivatives. + """ + f = np.asarray(field).reshape(grid_shape) + if deriv_field is None: + g = np.gradient(f, axis=axis) + else: + g = np.asarray(deriv_field).reshape(grid_shape) + num = float(np.mean(f * f)) + den = float(np.mean(g * g)) + if den <= eps or not np.isfinite(den): + return float('inf') + return float(np.sqrt(num / den)) + + +def _cached_primary_var_names() -> list: + """List bare variable labels currently in ``global_var.tensor_cache``. + + Bare labels are entries with no ``/`` or ``^`` -- ``u``, ``v``, ``p``, + etc. Derivative labels (``du/dx0``, ``d^2u/dx1^2``) are filtered out. + Returns ``[]`` when the cache is uninitialised or empty. + """ + import epde.globals as _gv + tc = getattr(_gv, 'tensor_cache', None) + if tc is None: + return [] + try: + numpy_dict = tc.memory_default.get('numpy', {}) + except Exception: + return [] + out = [] + for k in numpy_dict: + if not isinstance(k, tuple) or len(k) != 2: + continue + label = k[0] + if (isinstance(label, str) and '/' not in label + and '^' not in label and label not in out): + out.append(label) + return out + + +def resolve_vc_modes_from_input(grid_shape, main_var=None, + k_max: int = 6, k_min: int = 2): + """One-shot resolve of per-axis varying-coefficient basis modes ``K_d`` + from the cached source variable's Taylor microscale, cached per + ``(grid_shape, main_var)`` so every candidate equation in one run shares + the SAME basis resolution. + + ``K_d = clip(ceil(n_d / lambda_d) + 1, k_min, k_max)``: roughly one + cosine mode per coherence length along axis ``d``, ``+1`` for the + constant term. The microscale is taken as the *minimum* over variables + (fastest-varying wins) so the basis can resolve every variable's + structure; constant / near-constant fields (``lambda_d -> inf``) + collapse to ``k_min`` (constant + one mode). + + Returns a tuple ``(K_0, ..., K_{D-1})``; ``None`` on cache miss with no + cached variable (caller falls back to the target-field path). + """ + import epde.globals as _gv + key = (tuple(int(n) for n in grid_shape), main_var, 'vc') + cache = getattr(_gv, 'vc_modes_cache', None) + if cache is None: + return None + if key in cache: + return cache[key] + + if main_var is not None: + var_names = [main_var] + else: + var_names = _cached_primary_var_names() + if not var_names: + return None + + tc = _gv.tensor_cache + D = len(grid_shape) + lam_min = [float('inf')] * D + for var in var_names: + try: + field = tc.get((var, (1.0,))) + except Exception: + continue + for d in range(D): + lam = taylor_microscale(field, grid_shape, d) + if lam < lam_min[d]: + lam_min[d] = lam + if all(not np.isfinite(l) for l in lam_min): + return None + + modes = [] + for d in range(D): + n_d = int(grid_shape[d]) + lam = lam_min[d] + if not np.isfinite(lam) or lam <= 0: + K_d = k_min + else: + K_d = int(np.ceil(n_d / lam)) + 1 + K_d = max(k_min, min(int(K_d), int(k_max))) + modes.append(int(K_d)) + result = tuple(modes) + cache[key] = result + return result + + +class VaryingCoefSetup: + """Varying-coefficient stability estimator -- the default + ``gram_mode='vcoef'`` path, alternative to the axis-aligned + ``GramSetup`` (the ``gram_mode='axis'`` backup). + + Instead of measuring each coefficient's dispersion over local + sub-regions, model each term's coefficient as a smooth function of + position ``beta_j(x) = gamma_{j,0} + sum_d sum_{k>=1} gamma_{j,d,k} + B_k(x_d)`` (additive low-frequency cosine basis, so the column count is + LINEAR in the spatial dimension). A term is homogeneous (true) iff its + *non-constant* energy is small relative to its constant part; the + per-term stability score is + + score_j = (Var(gamma_{j,0}) + NC_j) / (gamma_{j,0}^2 + 1e-30) + + with ``NC_j = sum_{k>=1} max(gamma_{j,d,k}^2 - lam*Var(gamma_{j,d,k}), 0)`` + the noise-debiased non-constant energy: significance of the constant + coefficient plus region-variation, over the squared constant part. The + locality scale is the basis resolution ``K`` -- resolved once per + ``(grid_shape, main_var)`` from the Taylor microscale and shared by every + candidate -- so there is no locality hyperparameter to tune per dataset. + + Contract: exposes ``score(active_mask)`` returning a length-active vector + aligned to the active feature columns, a drop-in for + ``PhysicsInformedLasso.get_cv``'s return. The ``precompute_super`` / + ``from_full`` pair mirrors ``GramSetup`` so the EqRPS term-sweep builds + the expanded Gram ONCE and slices per candidate target. + + Noise handling: a frequency-scaled ridge ``rho*k^2`` on the gamma solve + (smoothing-spline analogue -- high modes shrunk hard) plus a chi-square + bias subtraction of the expected noise energy from ``NC_j``, both scaled + by the homogeneous-fit residual variance, so a true term's *expected* + score is ~0 rather than a positive noise floor. + """ + + is_vcoef = True + + def __init__(self, X, y, sample_weights, grid_shape, main_var: str = None, + modes=None, k_max=None, freq_coef=None, eps_rel: float = 1e-6, + fit_intercept: bool = True): + grid_shape = tuple(int(n) for n in grid_shape) + X = np.asarray(X, dtype=float) + if X.ndim == 1: + X = X[:, None] + n_samples = X.shape[0] + # Mirror calculate_weights' intercept policy: only append the constant + # (ones) column when the equation actually fits an intercept + # (weights_final[-1] != 0). A PDE has no constant term, so its intercept + # is zeroed; including it would fit gamma_0 ~ 0 and blow up the + # 1/t^2 score, dominating the summed stability objective. + X_aug = (np.hstack([X, np.ones((n_samples, 1))]) if fit_intercept + else X) + modes = self._resolve_modes(modes, grid_shape, main_var, k_max) + w = np.asarray(sample_weights, dtype=float).reshape(-1) + Bvals, mode_k = self._basis_values(grid_shape, modes) + G, Phiy, _, B = self._gram(X_aug, w, Bvals, mode_k, y=y) + + self.G = G + self.Phiy = Phiy + self.B = int(B) + self.basis_mode_k = mode_k + # Basis values (N, B) kept for beta(x) reconstruction in + # ``beta_field_stats`` (diagnostic report-form comparison). Only the + # direct-construction path stores it; ``from_full`` leaves it None. + self._Bvals = Bvals + self.n_features = X_aug.shape[1] + self.grid_shape = grid_shape + self.N_eff = float(np.sum(w)) + y_flat = np.asarray(y, dtype=float).reshape(-1) + self.yWy = float(np.dot(y_flat, w * y_flat)) + self.freq_coef = self._cfg_freq(freq_coef) + self.eps_rel = float(eps_rel) + + # ------------------------------------------------------------------ # + # Config / basis helpers + # ------------------------------------------------------------------ # + @staticmethod + def _cfg_freq(freq_coef): + if freq_coef is not None: + return float(freq_coef) + import epde.globals as _gv + return float(getattr(_gv, 'vc_freq_coef', 1.0)) + + @classmethod + def _resolve_modes(cls, modes, grid_shape, main_var, k_max): + if modes is not None: + return tuple(int(m) for m in modes) + import epde.globals as _gv + if k_max is None: + k_max = int(getattr(_gv, 'vc_k_max', 6)) + resolved = resolve_vc_modes_from_input( + grid_shape, main_var=main_var, k_max=k_max) + if resolved is not None: + return resolved + # Fallback (no cached field): a modest fixed resolution per axis. + return tuple(min(int(k_max), 3) for _ in grid_shape) + + @staticmethod + def _basis_values(grid_shape, modes): + """Return ``(Bvals (N, B), mode_k (B,))`` -- the additive cosine + design columns over the flattened grid. Column 0 is the constant; + the rest are zero-mean unit-mean-square DCT-II cosines + ``sqrt(2) cos(pi k (i+0.5)/n_d)`` per axis, ``k=1..K_d-1``. + """ + grid_shape = tuple(int(n) for n in grid_shape) + D = len(grid_shape) + N = int(np.prod(grid_shape)) + cols = [np.ones(N, dtype=float)] + mode_k = [0.0] + for d in range(D): + n_d = grid_shape[d] + Kd = int(modes[d]) + idx = np.arange(n_d, dtype=float) + for k in range(1, Kd): + m1d = np.sqrt(2.0) * np.cos(np.pi * k * (idx + 0.5) / n_d) + shp = [1] * D + shp[d] = n_d + full = np.broadcast_to(m1d.reshape(shp), grid_shape) + cols.append(np.ascontiguousarray(full).reshape(-1)) + mode_k.append(float(k)) + Bvals = np.stack(cols, axis=1) + return Bvals, np.asarray(mode_k, dtype=float) + + @staticmethod + def _gram(F_aug, w, Bvals, mode_k, y=None): + """Build the expanded Gram ``G = Phi^T W Phi`` (and ``Phi^T W y`` if + ``y`` given) where ``Phi[:, f*B + b] = F_aug[:, f] * Bvals[:, b]`` + (block-major by feature, so feature ``f`` owns columns + ``[f*B : (f+1)*B]`` and its constant column is ``f*B``). + """ + N, Fc = F_aug.shape + B = Bvals.shape[1] + Phi = (F_aug[:, :, None] * Bvals[:, None, :]).reshape(N, Fc * B) + WPhi = w[:, None] * Phi + G = Phi.T @ WPhi + Phiy = None + if y is not None: + Phiy = Phi.T @ (w * np.asarray(y, dtype=float).reshape(-1)) + return G, Phiy, mode_k, B + + # ------------------------------------------------------------------ # + # Super-Gram precompute / per-target slicing (EqRPS term sweep) + # ------------------------------------------------------------------ # + @classmethod + def precompute_super(cls, Z, sample_weights, grid_shape, main_var=None, + modes=None, k_max=None, freq_coef=None, + eps_rel: float = 1e-6): + grid_shape = tuple(int(n) for n in grid_shape) + n_samples, n_terms = Z.shape + Z_aug = np.hstack([Z, np.ones((n_samples, 1))]) + modes = cls._resolve_modes(modes, grid_shape, main_var, k_max) + w = np.asarray(sample_weights, dtype=float).reshape(-1) + Bvals, mode_k = cls._basis_values(grid_shape, modes) + G_super, _, _, B = cls._gram(Z_aug, w, Bvals, mode_k, y=None) + return { + 'mode': 'vcoef', + 'G_super': G_super, + 'B': int(B), + 'basis_mode_k': mode_k, + 'n_features_aug': n_terms + 1, + 'grid_shape': grid_shape, + 'n_terms': n_terms, + 'N_eff': float(np.sum(w)), + 'freq_coef': cls._cfg_freq(freq_coef), + 'eps_rel': float(eps_rel), + 'Z': Z, + } + + @classmethod + def from_full(cls, super_data, target_idx_in_terms): + """Per-target view: features = all terms except ``target`` plus the + intercept; ``Phi^T W y`` is the target's constant column of the + super-Gram (since the target equals its own constant-expansion + column). Mirrors ``GramSetup.from_full``. + """ + G_super = super_data['G_super'] + B = int(super_data['B']) + n_feat_super = int(super_data['n_features_aug']) # n_terms + 1 + n_terms = int(super_data['n_terms']) + if not (0 <= target_idx_in_terms < n_terms): + raise IndexError( + f'target_idx_in_terms={target_idx_in_terms} out of range ' + f'[0, {n_terms}) for vcoef super-Gram.') + + active_global = [i for i in range(n_feat_super) + if i != target_idx_in_terms] + cols = np.concatenate( + [np.arange(i * B, (i + 1) * B) for i in active_global]) + target_const = target_idx_in_terms * B + + inst = cls.__new__(cls) + inst.is_vcoef = True + inst.G = G_super[np.ix_(cols, cols)] + inst.Phiy = G_super[cols, target_const].copy() + inst.yWy = float(G_super[target_const, target_const]) + inst.B = B + inst.basis_mode_k = super_data['basis_mode_k'] + inst.n_features = len(active_global) + inst.grid_shape = super_data['grid_shape'] + inst.N_eff = float(super_data['N_eff']) + inst.freq_coef = float(super_data['freq_coef']) + inst.eps_rel = float(super_data['eps_rel']) + inst._Bvals = None # super path: beta(x) reconstruction unavailable + return inst + + # ------------------------------------------------------------------ # + # Per-term stability score + # ------------------------------------------------------------------ # + def _solve_gammas(self, active_mask=None): + """Frisch-Waugh block solve for the per-active-feature gammas: the + constant part ``gamma_0`` is the features-only weighted OLS coefficient + (decoupled from the basis modes, so collinear systems recover it + correctly), and the modes are fit to the constant-fit residual with a + noise-adaptive frequency ridge. Returns a dict with ``gamma`` (basis + coefficients, block-major by feature), ``var`` (Cov diagonal), and the + layout (``nf``, ``B``, ``mk``, ``mean_power``). ``None`` if no active + feature. + """ + n = self.n_features + if active_mask is None: + active_mask = np.ones(n, dtype=bool) + active_feats = np.where(active_mask)[0] + nf = int(active_feats.size) + if nf == 0: + return None + + B = self.B + mk = self.basis_mode_k + cols = (active_feats[:, None] * B + + np.arange(B)[None, :]).reshape(-1) + G_a = self.G[np.ix_(cols, cols)].astype(float).copy() + Phiy_a = self.Phiy[cols].astype(float).copy() + col_mode = np.tile(mk, nf) + const_local = np.arange(nf) * B + is_mode = np.ones(nf * B, dtype=bool) + is_mode[const_local] = False + mode_local = np.where(is_mode)[0] + + # --- Constant block: features-only weighted OLS (the standard SINDy + # coefficient gamma_0). Decoupling it from the basis modes (Frisch- + # Waugh: fit constants, then fit modes to the constant-fit residual) is + # what keeps the recovered coefficient correct on collinear systems. The + # joint expanded solve let the near-collinear modulated columns steal + # signal from the constant (e.g. lorenz g0=2.3 instead of 10, despite a + # well-conditioned const block) -> spurious non-constant energy and a + # wrong magnitude. Here g0 = exactly the features-only OLS, untouched by + # the modes. + Ac = G_a[np.ix_(const_local, const_local)] + bc = Phiy_a[const_local] + dC = np.sqrt(np.maximum(np.diag(Ac), 1e-30)) + AcN = Ac / np.outer(dC, dC) + AcN[np.diag_indices_from(AcN)] += 1e-10 + try: + AcN_inv = np.linalg.inv(AcN) + except np.linalg.LinAlgError: + AcN_inv = np.linalg.pinv(AcN) + g0_vec = (AcN_inv @ (bc / dC)) / dC + + rss = self.yWy - float(g0_vec @ bc) + sigma2 = max(rss, 0.0) / max(self.N_eff - nf, 1.0) + mean_power = self.yWy / max(self.N_eff, 1.0) + noise_rel = min(max(sigma2 / (mean_power + 1e-30), 0.0), 1.0) + # OLS coefficient variances Var(gamma_0) = sigma^2 (X^T W X)^{-1}. + var0_vec = sigma2 * (np.diag(AcN_inv) / (dC * dC)) + + gamma = np.zeros(nf * B) + var = np.zeros(nf * B) + gamma[const_local] = g0_vec + var[const_local] = var0_vec + + # --- Mode block: fit the constant-fit residual r = y - X gamma_0 onto + # the basis-modulated columns (r is the OLS residual, so it is already + # W-orthogonal to the constant columns -> gamma_0 stays untouched). The + # noise-adaptive frequency ridge rho*k^2 still shrinks high modes; on + # clean data the modes simply stay ~0 (a true constant coefficient). + if mode_local.size: + Amm = G_a[np.ix_(mode_local, mode_local)] + Amc = G_a[np.ix_(mode_local, const_local)] + b_m = Phiy_a[mode_local] - Amc @ g0_vec + dM = np.sqrt(np.maximum(np.diag(Amm), 1e-30)) + AmmN = Amm / np.outer(dM, dM) + kmax2 = max(float(np.max(mk)) ** 2, 1.0) + cm = col_mode[mode_local] + ridge_m = 1e-6 + self.freq_coef * noise_rel * (cm ** 2 / kmax2) + AmmN[np.diag_indices_from(AmmN)] += ridge_m + import epde.globals as _gv + if getattr(_gv, 'vc_mode_decouple', False): + # Block-diagonalise by feature: drop the cross-feature mode + # collinearity blocks so each term's modes are fit only to what + # ITS OWN modulated columns explain of the constant-fit + # residual. A true constant-coef term then keeps modes ~0 (B + # small) even when collinear grid-modulated cousins are present, + # instead of borrowing their variation. Within-feature mode + # correlations (same feature, different k) are preserved. + feat_of_mode = mode_local // B + AmmN = AmmN * (feat_of_mode[:, None] == feat_of_mode[None, :]) + try: + AmmN_inv = np.linalg.inv(AmmN) + except np.linalg.LinAlgError: + AmmN_inv = np.linalg.pinv(AmmN) + gamma[mode_local] = (AmmN_inv @ (b_m / dM)) / dM + var[mode_local] = sigma2 * (np.diag(AmmN_inv) / (dM * dM)) + + return {'gamma': gamma, 'var': var, 'active_feats': active_feats, + 'nf': nf, 'B': B, 'mk': mk, 'mean_power': mean_power} + + def beta_field_stats(self, active_mask=None): + """Per-active-feature stats of the reconstructed coefficient field + ``beta_j(x) = sum_b gamma_{j,b} B_b(x)`` over the grid -- ``mu``, ``std``, + ``median``, ``mad`` (length nf). Lets callers compare report-form + variants (std/mu, mad/median, squared) on the SAME gammas. Requires the + basis values (direct-construction path only); ``None`` otherwise. + """ + Bvals = getattr(self, '_Bvals', None) + sol = self._solve_gammas(active_mask) + if sol is None or Bvals is None: + return None + gamma, nf, B = sol['gamma'], sol['nf'], sol['B'] + mu = np.empty(nf); sd = np.empty(nf) + med = np.empty(nf); mad = np.empty(nf) + for i in range(nf): + beta = Bvals @ gamma[i * B:(i + 1) * B] # (N,) coefficient field + mu[i] = float(np.mean(beta)) + sd[i] = float(np.std(beta)) + m = float(np.median(beta)) + med[i] = m + mad[i] = float(np.median(np.abs(beta - m))) + return {'mu': mu, 'std': sd, 'median': med, 'mad': mad} + + def score(self, active_mask=None): + """Per-active-feature stability score (the ``report`` form), uncapped: + + score_j = (Var(gamma_0) + NC_deb) / (gamma_0^2 + 1e-30), + NC_deb = sum_{k>=1} max(gamma_{0,k}^2 - lam*Var(gamma_{0,k}), 0). + + Two normalised instabilities over the squared constant coefficient: the + significance term ``Var(gamma_0)/gamma_0^2`` (a noisy / weakly identified + coefficient) plus the noise-DEBIASED non-constant energy + ``NC_deb/gamma_0^2`` of the varying coefficient ``beta_j(x)`` (each + mode's own sampling variance subtracted, so a pure-noise modulation + contributes ~0). A true (homogeneous) term fits with a well-identified + CONSTANT coefficient -> both terms ~0 -> score ~0 (kept); the debiasing + is what saves a genuinely weak true term (Allen-Cahn's 1e-4 diffusion). + A spurious term either VARIES across the region (``gamma_k^2 >> Var`` -> + large ``NC_deb``) or is poorly identified (large ``Var(gamma_0)``) -> + large score. + + Uncapped: the former 1e6 cap only masked a ``gamma_0 ~ 0`` blow-up from a + zeroed intercept appended as a feature column. The Frisch-Waugh solve + keeps ``gamma_0`` the clean features-only OLS coefficient and the + intercept-exclusion fit (``fit_intercept = weights[-1] != 0``) removes + that source, so a surviving real term's ``gamma_0`` stays away from 0 and + the ratio stays finite without a cap. + + Aligned to the active feature columns: a drop-in for the + ``get_cv`` return (``active_thresholds = score * max_corr`` in + ``PhysicsInformedLasso.fit``); the equation-level stability objective is + the sum over the non-zero real terms (see ``vc_stability_total_lr``). + """ + sol = self._solve_gammas(active_mask) + if sol is None: + return np.zeros(0) + gamma, var = sol['gamma'], sol['var'] + nf, B, mk = sol['nf'], sol['B'], sol['mk'] + import epde.globals as _gv + lam = float(getattr(_gv, 'vc_debias_lambda', 1.0)) + is_const = mk == 0 + nonconst = ~is_const + scores = np.empty(nf) + for i in range(nf): + sl = slice(i * B, (i + 1) * B) + g = gamma[sl] + v = var[sl] + g0 = float(g[is_const][0]) if np.any(is_const) else 0.0 + var0 = float(v[is_const][0]) if np.any(is_const) else 0.0 + C = g0 ** 2 + nc_g2 = g[nonconst] ** 2 + nc_deb = float(np.sum(np.maximum(nc_g2 - lam * v[nonconst], 0.0))) + # Significance + debiased region-variation over the squared constant + # coefficient (floored by 1e-30), uncapped. + scores[i] = (var0 + nc_deb) / (C + 1e-30) + return np.nan_to_num(scores) + + +def vc_stability_total_lr(features, target, sample_weights, grid_shape, + main_var: str = None, fit_intercept: bool = True): + """Equation-level varying-coefficient stability: the SUM over the equation's + terms of the per-term ``score`` ``(Var(gamma_0) + NC_deb)/gamma_0^2`` (the + ``gram_mode='vcoef'`` replacement for the inline ``total_lr``). + Lower = more stable. Pass non-zero-term features (``evaluate(normalize= + False)``) and ``fit_intercept = weights_final[-1] != 0`` so neither + zero-weight terms nor a zeroed intercept (a ~0 coefficient that would blow + up the 1/gamma_0^2 ratio) enter the sum. + """ + X = np.asarray(features) + if X.ndim == 1: + X = X[:, None] + setup = VaryingCoefSetup(X, target, sample_weights, grid_shape, + main_var=main_var, fit_intercept=fit_intercept) + return float(np.sum(setup.score(None))) + + +def calculate_weights(X, y, sample_weights, grid_shape, fit_intercept=True, + gram_cls=None, gram_kwargs=None): + """ + Vectorized calculation of weights across sliding windows. + Dynamically handles whether the intercept should be fit. + + Single-shot wrapper over :class:`GramSetup`: builds the precomputed + Gram once and immediately solves with the requested intercept policy. + Callers that solve the same Gram against many active masks (e.g. + :class:`PhysicsInformedLasso.fit`) should instantiate ``GramSetup`` + directly and call ``.solve(active_mask)`` per iteration to avoid + re-running the expensive ``X^T diag(w) X`` matmul. + + ``gram_cls`` selects the construction strategy: default ``None`` -> + ``GramSetup`` (axis-aligned sliding windows, the backup path). + ``gram_kwargs`` is forwarded to the chosen class's constructor. The + varying-coefficient default (``gram_mode='vcoef'``) does not route + through here -- ``PhysicsInformedLasso.fit`` instantiates + ``VaryingCoefSetup`` directly. + """ + if gram_cls is None: + gram_cls = GramSetup + gram_kwargs = gram_kwargs or {} + setup = gram_cls(X, y, sample_weights, grid_shape, **gram_kwargs) + active_mask = np.ones(setup.n_features_aug, dtype=bool) + if not fit_intercept: + # GramSetup always augments with the intercept column; drop it + # from the active set to mimic the legacy ``fit_intercept=False`` + # branch (which never augmented in the first place). + active_mask[-1] = False + return setup.solve(active_mask) diff --git a/epde/operators/multiobjective/moeadd_specific.py b/epde/operators/multiobjective/moeadd_specific.py index 65bd3f68..c6c91276 100644 --- a/epde/operators/multiobjective/moeadd_specific.py +++ b/epde/operators/multiobjective/moeadd_specific.py @@ -8,6 +8,7 @@ import copy import numpy as np import time +import warnings from typing import Union, Tuple from functools import reduce, partial @@ -502,24 +503,47 @@ def apply(self, objective : ParetoLevels, arguments : dict): attempts, uniqueness_attempt_limit, ) if hit_cap: - # Initial population must be duplicate-free for MOEA/D's - # per-sector uniqueness invariant. If the candidate pool - # is too small to satisfy pop_size, fail loud rather than - # silently corrupt the initial Pareto layer. - raise RuntimeError( - f"InitialParetoLevelSorting: could not generate a unique " - f"initial candidate after {uniqueness_attempt_limit} attempts " - f"(candidate index {idx}, {len(objective.history)} unique " - f"systems already placed). The search space appears smaller " - f"than the requested population. Reduce pop_size, widen the " - f"token pool, or raise InitialParetoLevelSorting's " - f"'uniqueness_attempt_limit' parameter." + # Search-space collapse: the RPS+simplify pipeline canonicalises + # ``uniqueness_attempt_limit`` consecutive create() outputs into + # one of the already-placed structures. Diagnostic builds want + # to observe the partial population's objective values rather + # than abort. Policy: + # warn, register the placed candidates in ``population`` / + # ``levels[0]`` so downstream consumers can read their + # ``text_form`` + ``obj_fun``, set ``_init_collapsed=True``, + # and stop -- MOEA/D's ``optimize`` checks this flag after + # init and skips the epoch loop so mutation/crossover never + # runs on the degenerate population. + warnings.warn( + f"InitialParetoLevelSorting: search-space collapse at " + f"candidate {idx} after {uniqueness_attempt_limit} attempts " + f"({len(objective.history)} unique systems placed of " + f"{len(objective.unplaced_candidates)} requested). Stopping " + f"the search; placed candidates are accessible via " + f"pareto_levels.levels[0]." ) + placed = list(objective.unplaced_candidates[:idx]) + objective.population = placed + objective.unplaced_candidates = [] + # Single-level dump: Pareto-correctness is moot here because + # we're terminating. Consumers reading ``levels[0]`` get + # every placed candidate. + objective.levels = [placed] if placed else [[]] + objective._init_collapsed = True + init_collapsed = True + break self.suboperators['chromosome_fitness'].apply(objective=candidate, arguments=subop_args['chromosome_fitness']) objective.history.add(system) if global_var.verbose.candidate_objectives: print(candidate.obj_fun) + else: + init_collapsed = False + if init_collapsed: + if global_var.verbose.show_iter_idx: + print(f'\n*** Search-space collapse: stopping early with ' + f'{len(objective.levels[0])} placed candidates. ***') + return if global_var.verbose.show_iter_idx: print('\n========== Marriage (weight assignment) ==========') objective.associate_weights() diff --git a/epde/operators/multiobjective/mutations.py b/epde/operators/multiobjective/mutations.py index c89b0b0a..b5019c70 100644 --- a/epde/operators/multiobjective/mutations.py +++ b/epde/operators/multiobjective/mutations.py @@ -26,6 +26,7 @@ class SystemMutation(CompoundOperator): key = 'SystemMutation' + @_loop_stats.timed('SystemMutation.apply') def apply(self, objective : SoEq, arguments : dict): # TODO: add setter for best_individuals & worst individuals self_args, subop_args = self.parse_suboperator_args(arguments = arguments) @@ -69,6 +70,7 @@ def use_default_tags(self): class EquationMutation(CompoundOperator): key = 'EquationMutation' + @_loop_stats.timed('EquationMutation.apply') @HistoryExtender(f'\n -> mutating equation', 'ba') def apply(self, objective : Equation, arguments : dict): self_args, subop_args = self.parse_suboperator_args(arguments = arguments) @@ -91,7 +93,11 @@ def apply(self, objective : Equation, arguments : dict): r_mutation = self.params['r_mutation'] replace_attempts = 0 mutable_count = max(1, len(equation.structure) - equation.n_immutable) - for term_idx in range(equation.n_immutable, len(equation.structure)): + # Reverse order so a full-drop dedup inside TermMutation (which can + # remove ``term`` and shrink ``structure``) only shifts indices we + # have already processed -- forward iteration would skip terms or + # run off the end of the now-shorter structure. + for term_idx in reversed(range(equation.n_immutable, len(equation.structure))): if np.random.uniform(0, 1) <= r_mutation: replace_attempts += 1 self.suboperators['mutation'].apply( @@ -175,45 +181,34 @@ def apply(self, objective : tuple, arguments : dict): #term_idx, equation): # optimization. Saves a ~10ms Term deepcopy on every replace_terms # iter (was the largest remaining mutation-path deepcopy cost). original_structure = term.structure - original_labels = term.factors_labels term.randomize() term.reset_saved_state() equation._invalidate_label_cache() - # Re-randomize while the mutation produced a duplicate term within - # the equation OR no actual change vs the previous term. Cap the - # retries so a tight token pool can't deadlock the optimizer (same - # hazard fixed in ``enforce_rps_uniqueness`` / ``simplify_equation``). - # On cap-hit, revert to the pre-mutation term: silently committing a - # duplicate or no-op violates the structure-dedup rule and lets - # population diversity drift unobservably. - max_iter = 100 - attempts = 0 - hit_cap = True - for _ in range(max_iter): - attempts += 1 - signatures = {t.factors_labels for t in equation.structure} - duplicate = len(signatures) != len(equation.structure) - unchanged = term.factors_labels == original_labels - if not (duplicate or unchanged): - hit_cap = False - break - term.randomize() - term.reset_saved_state() + # Full-drop dedup (no regeneration, no retry loop): if the freshly + # mutated term now duplicates another term in the equation, remove + # it outright. Only at the 2-term floor (nothing safe to drop) do we + # restore the unique pre-mutation term. This is the mutation-path + # policy -- distinct from simplify's regenerate-n-then-drop -- and + # keeps duplicates out of the population (otherwise caught a + # generation later at the RPS entry / crossover assert). + signatures = {t.factors_labels for t in equation.structure} + duplicate = len(signatures) != len(equation.structure) + if duplicate and len(equation.structure) > 2: + tgt = getattr(equation, 'target_idx', None) + equation.structure = [t for t in equation.structure if t is not term] + if tgt is not None and term_idx < tgt: + equation.target_idx -= 1 equation._invalidate_label_cache() - if hit_cap: - # Revert by restoring the original Factor list. The Term - # instance itself is the same one, and only its ``structure`` - # slot was mutated by randomize() -- restoring that slot - # restores its observable identity for downstream consumers - # (factors_labels, __eq__ both read ``structure`` live). + _loop_stats.record('TermMutation.unique_term.DROP', 1, 1) + elif duplicate: + # Floor: restore the (unique) pre-mutation term in place. term.structure = original_structure term.reset_saved_state() equation._invalidate_label_cache() - _loop_stats.record( - 'TermMutation.unique_term' + ('.FAIL' if hit_cap else ''), - attempts, max_iter, - ) + _loop_stats.record('TermMutation.unique_term.FLOOR_REVERT', 1, 1) + else: + _loop_stats.record('TermMutation.unique_term', 1, 1) return term def use_default_tags(self): diff --git a/epde/operators/multiobjective/variation.py b/epde/operators/multiobjective/variation.py index 479d5633..a37495ef 100644 --- a/epde/operators/multiobjective/variation.py +++ b/epde/operators/multiobjective/variation.py @@ -46,7 +46,8 @@ class ParetoLevelsCrossover(CompoundOperator): copy_properties_to """ key = 'ParetoLevelsCrossover' - + + @_loop_stats.timed('ParetoLevelsCrossover.apply') def apply(self, objective : ParetoLevels, arguments : dict): """ Method to obtain a new population by selection of parent individuals (equations) and performing a crossover between them to get the offsprings. @@ -167,6 +168,7 @@ def use_default_tags(self): class EquationCrossover(CompoundOperator): key = 'EquationCrossover' + @_loop_stats.timed('EquationCrossover.apply') @HistoryExtender(f'\n -> performing equation crossover', 'ba') def apply(self, objective : tuple, arguments : dict): """Hybrid random-partition + parameter-blend crossover. diff --git a/epde/optimizers/moeadd/moeadd.py b/epde/optimizers/moeadd/moeadd.py index 77e1f26b..cf71065a 100644 --- a/epde/optimizers/moeadd/moeadd.py +++ b/epde/optimizers/moeadd/moeadd.py @@ -27,6 +27,8 @@ def flatten_chain(matrix): from epde.optimizers.moeadd.supplementary import fast_non_dominated_sorting, ndl_update, Equality, Inequality, acute_angle from scipy.spatial import ConvexHull +from epde import _loop_stats + def clear_list_of_lists(inp_list) -> list: ''' Delete elements-lists with len(0) from the list @@ -671,23 +673,45 @@ def optimize(self, epochs, early_stopping_callback=None): linked.reset_traversal_cond() linked.initial[0][1].set_output(self.pareto_levels) init_block.apply(self.form_processer_args(0)) - for epoch_idx in np.arange(epochs): + if getattr(self.pareto_levels, '_init_collapsed', False): + # Init flagged a search-space collapse (see + # ``InitialParetoLevelSorting.apply``). Skip the epoch loop -- + # mutation/crossover on a degenerate truncated population would + # only contaminate the results. Snapshot the placed level-0 + # candidates once so downstream consumers + # (``pareto_history`` / ``thesis_runner._extract_*``) still + # see something. + snapshot = [] + for sol in self.pareto_levels.levels[0]: + try: + obj = sol.obj_fun.tolist() if hasattr(sol.obj_fun, 'tolist') else list(sol.obj_fun) + except Exception: + obj = None + snapshot.append({'text_form': sol.text_form, 'obj_fun': obj}) + self._pareto_history.append(snapshot) if global_var.verbose.show_iter_idx: - print(f'\n----- Multiobjective optimization : {epoch_idx + 1}-th epoch -----') - # Shuffle sector order each epoch. The prior fixed - # 0..N-1 traversal gave early sectors a population - # advantage every epoch (their offspring entered the - # global Pareto pool before late sectors saw the - # evolved chromosomes). Reproducibility is preserved - # because np.random is seeded by the caller via - # ``_set_seeds`` before optimization. - n_sectors = len(self.weights) - for order_idx, weight_idx in enumerate(np.random.permutation(n_sectors)): + print(f'\n*** MOEADD.optimize: early stop after init collapse ' + f'with {len(self.pareto_levels.levels[0])} candidates. ***') + return + for epoch_idx in np.arange(epochs): + with _loop_stats.timer('MOEADD.epoch'): if global_var.verbose.show_iter_idx: - print(f'During MO : processing {order_idx + 1}-th sector (of {n_sectors}).') - sp_kwargs = self.form_processer_args(weight_idx) - self.sector_processer.run(population_subset = self.pareto_levels, - EA_kwargs = sp_kwargs) + print(f'\n----- Multiobjective optimization : {epoch_idx + 1}-th epoch -----') + # Shuffle sector order each epoch. The prior fixed + # 0..N-1 traversal gave early sectors a population + # advantage every epoch (their offspring entered the + # global Pareto pool before late sectors saw the + # evolved chromosomes). Reproducibility is preserved + # because np.random is seeded by the caller via + # ``_set_seeds`` before optimization. + n_sectors = len(self.weights) + for order_idx, weight_idx in enumerate(np.random.permutation(n_sectors)): + if global_var.verbose.show_iter_idx: + print(f'During MO : processing {order_idx + 1}-th sector (of {n_sectors}).') + sp_kwargs = self.form_processer_args(weight_idx) + with _loop_stats.timer('MOEADD.sector'): + self.sector_processer.run(population_subset = self.pareto_levels, + EA_kwargs = sp_kwargs) stats = self.pareto_levels.get_stats() self._hist.append(stats) # Snapshot the current Pareto-0 structures so consumers can diff --git a/epde/structure/main_structures.py b/epde/structure/main_structures.py index b4c40d80..fead648f 100644 --- a/epde/structure/main_structures.py +++ b/epde/structure/main_structures.py @@ -223,6 +223,7 @@ def descr_variable_marker(self, marker: False): else: raise ValueError('Described variable marker shall be a family label (i.e. "u") of "False"') + @_loop_stats.timed('Term.evaluate') def evaluate(self, structural, grids=None): assert global_var.tensor_cache is not None, 'Currently working only with connected cache' normalize = structural @@ -407,7 +408,7 @@ class Equation(ComplexStructure): 'target_idx', 'right_part_selected', '_weights_final', 'weights_final_evald', 'simplified', 'is_correct_right_part', '_weights_internal', 'weights_internal_evald', 'fitness_calculated', 'stability_calculated', 'aic_calculated', 'solver_form_defined', '_fitness_value', '_coefficients_stability', '_aic', 'metaparameters', 'main_var_to_explain', - '_eval_cache', '_cached_sw_weights', + '_eval_cache', '_cached_sw_weights', '_cached_vc_score', '_terms_labels_cache', '_terms_labels_without_power_cache', '_gram_super'] # , '_solver_form' @@ -568,6 +569,7 @@ def __eq__(self, other): return (all([any([other_elem == self_elem for other_elem in other.structure]) for self_elem in self.structure]) and all([any([other_elem == self_elem for self_elem in self.structure]) for other_elem in other.structure]) and len(other.structure) == len(self.structure) + and len(self.weights_final) == len(other.weights_final) and np.all(np.isclose(self.weights_final, other.weights_final))) else: return (all([any([other_elem == self_elem for other_elem in other.structure]) for self_elem in self.structure]) @@ -701,6 +703,7 @@ def reconstruct_by_right_part(self, right_part_idx): new_eq.reset_saved_state() return new_eq + @_loop_stats.timed('Equation.evaluate') def evaluate(self, normalize: bool = True, return_val: bool = False, grids: list = None) -> Tuple: """Evaluate the equation and return (value, target, features). @@ -804,6 +807,9 @@ def reset_state(self, reset_right_part: bool = True) -> None: self._eval_cache = {} # consumed by epde.operators.common.fitness.L2LRFitness; resets here. self._cached_sw_weights = None + # vcoef analogue of _cached_sw_weights: per-term stability scores from + # the sparsity gram_setup, summed as the stability objective in fitness. + self._cached_vc_score = None self._terms_labels_cache = None self._terms_labels_without_power_cache = None # Tier 3 super-Gram cache (set by EqRightPartSelector for the @@ -845,7 +851,7 @@ def __deepcopy__(self, memo=None): new_struct = _deepcopy_slots( self, memo, attrs_to_avoid_copy=( - '_cached_sw_weights', + '_cached_sw_weights', '_cached_vc_score', '_terms_labels_cache', '_terms_labels_without_power_cache', '_gram_super', ), @@ -878,7 +884,7 @@ def clone_shell(self): self, memo={}, attrs_to_avoid_copy=( 'structure', - '_cached_sw_weights', + '_cached_sw_weights', '_cached_vc_score', '_terms_labels_cache', '_terms_labels_without_power_cache', '_gram_super', ), @@ -1026,7 +1032,7 @@ def text_form(self): if term_idx != self.target_idx: form += str(self.weights_final[term_idx]) if term_idx < self.target_idx else str(self.weights_final[term_idx-1]) form += ' * ' + self.structure[term_idx].name + ' + ' - form += str(self.weights_final[-1]) + ' = ' + \ + form += str(self.weights_internal[-1]) + ' = ' + \ self.structure[self.target_idx].name else: for term_idx in range(len(self.structure)): @@ -1050,7 +1056,7 @@ def latex_form(self): exp_str = r'\cdot 10^{{{0}}} '.format(str(exp)) if exp != 0 else '' form += str(mnt) + exp_str + term.latex_form + r' + ' - mnt, exp = exp_form(self.weights_final[-1], digits_rounding_max) + mnt, exp = exp_form(self.weights_internal[-1], digits_rounding_max) exp_str = r'\cdot 10^{{{0}}} '.format(str(exp)) if exp != 0 else '' form += str(mnt) + exp_str diff --git a/epde/supplementary.py b/epde/supplementary.py index 10008cab..21ebd316 100644 --- a/epde/supplementary.py +++ b/epde/supplementary.py @@ -12,6 +12,7 @@ import numpy as np from functools import reduce import copy +import re import torch # device = torch.device('cpu') @@ -20,40 +21,10 @@ from epde.solver.data import Domain from epde.solver.models import Fourier_embedding, mat_model from epde.preprocessing.smoothers import NN -from numpy.lib.stride_tricks import sliding_window_view from epde import _loop_stats -# Default behaviour for GramSetup's sliding-window CV: ``True`` treats -# each axis as periodic so windows near the boundary wrap around (the -# last few windows span data[end-k:end] + data[0:window_size-k]). -# ``False`` is the legacy linear semantics where the number of windows -# along an axis is ``N - window_size + 1`` and no wrap occurs. -_DEFAULT_CIRCULAR_CV = True - - -def _windowed_take(arr: np.ndarray, dim: int, window_size: int, - num_horizons: int, step_size: int, - circular: bool) -> np.ndarray: - """Return ``sliding_window_view`` along ``dim``, optionally with - circular padding so windows near the boundary wrap to the start. - - Caller computes ``num_horizons`` (the number of valid start positions - along ``dim``) and ``step_size`` (subsampling stride). Under circular - mode, ``num_horizons == arr.shape[dim]`` and the input is padded by - ``window_size - 1`` samples copied from the start at the end via - ``np.pad(..., mode='wrap')``. Under linear mode, ``num_horizons == - arr.shape[dim] - window_size + 1`` and no padding is applied. - """ - if circular: - pad = [(0, 0)] * arr.ndim - pad[dim] = (0, window_size - 1) - arr = np.pad(arr, pad, mode='wrap') - windows = sliding_window_view(arr, window_shape=window_size, axis=dim) - return windows.take(indices=range(0, num_horizons, step_size), axis=dim) - - def retry_until_unique(*, predicate, mutate, max_iter: int, stats_name: str): """Bounded retry loop: keep mutating a candidate until ``predicate`` holds. @@ -475,339 +446,3 @@ def minmax_normalize(matrix): else: matrix[i] = np.zeros_like(matrix[i]) return matrix - - -def _cholesky_solve_batched(A, b): - """Solve ``A @ x = b`` batched over the leading axis using Cholesky. - - ``A`` is assumed symmetric positive-definite (shape ``(batch, n, n)``); - ``b`` is the RHS ``(batch, n, 1)``. Returns ``(x, L)`` where ``x`` is - the solution and ``L`` is the lower-triangular factor (so the caller - can reuse it for iterative refinement). If Cholesky fails on any batch - entry, returns ``(None, None)`` to signal "use the lstsq fallback". - - numpy doesn't ship a batched triangular solver, so the two triangular - solves go through ``np.linalg.solve`` -- still SPD-stable and ~1.5x - faster than feeding the full ``A`` to ``np.linalg.solve``. - """ - try: - L = np.linalg.cholesky(A) - except np.linalg.LinAlgError: - return None, None - try: - z = np.linalg.solve(L, b) - x = np.linalg.solve(L.transpose(0, 2, 1), z) - except np.linalg.LinAlgError: - return None, L - return x, L - - -def _per_batch_lstsq(A, b): - """Per-batch SVD-based least-squares solve. Used as the safety net - when Cholesky reports the equilibrated batch is non-SPD. Returns - weights of shape ``(batch, n, 1)`` matching the input RHS layout so - the caller can compose with subsequent matrix products without - reshaping. - """ - batch_size = A.shape[0] - n = A.shape[1] - out = np.empty((batch_size, n, 1)) - for i in range(batch_size): - sol, *_ = np.linalg.lstsq(A[i], b[i, :, 0], rcond=None) - out[i, :, 0] = sol - return out - - -class GramSetup: - """Precomputed batched normal-equation matrices for fast active-mask - solves. Splits :func:`calculate_weights` into a setup phase (compute - ``X^T diag(w) X`` and ``X^T diag(w) y`` per window-batch per dimension, - using the FULL augmented feature matrix) and a solve phase (slice each - full Gram matrix by an active-feature mask and solve). The setup is - mask-independent; only the solve depends on which columns are active. - - Used by :class:`PhysicsInformedLasso.fit`, whose outer RFE loop calls - ``calculate_weights`` per shrinking column subset. With this split the - expensive ``X^T diag(w) X`` matmul runs ONCE per fit and each outer - iter only pays the cost of an (active × active) solve. The math is - exact: a sub-block of a Gram matrix equals the Gram of the - corresponding sub-columns. - """ - - def __init__(self, X, y, sample_weights, grid_shape, - circular_cv: bool = _DEFAULT_CIRCULAR_CV): - n_samples = X.shape[0] - # Always augment X with the intercept column so callers can toggle - # ``fit_intercept`` via the active mask's last bit rather than - # re-running setup. - X_aug = np.hstack([X, np.ones((n_samples, 1))]) - n_features_aug = X_aug.shape[1] - - X_grid = X_aug.reshape(*grid_shape, n_features_aug) - y_grid = y.reshape(*grid_shape) - sample_weights_grid = sample_weights.reshape(*grid_shape) - - self.n_features_aug = n_features_aug - self.grid_shape = grid_shape - self._per_dim = [] - - for dim in range(len(grid_shape)): - window_size = grid_shape[dim] // 2 - # Circular: every position along the axis is a valid window - # start (the input is virtually periodic); linear: only the - # first ``window_size + 1`` positions yield a full window. - num_horizons = grid_shape[dim] if circular_cv else window_size + 1 - step_size = max(1, num_horizons // 30) - - X_windows = _windowed_take(X_grid, dim, window_size, - num_horizons, step_size, circular_cv) - y_windows = _windowed_take(y_grid, dim, window_size, - num_horizons, step_size, circular_cv) - w_windows = _windowed_take(sample_weights_grid, dim, window_size, - num_horizons, step_size, circular_cv) - - X_windows = np.moveaxis(X_windows, dim, 0) - y_windows = np.moveaxis(y_windows, dim, 0) - w_windows = np.moveaxis(w_windows, dim, 0) - X_windows = np.moveaxis(X_windows, -2, -1) - - batch_size = X_windows.shape[0] - X_batch = X_windows.reshape(batch_size, -1, n_features_aug) - y_batch = y_windows.reshape(batch_size, -1) - weights_batch = w_windows.reshape(batch_size, -1, 1) - - XTW = X_batch.transpose(0, 2, 1) * weights_batch.transpose(0, 2, 1) - XTWX_full = XTW @ X_batch - XTWy_full = XTW @ y_batch[..., None] - - # Per-batch column scales for equilibration in :meth:`solve`. - # ``diag`` is the per-feature L2 norm squared (weighted) of the - # underlying X columns; ``sqrt`` brings it back to a column- - # norm scale. The ``1e-30`` floor is a degenerate-column guard - # (well below any meaningful data scale) so ``1/scale`` stays - # finite for near-zero columns. - diag = np.diagonal(XTWX_full, axis1=1, axis2=2) - scales = np.sqrt(np.maximum(np.abs(diag), 1e-30)) - - self._per_dim.append((XTWX_full, XTWy_full, scales)) - - def solve(self, active_mask=None, ridge_rel=None, ridge_floor=None): - """Solve the normal equations for the active-feature subset across - every window-batch in every spatial dimension. ``active_mask`` is a - length-``n_features_aug`` boolean array; pass ``None`` for the full - set (equivalent to the legacy ``fit_intercept=True`` path). Returns - weights of shape ``(total_windows_across_dims, active_count)``. - - Stability strategy (preserves the Gram-sub-block precompute trick): - - 1. **Column equilibration**: rescale columns by - ``1/sqrt(diag(XTWX))`` so the equilibrated Gram has unit - diagonals and a much smaller effective condition number than - the raw ``XTWX`` (which carries the squared condition number - of the underlying ``sqrt(W) X``). - 2. **Cholesky on the equilibrated SPD batch** (with batched LU - fallback if scipy's batched triangular solve isn't available - on this numpy). Cholesky has tighter backward error than LU - and is ~2x faster on SPD inputs. - 3. **One step of iterative refinement** on the original (un- - equilibrated) system, recovering 6-8 decimal digits that - normal-equation conditioning costs. - 4. **Per-batch lstsq safety net** for any window-batch where - Cholesky fails (non-SPD after equilibration -- rare). - - ``ridge_rel`` / ``ridge_floor`` are kept as no-op kwargs for - backward compatibility with callers from the previous adaptive- - ridge era; the equilibrated solve does not need a per-feature - ridge, only a tiny flat ``1e-10`` on the unit-diagonal matrix. - """ - if active_mask is None: - active_mask = np.ones(self.n_features_aug, dtype=bool) - active_size = int(active_mask.sum()) - - all_weights = [] - for XTWX_full, XTWy_full, scales_full in self._per_dim: - # Two-step boolean slice. Boolean indexing copies, so the - # result is a fresh array we can modify in place without - # corrupting the cached full Gram. - XTWX_a = XTWX_full[:, active_mask, :][:, :, active_mask] - XTWy_a = XTWy_full[:, active_mask, :] - s_a = scales_full[:, active_mask] # (batch, k) - inv_s = 1.0 / s_a # (batch, k) - - # Equilibrate: A = D^-1 XTWX D^-1, b = D^-1 XTWy. After this - # the diagonal of A is 1 by construction; the off-diagonals - # are the correlation coefficients between the underlying - # columns of sqrt(W) X. - A = XTWX_a * inv_s[:, :, None] * inv_s[:, None, :] - b = XTWy_a * inv_s[:, :, None] - - # Tiny flat ridge on the equilibrated diagonal (now ~1 by - # construction) to keep Cholesky well-defined when columns - # are exactly collinear. - idx = np.arange(active_size) - A[:, idx, idx] += 1e-10 - - batch_size = A.shape[0] - w_norm, L = _cholesky_solve_batched(A, b) - if w_norm is None: - # Cholesky failed somewhere in the batch; per-entry - # lstsq safety net on the equilibrated system. - w_norm = _per_batch_lstsq(A, b) - - # Iterative refinement on the ORIGINAL system to claw back - # digits lost to normal-equation condition squaring. - # w0 = D^-1 w_norm is the candidate solution in original - # coordinates; the residual r = XTWy - XTWX @ w0 measures - # how much it misses the original equation; the correction - # dw_norm solves the same equilibrated system on D^-1 r and - # is unscaled back to dw. - w0 = w_norm * inv_s[:, :, None] - r = XTWy_a - XTWX_a @ w0 - r_norm = r * inv_s[:, :, None] - if L is not None: - try: - z = np.linalg.solve(L, r_norm) - dw_norm = np.linalg.solve(L.transpose(0, 2, 1), z) - except np.linalg.LinAlgError: - dw_norm = _per_batch_lstsq(A, r_norm) - else: - dw_norm = _per_batch_lstsq(A, r_norm) - w = w0 + dw_norm * inv_s[:, :, None] - - all_weights.append(w.squeeze(-1)) - return np.vstack(all_weights) - - @classmethod - def precompute_super(cls, Z, sample_weights, grid_shape, - circular_cv: bool = _DEFAULT_CIRCULAR_CV): - """Build a per-dim super-Gram over ``Z_aug = column_stack(Z, ones)`` - once. Returns an opaque dict consumed by :meth:`from_full` to - derive per-target GramSetup views via pure slicing. - - Used by EqRightPartSelector's term-sweep: instead of rebuilding - the windowed XTWX matrix for each candidate target column (which - does the same reshape/window/matmul on the SAME underlying Z), - precompute once over the full Z and slice out target-specific - sub-blocks. The math is exact -- (Z[:, ~t])^T W Z[:, ~t] is the - sub-block of (Z_aug)^T W Z_aug at rows/cols ``~t U intercept``. - - ``Z`` is shape (n_samples, n_terms); ``sample_weights`` is the - flat per-sample weight vector; ``grid_shape`` is the same shape - ``GramSetup.__init__`` consumes. ``circular_cv`` mirrors the - ``__init__`` flag -- callers that flow through both paths - (PhysicsInformedLasso single-call + EqRPS super-Gram sweep) must - keep these values aligned for the per-target views to remain - sub-blocks of the super-Gram (the math requires identical window - sets across the two passes). - """ - n_samples, n_terms = Z.shape - Z_aug = np.hstack([Z, np.ones((n_samples, 1))]) - n_features_aug = n_terms + 1 - - Z_grid = Z_aug.reshape(*grid_shape, n_features_aug) - sw_grid = sample_weights.reshape(*grid_shape) - per_dim_super = [] - - for dim in range(len(grid_shape)): - window_size = grid_shape[dim] // 2 - num_horizons = grid_shape[dim] if circular_cv else window_size + 1 - step_size = max(1, num_horizons // 30) - - Z_windows = _windowed_take(Z_grid, dim, window_size, - num_horizons, step_size, circular_cv) - w_windows = _windowed_take(sw_grid, dim, window_size, - num_horizons, step_size, circular_cv) - - Z_windows = np.moveaxis(Z_windows, dim, 0) - w_windows = np.moveaxis(w_windows, dim, 0) - Z_windows = np.moveaxis(Z_windows, -2, -1) - - batch_size = Z_windows.shape[0] - Z_batch = Z_windows.reshape(batch_size, -1, n_features_aug) - weights_batch = w_windows.reshape(batch_size, -1, 1) - - ZTW = Z_batch.transpose(0, 2, 1) * weights_batch.transpose(0, 2, 1) - XTWX_super = ZTW @ Z_batch - - diag = np.diagonal(XTWX_super, axis1=1, axis2=2) - scales_super = np.sqrt(np.maximum(np.abs(diag), 1e-30)) - - per_dim_super.append((XTWX_super, scales_super)) - - return { - 'per_dim_super': per_dim_super, - 'n_features_aug': n_features_aug, - 'grid_shape': grid_shape, - 'n_terms': n_terms, - # Cached so downstream VWSRSparsity can derive per-target - # ``target`` / ``features`` by slicing instead of re-calling - # objective.evaluate(normalize=True) -- which would force - # another vstack + transpose of the same term evaluations - # for every candidate target_idx in the sweep. - 'Z': Z, - } - - @classmethod - def from_full(cls, super_data, target_idx_in_terms): - """Construct a per-target GramSetup view from precomputed - super-Gram data via slicing. - - ``super_data`` is the dict returned by :meth:`precompute_super`. - ``target_idx_in_terms`` is the column index within Z (the terms - portion) that the caller wants as the regression target; the - intercept column stays in the feature set automatically. - - The returned object has the same ``n_features_aug``, ``grid_shape`` - and ``_per_dim`` shape contract as a regular ``GramSetup`` so - downstream code (``PhysicsInformedLasso.fit`` -> ``solve``) is - unchanged. - """ - per_dim_super = super_data['per_dim_super'] - n_features_aug_super = super_data['n_features_aug'] - grid_shape = super_data['grid_shape'] - - if not (0 <= target_idx_in_terms < n_features_aug_super - 1): - raise IndexError( - f'target_idx_in_terms={target_idx_in_terms} out of range ' - f'[0, {n_features_aug_super - 1}) for super-Gram with ' - f'{n_features_aug_super - 1} terms (+1 intercept).' - ) - - active = np.ones(n_features_aug_super, dtype=bool) - active[target_idx_in_terms] = False - - instance = cls.__new__(cls) - instance.n_features_aug = int(active.sum()) # n_terms (incl. intercept) - instance.grid_shape = grid_shape - instance._per_dim = [] - for XTWX_super, scales_super in per_dim_super: - XTWX_target = XTWX_super[:, active, :][:, :, active] - # XTWy = (Z_aug[:, ~t])^T W Z[:, t] = column t of XTWX_super - # at rows in ``active``. - XTWy_target = XTWX_super[:, active, - target_idx_in_terms:target_idx_in_terms + 1] - scales_target = scales_super[:, active] - instance._per_dim.append((XTWX_target, XTWy_target, scales_target)) - return instance - - -def calculate_weights(X, y, sample_weights, grid_shape, fit_intercept=True): - """ - Vectorized calculation of weights across sliding windows. - Dynamically handles whether the intercept should be fit. - - Single-shot wrapper over :class:`GramSetup`: builds the precomputed - Gram once and immediately solves with the requested intercept policy. - Callers that solve the same Gram against many active masks (e.g. - :class:`PhysicsInformedLasso.fit`) should instantiate ``GramSetup`` - directly and call ``.solve(active_mask)`` per iteration to avoid - re-running the expensive ``X^T diag(w) X`` matmul. - """ - setup = GramSetup(X, y, sample_weights, grid_shape) - active_mask = np.ones(setup.n_features_aug, dtype=bool) - if not fit_intercept: - # GramSetup always augments with the intercept column; drop it - # from the active set to mimic the legacy ``fit_intercept=False`` - # branch (which never augmented in the first place). - active_mask[-1] = False - return setup.solve(active_mask) diff --git a/projects/pic/data/ac/ac.py b/projects/pic/data/ac/ac.py index b51f1b86..d8900939 100644 --- a/projects/pic/data/ac/ac.py +++ b/projects/pic/data/ac/ac.py @@ -136,7 +136,7 @@ def ac_discovery(foldername, noise_level): dimensionality = data.ndim - 1 - epde_search_obj = EpdeSearch(use_solver=True, multiobjective_mode=True, + epde_search_obj = EpdeSearch(use_solver=False, multiobjective_mode=True, use_pic=True, boundary=(5, 10), verbose_params = {'show_iter_idx' : True, 'show_iter_fitness' : True}, coordinate_tensors=grid, device='cuda') @@ -218,5 +218,5 @@ def ac_discovery(foldername, noise_level): directory = os.path.dirname(os.path.realpath(__file__)) ac_folder_name = os.path.join(directory) - AC_test(fit_operator, ac_folder_name, 0) - # ac_discovery(ac_folder_name, 0) + # AC_test(fit_operator, ac_folder_name, 0) + ac_discovery(ac_folder_name, 0) diff --git a/projects/pic/data/heat_solar/heat_solar.py b/projects/pic/data/heat_solar/heat_solar.py index 4480e9d0..298b903c 100644 --- a/projects/pic/data/heat_solar/heat_solar.py +++ b/projects/pic/data/heat_solar/heat_solar.py @@ -378,6 +378,7 @@ def hs_3d_discovery(foldername, noise_level): ac_folder_name = os.path.join(directory) # hs_test(fit_operator, ac_folder_name, 0) - hs_discovery(ac_folder_name, 0) - # hs_2d_discovery(ac_folder_name, 0) + # hs_discovery(ac_folder_name, 0) + hs_2d_discovery(ac_folder_name, 0) # hs_3d_discovery(ac_folder_name, 0) + diff --git a/projects/pic/data/kdv/kdv.py b/projects/pic/data/kdv/kdv.py index fab95f0a..0c44fe48 100644 --- a/projects/pic/data/kdv/kdv.py +++ b/projects/pic/data/kdv/kdv.py @@ -421,7 +421,7 @@ def kdv_sga_discovery(foldername, noise_level): def kdv_sindy_discovery(foldername, noise_level): grid, data = kdv_sindy_data(os.path.join(foldername, 'kdv_sindy.mat')) noised_data = noise_data(data, noise_level) - data_nn = load_pretrained_PINN(os.path.join(foldername, f'kdv_{noise_level}_ann.pickle')) + # data_nn = load_pretrained_PINN(os.path.join(foldername, f'kdv_{noise_level}_ann.pickle')) dimensionality = data.ndim - 1 @@ -436,7 +436,7 @@ def kdv_sindy_discovery(foldername, noise_level): popsize = 16 epde_search_obj.set_moeadd_params(population_size=popsize, - training_epochs=1) + training_epochs=3) custom_trigonometric_eval_fun = { 'cos(t)sin(x)': lambda *grids, **kwargs: (np.cos(grids[0]) * np.sin(grids[1])) ** kwargs['power']} @@ -476,8 +476,8 @@ def kdv_sindy_discovery(foldername, noise_level): from epde.operators.utils.default_parameter_loader import EvolutionaryParams print("CUDA available:", torch.cuda.is_available()) # Operator = fitness.SolverBasedFitness # Replace by the developed PIC-based operator. - Operator = fitness.PIC - # Operator = fitness.L2LRFitness + # Operator = fitness.PIC + Operator = fitness.L2LRFitness params = EvolutionaryParams() operator_params = params.get_default_params_for_operator('DiscrepancyBasedFitnessWithCV') #{"penalty_coeff": 0.2, "pinn_loss_mult": 1e4} # operator_params = {"penalty_coeff": 0.2, "pinn_loss_mult": 1e4} @@ -488,11 +488,11 @@ def kdv_sindy_discovery(foldername, noise_level): directory = os.path.dirname(os.path.realpath(__file__)) kdv_folder_name = os.path.join(directory) - KdV_test(fit_operator, kdv_folder_name, 0) + # KdV_test(fit_operator, kdv_folder_name, 0) # KdV_h_test(fit_operator, kdv_folder_name, 0) # KdV_sga_test(fit_operator, kdv_folder_name, 0) # kdv_discovery(kdv_folder_name, 0) # kdv_h_discovery(kdv_folder_name, 0) # kdv_sga_discovery(kdv_folder_name, 5) - # kdv_sindy_discovery(kdv_folder_name, 0) \ No newline at end of file + kdv_sindy_discovery(kdv_folder_name, 0) \ No newline at end of file diff --git a/projects/pic/data/lorenz/lorenz.py b/projects/pic/data/lorenz/lorenz.py index c78e9c90..f667c6d0 100644 --- a/projects/pic/data/lorenz/lorenz.py +++ b/projects/pic/data/lorenz/lorenz.py @@ -279,8 +279,8 @@ def lorenz_discovery(noise_level): # print('operator_params ', operator_params) fit_operator = prepare_suboperators(Operator(list(operator_params.keys())), operator_params) - #lorenz_discovery(0) - lorenz_test(fit_operator, noise_level=0) + lorenz_discovery(0) + # lorenz_test(fit_operator, noise_level=0) def get_pic_network_summary(operator): diff --git a/projects/pic/data/lv/lv.py b/projects/pic/data/lv/lv.py index fa0821d8..3901c9e1 100644 --- a/projects/pic/data/lv/lv.py +++ b/projects/pic/data/lv/lv.py @@ -107,12 +107,12 @@ def lv_discovery(noise_level): epde_search_obj.set_preprocessor(default_preprocessor_type='FD', preprocessor_kwargs={}) - popsize = 16 - epde_search_obj.set_moeadd_params(population_size=popsize, training_epochs=1) + popsize = 32 + epde_search_obj.set_moeadd_params(population_size=popsize, training_epochs=5) factors_max_number = {'factors_num': [1, 2], 'probas' : [0.8, 0.2]} - epde_search_obj.fit(data=[x, y], variable_names=['u', 'v'], max_deriv_order=(1,), + epde_search_obj.fit(data=[x, y], variable_names=['u', 'v'], max_deriv_order=(2,), equation_terms_max_number=7, data_fun_pow=3, additional_tokens=[trig_tokens, grid_tokens], equation_factors_max_number=factors_max_number, eq_sparsity_interval=(1e-8, 1e-0)) # @@ -135,4 +135,4 @@ def lv_discovery(noise_level): print('operator_params ', operator_params) fit_operator = prepare_suboperators(Operator(list(operator_params.keys())), operator_params) - #lv_discovery(0) + lv_discovery(0) diff --git a/projects/pic/data/ns/ns.py b/projects/pic/data/ns/ns.py index ce0635e4..d6ec322c 100644 --- a/projects/pic/data/ns/ns.py +++ b/projects/pic/data/ns/ns.py @@ -242,10 +242,10 @@ def ns_discovery(foldername, noise_level): # preprocessor_kwargs={'epochs_max' : 1e3}) epde_search_obj.set_preprocessor(default_preprocessor_type='FD', preprocessor_kwargs={}) - popsize = 32 + popsize = 64 epde_search_obj.set_moeadd_params(population_size=popsize, - training_epochs=5) + training_epochs=30) custom_grid_tokens = CacheStoredTokens(token_type='grid', token_labels=['t', 'x'], @@ -289,5 +289,5 @@ def ns_discovery(foldername, noise_level): directory = os.path.dirname(os.path.realpath(__file__)) ns_folder_name = os.path.join(directory) - ns_test(fit_operator, ns_folder_name, 0) - # ns_discovery(ns_folder_name, 0) + # ns_test(fit_operator, ns_folder_name, 0) + ns_discovery(ns_folder_name, 0) diff --git a/projects/pic/data/ode/ode.py b/projects/pic/data/ode/ode.py index a9997609..5b3f20c5 100644 --- a/projects/pic/data/ode/ode.py +++ b/projects/pic/data/ode/ode.py @@ -227,6 +227,6 @@ def ODE_simple_discovery(foldername, noise_level): ode_folder_name = os.path.join(directory) # ODE_test(fit_operator, ode_folder_name, 0) - # ODE_discovery(ode_folder_name, 0) - ODE_simple_discovery(ode_folder_name, 0) + ODE_discovery(ode_folder_name, 0) + # ODE_simple_discovery(ode_folder_name, 0) diff --git a/projects/pic/data/vdp/vdp.py b/projects/pic/data/vdp/vdp.py index 84484c0c..89ed88d6 100644 --- a/projects/pic/data/vdp/vdp.py +++ b/projects/pic/data/vdp/vdp.py @@ -147,7 +147,7 @@ def vdp_discovery(foldername, noise_level): preprocessor_kwargs={}) popsize = 16 - epde_search_obj.set_moeadd_params(population_size=popsize, training_epochs=1) + epde_search_obj.set_moeadd_params(population_size=popsize, training_epochs=5) factors_max_number = {'factors_num': [1, 2], 'probas': [0.65, 0.35]} diff --git a/projects/pic/data/wave/wave.py b/projects/pic/data/wave/wave.py index 501d3e12..569a032b 100644 --- a/projects/pic/data/wave/wave.py +++ b/projects/pic/data/wave/wave.py @@ -197,5 +197,5 @@ def wave_discovery(foldername, noise_level): directory = os.path.dirname(os.path.realpath(__file__)) wave_folder_name = os.path.join(directory) - wave_test(fit_operator, wave_folder_name, 0) - # wave_discovery(wave_folder_name, 0) \ No newline at end of file + # wave_test(fit_operator, wave_folder_name, 0) + wave_discovery(wave_folder_name, 0) \ No newline at end of file diff --git a/projects/thesis/_numcheck.py b/projects/thesis/_numcheck.py new file mode 100644 index 00000000..5f9a3bab --- /dev/null +++ b/projects/thesis/_numcheck.py @@ -0,0 +1,168 @@ +"""Numerical-correctness check of VaryingCoefSetup on the REAL 14-dataset +inputs (the seeded-truth features/grids/g_func weights each system feeds the +estimator). Verifies, per (system, variable): + + A. G symmetric & finite (direct + super-Gram). + B. super/from_full == direct construction (G, Phiy, yWy, score) -- the + EqRPS fast path used in the live search. + C. gamma_0 solves the weighted normal equations (|X^T W (y - X g0)| ~ 0) -- + robust to collinearity, unlike comparing to a separate WLS solve. + D. Var(gamma_0) matches sigma^2 * diag((X^T W X)^-1) (reported; the + equilibration floor makes it approximate on near-singular blocks). + E. score finite, >= 0, no NaN/Inf anywhere. + F. Parseval (exact, basis property): var(beta(x)) == NC_raw per feature. + +Usage: python _numcheck.py [system ...] (default: all 14) +Run, then delete.""" +from __future__ import annotations +import os, sys, traceback + +_THIS = os.path.dirname(os.path.abspath(__file__)) +_ROOT = os.path.abspath(os.path.join(_THIS, '..', '..')) +for _p in (_ROOT, _THIS): + if _p not in sys.path: + sys.path.insert(0, _p) + +import numpy as np +import yaml +import epde.globals as gv +from epde.operators.common.stability import VaryingCoefSetup as VC +from kdv_sindy_test import build_pool_only, _normalize_grid_labels +from thesis_runner import load_config, pipeline_settings, _set_seeds +from vcoef_stat_compare import _ALL +from epde.interface.equation_translator import translate_equation + + +def _inputs(system): + """[(var, Z (N,n_terms), target_idx, w (N,), grid_shape), ...] for truth.""" + cfg = load_config(system) + _set_seeds(0) + gv.set_gram_config('vcoef') + search = build_pool_only(cfg, pipeline_settings('new')) + coords, data, variable_names, dim = cfg.load_data() + all_vars = list(variable_names) + truth = yaml.safe_load(open(os.path.join(_THIS, 'configs', f'{system}.yaml'))) + teqs = truth.get('truth_equations') or [] + seeded = (teqs[0] if len(all_vars) == 1 + else {v: teqs[i] for i, v in enumerate(all_vars)}) + seeded = _normalize_grid_labels(seeded) + soeq = translate_equation(seeded, search.pool, all_vars=all_vars) + out = [] + for v in all_vars: + eq = soeq.vals[v] + eq.main_var_to_explain = v + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + eq.evaluate(normalize=False, return_val=False) # populate grid cache + Z = np.vstack([t.evaluate(False, grids=None) + for t in eq.structure]).T.astype(float) + w = np.asarray(gv.grid_cache.g_func[gv.grid_cache.g_func_mask], float).reshape(-1) + gshape = tuple(int(n) for n in gv.grid_cache.inner_shape) + out.append((v, Z, int(eq.target_idx), w, gshape)) + return out + + +def _check(v, Z, tgt, w, gshape): + N, n_terms = Z.shape + feat_idx = [i for i in range(n_terms) if i != tgt] + Xf = Z[:, feat_idx] + yt = Z[:, tgt] + direct = VC(Xf, yt, w, gshape, main_var=v, fit_intercept=True) + sup = VC.precompute_super(Z, w, gshape, main_var=v) + ff = VC.from_full(sup, tgt) + + m = {} + # A. symmetric & finite + G = direct.G + m['Gsym'] = float(np.abs(G - G.T).max() / (np.abs(G).max() + 1e-30)) + m['finite'] = bool(np.all(np.isfinite(G)) and np.all(np.isfinite(direct.Phiy)) + and np.all(np.isfinite(sup['G_super']))) + # B. super/from_full == direct + m['dG'] = float(np.abs(ff.G - direct.G).max() / (np.abs(direct.G).max() + 1e-30)) + m['dPhiy'] = float(np.abs(ff.Phiy - direct.Phiy).max() / (np.abs(direct.Phiy).max() + 1e-30)) + m['dyWy'] = float(abs(ff.yWy - direct.yWy) / (abs(direct.yWy) + 1e-30)) + sc_d = direct.score(None) + sc_f = ff.score(None) + m['dscore'] = float(np.abs(sc_f - sc_d).max()) + # C. gamma_0 weighted normal equations + sol = direct._solve_gammas(None) + B = sol['B']; nf = sol['nf'] + g0 = sol['gamma'][np.arange(nf) * B] # const per feature (incl intercept) + Xa = np.column_stack([Xf, np.ones(N)]) + XtWy = Xa.T @ (w * yt) + resid = Xa.T @ (w * (yt - Xa @ g0)) + m['normeq'] = float(np.abs(resid).max() / (np.abs(XtWy).max() + 1e-30)) + # D. Var(gamma_0) vs sigma^2 diag((X^T W X)^-1) + A = Xa.T @ (w[:, None] * Xa) + Neff = float(w.sum()) + rss = max(float(yt @ (w * yt) - g0 @ XtWy), 0.0) + sigma2 = rss / max(Neff - nf, 1.0) + try: + var_ref = sigma2 * np.diag(np.linalg.inv(A)) + var_vc = sol['var'][np.arange(nf) * B] + rel = np.abs(var_vc - var_ref) / (np.abs(var_ref) + 1e-30) + m['dVar'] = float(np.nanmax(rel)) + except np.linalg.LinAlgError: + m['dVar'] = float('nan') + m['condA'] = float(np.linalg.cond(A)) + # E. score sane + m['score_ok'] = bool(np.all(np.isfinite(sc_d)) and np.all(sc_d >= -1e-9)) + # F. Parseval exact + st = direct.beta_field_stats(None) + g = sol['gamma'] + par = 0.0 + for i in range(nf): + nc_raw = float(np.sum(g[i * B + 1:(i + 1) * B] ** 2)) + par = max(par, abs(st['std'][i] ** 2 - nc_raw) / (nc_raw + 1e-12)) + m['parseval'] = float(par) + return m + + +# tolerances +TOL = dict(Gsym=1e-10, dG=1e-7, dPhiy=1e-7, dyWy=1e-9, dscore=1e-6, + normeq=1e-6, parseval=1e-6) + + +def verdict(m): + bad = [] + if not m['finite']: + bad.append('NONFINITE') + if not m['score_ok']: + bad.append('score') + for k, t in TOL.items(): + if not (m[k] <= t): + bad.append(f'{k}={m[k]:.1e}') + return bad + + +def main(): + systems = sys.argv[1:] or list(_ALL) + n_ok = n_tot = 0 + for s in systems: + try: + rows = _inputs(s) + except Exception as e: + print(f"{s:18s} INPUT-ERROR {type(e).__name__}: {str(e)[:60]}") + continue + for (v, Z, tgt, w, gshape) in rows: + n_tot += 1 + try: + m = _check(v, Z, tgt, w, gshape) + except Exception as e: + print(f"{s:14s}/{v:3s} CHECK-ERROR {type(e).__name__}: {str(e)[:50]}") + traceback.print_exc() + continue + bad = verdict(m) + tag = f"{s}/{v}" if len(rows) > 1 else s + if not bad: + n_ok += 1 + print(f"{tag:18s} OK superGmax={m['dG']:.0e} normeq={m['normeq']:.0e} " + f"parseval={m['parseval']:.0e} dVar={m['dVar']:.0e} condA={m['condA']:.0e}") + else: + print(f"{tag:18s} FAIL {', '.join(bad)} (condA={m['condA']:.0e})") + print(f"\n==== numerically correct: {n_ok}/{n_tot} (system,var) blocks ====") + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/projects/thesis/_vc_cache_gate.py b/projects/thesis/_vc_cache_gate.py new file mode 100644 index 00000000..1ee4ae98 --- /dev/null +++ b/projects/thesis/_vc_cache_gate.py @@ -0,0 +1,118 @@ +"""Equivalence gate for the cv_cache-style vcoef stability reuse. + +For each system, seed the truth equation, run the real sparsity pass (which now +caches per-term stability scores on ``eq._cached_vc_score``), then check: + + (a) WIRING / no-super equality: + sum(eq._cached_vc_score) == vc_stability_total_lr(normalize=True feats) + Both use the same normalize=True regime + the same score(); must match to + ~fp. This is exactly what the fitness reuse branch sums. + + (b) SCALE-INVARIANCE (the super-Gram path assumption): + vc_stability_total_lr(normalize=True) ~= vc_stability_total_lr(normalize=False) + The super-Gram is built from RAW term values while the in-place fitness + uses normalized features; reuse is lossless only if the score is invariant + to per-column scaling. This checks it on real equations. + +PASS => the cache reuse is lossless under the active vc_score_formula. +""" +from __future__ import annotations +import os +import sys + +_THIS = os.path.dirname(os.path.abspath(__file__)) +_ROOT = os.path.abspath(os.path.join(_THIS, '..', '..')) +for _p in (_ROOT, _THIS): + if _p not in sys.path: + sys.path.insert(0, _p) + +import numpy as np +import yaml +import epde.globals as gv +from epde.interface.equation_translator import translate_equation +from epde.operators.common.fitness import vc_stability_total_lr +from kdv_sindy_test import build_pool_only, make_fit_operator, _normalize_grid_labels +from thesis_runner import load_config, pipeline_settings, _set_seeds + +_SYSTEMS = ['lorenz', 'ks', 'ac'] +_TOL = 1e-9 + + +def _names(eq): + return [t.name for i, t in enumerate(eq.structure) if i != eq.target_idx] + + +def check(system): + cfg = load_config(system) + _set_seeds(0) + gv.set_gram_config('vcoef') + search = build_pool_only(cfg, pipeline_settings('new')) + fit_op = make_fit_operator() + coords, data, variable_names, dim = cfg.load_data() + all_vars = list(variable_names) + truth = yaml.safe_load(open(os.path.join(_THIS, 'configs', f'{system}.yaml'))) + teqs = truth.get('truth_equations') or [] + seeded = (teqs[0] if len(all_vars) == 1 + else {v: teqs[i] for i, v in enumerate(all_vars)}) + seeded = _normalize_grid_labels(seeded) + tsoeq = translate_equation(seeded, search.pool, all_vars=all_vars) + + g_fun = gv.grid_cache.g_func[gv.grid_cache.g_func_mask].reshape(-1) + data_shape = gv.grid_cache.inner_shape + + rows = [] + for v in all_vars: + eq = tsoeq.vals[v] + eq.main_var_to_explain = v + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + # Run the real sparsity (force_out_of_place) -> sets _cached_vc_score. + fit_op.apply(eq, {}, force_out_of_place=True) + + cached_vec = getattr(eq, '_cached_vc_score', None) + cached = None if cached_vec is None else float(np.sum(cached_vec)) + + fit_int = bool(eq.weights_internal[-1] != 0) + _, t_n, f_n = eq.evaluate(normalize=True, return_val=False) + _, t_r, f_r = eq.evaluate(normalize=False, return_val=False) + fresh_norm = (None if f_n is None else + vc_stability_total_lr(f_n, t_n, g_fun, data_shape, + main_var=v, fit_intercept=fit_int)) + fresh_raw = (None if f_r is None else + vc_stability_total_lr(f_r, t_r, g_fun, data_shape, + main_var=v, fit_intercept=fit_int)) + rows.append((v, cached, fresh_norm, fresh_raw)) + return rows + + +def main(): + print(f"vc_score_formula = {getattr(gv, 'vc_score_formula', '?')} tol={_TOL}\n") + all_ok = True + for system in _SYSTEMS: + try: + rows = check(system) + except Exception as e: + import traceback + print(f"{system:8s} ERROR {type(e).__name__}: {str(e)[:80]}") + traceback.print_exc() + all_ok = False + continue + for (v, cached, fn, fr) in rows: + # (a) wiring equality: cached == fresh_norm + wire_ok = (cached is not None and fn is not None + and abs(cached - fn) <= _TOL * (1 + abs(fn))) + # (b) scale-invariance: fresh_norm ~= fresh_raw + si_rel = (abs(fn - fr) / (1 + abs(fn))) if (fn is not None and fr is not None) else float('nan') + si_ok = si_rel <= 1e-6 + flag = '' if (wire_ok and si_ok) else ' <-- FAIL' + if not (wire_ok and si_ok): + all_ok = False + print(f"{system:8s}[{v}] cached={cached!r} fresh_norm={fn!r} " + f"fresh_raw={fr!r} wire_ok={wire_ok} si_rel={si_rel:.2e}{flag}") + print("\n" + ("ALL PASS" if all_ok else "SOME FAILED")) + return 0 if all_ok else 1 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/projects/thesis/_wave_why.py b/projects/thesis/_wave_why.py new file mode 100644 index 00000000..77cedad3 --- /dev/null +++ b/projects/thesis/_wave_why.py @@ -0,0 +1,89 @@ +"""Diagnose why the spurious wave solution beats the true wave equation. + +Spurious (Pareto-0): 48.679*u_xx*u_tt - 591.98*u_tt^2 = u_xx^2 (target = u_xx^2) +True: 0.04*u_xx = u_tt (u_tt = c^2 u_xx, c^2=0.04) + +Hypothesis: on the wave solution u_tt = 0.04 u_xx, the three quadratic-in-2nd- +derivative terms u_xx^2, u_xx*u_tt, u_tt^2 are all proportional, so the spurious +equation is an EXACT algebraic identity -> fits ~perfectly with constant coefs, +dominating the (FD-limited) true equation on (fitness, stability). Run once, delete. +""" +import os, sys +_THIS = os.path.dirname(os.path.abspath(__file__)) +_ROOT = os.path.abspath(os.path.join(_THIS, "..", "..")) +for _p in (_ROOT, _THIS): + if _p not in sys.path: + sys.path.insert(0, _p) +import numpy as np +import epde.globals as gv +from epde.interface.equation_translator import translate_equation +from kdv_sindy_test import build_pool_only, make_fit_operator, _normalize_grid_labels +from thesis_runner import load_config, pipeline_settings, _set_seeds + +DISCOVERED = ("48.67869238111131 * d^2u/dx1^2{power: 1.0} * d^2u/dx0^2{power: 1.0} " + "+ -591.9797844706457 * d^2u/dx0^2{power: 2.0} + 0.0 = d^2u/dx1^2{power: 2.0}") +TRUE = "0.04 * d^2u/dx1^2{power: 1.0} = d^2u/dx0^2{power: 1.0}" + + +def evaluate_eq(search, fit_op, eq_str): + soeq = translate_equation(_normalize_grid_labels(eq_str), search.pool, all_vars=["u"]) + eq = soeq.vals["u"] + eq.main_var_to_explain = "u" + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + fit_op.apply(eq, {}, force_out_of_place=True) + return eq + + +def main(): + cfg = load_config("wave") + _set_seeds(0) + gv.set_gram_config("vcoef") + search = build_pool_only(cfg, pipeline_settings("new")) + fit_op = make_fit_operator() + + print("=" * 70) + print("OBJECTIVE VECTORS (both objectives: lower = better)") + print("=" * 70) + for name, s in [("DISCOVERED (spurious identity)", DISCOVERED), ("TRUE wave", TRUE)]: + eq = evaluate_eq(search, fit_op, s) + print(f"\n{name}") + print(f" fitness_value = {getattr(eq, 'fitness_value', None)!r}") + print(f" coefficients_stability = {getattr(eq, 'coefficients_stability', None)!r}") + print(f" obj_fun = {getattr(eq, 'obj_fun', None)!r}") + + # ---- identity check on the actual solution fields ---- + soeq = translate_equation(_normalize_grid_labels(TRUE), search.pool, all_vars=["u"]) + teq = soeq.vals["u"] + teq.main_var_to_explain = "u" + teq.weights_internal = np.ones(len(teq.structure) - 1) + teq.weights_internal_evald = True + teq.weights_final_evald = True + _, t, f = teq.evaluate(normalize=False, return_val=False) + u_tt = np.asarray(t, float).reshape(-1) + f = np.asarray(f, float) + u_xx = f[:, 0] if f.ndim > 1 else f.reshape(-1) + + a = float(np.linalg.lstsq(u_xx[:, None], u_tt, rcond=None)[0][0]) + A, B, C = u_xx ** 2, u_xx * u_tt, u_tt ** 2 + M = np.column_stack([B, C]) + coef, *_ = np.linalg.lstsq(M, A, rcond=None) + resid = A - M @ coef + r2 = 1.0 - np.sum(resid ** 2) / np.sum((A - A.mean()) ** 2) + + print("\n" + "=" * 70) + print("IDENTITY CHECK on the wave solution") + print("=" * 70) + print(f" u_tt = a * u_xx -> a = {a:.6f} (truth c^2 = 0.04)") + print(f" u_xx^2 = b1*(u_xx*u_tt) + b2*(u_tt^2):") + print(f" b1 = {coef[0]:.4f} (discovered 48.679)") + print(f" b2 = {coef[1]:.4f} (discovered -591.98)") + print(f" R^2 = {r2:.8f} ||resid||/||u_xx^2|| = {np.linalg.norm(resid)/np.linalg.norm(A):.3e}") + print(f" arithmetic: 48.6787*0.04 - 591.98*0.04^2 = " + f"{48.6787*0.04 - 591.98*0.04**2:.6f} (should be ~1)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/projects/thesis/analyze_ac_ode_cv.py b/projects/thesis/analyze_ac_ode_cv.py new file mode 100644 index 00000000..c8b7de56 --- /dev/null +++ b/projects/thesis/analyze_ac_ode_cv.py @@ -0,0 +1,150 @@ +"""Find a per-term CV that KEEPS ac's weak diffusion (true, coef 1e-4) and +PRUNES ode's spurious u^3*u' (false, coef 1.2e-4). Both have tiny, highly +significant coefficients, so the only signal is the debiased region-variation +NC_deb -- but the normaliser decides whether the actual Lasso prune fires. + +The Lasso keeps term j iff |rho_j| >= cv_j * max_corr, i.e. cv_j <= r_j, +where r_j = |rho_j|/max_corr is the term's relative correlation (rho computed +exactly as PhysicsInformedLasso iter-1: rho = X_aug^T (w*y)). So we want, for +each candidate cv: + ac-diffusion : cv <= r (KEEP) + ode-spurious : cv > r (PRUNE) + +Per-term quantities (gamma_0, var0, NC_raw, NC_deb, contrib=||gamma_0*phi||_w), +candidate cv formulas, and the verdict are tabulated for the two critical terms. +""" +from __future__ import annotations +import os +import sys + +_THIS = os.path.dirname(os.path.abspath(__file__)) +_ROOT = os.path.abspath(os.path.join(_THIS, '..', '..')) +for _p in (_ROOT, _THIS): + if _p not in sys.path: + sys.path.insert(0, _p) + +import numpy as np +import epde.globals as gv +from epde.interface.equation_translator import translate_equation +from epde.operators.common.stability import VaryingCoefSetup +from kdv_sindy_test import build_pool_only, _normalize_grid_labels +from thesis_runner import load_config, pipeline_settings, _set_seeds + +_EPS = 1e-12 + +# (system, equation string, set of TRUE non-target term names, critical term +# substring, want_keep) +ODE_SOL0 = ("-0.9936633679622026 * du/dx0{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 0.0} + -3.9868949002903613 * u{power: 1.0} + 0.00011940105848764598 * u{power: 3.0} * du/dx0{power: 1.0} + 1.4973740096088297 * x{power: 1.0, dim: 0.0} + 0.0 = d^2u/dx0^2{power: 1.0}") +AC_TRUTH = ("0.0001 * d^2u/dx1^2{power: 1.0} + -5.0 * u{power: 3.0} + 5.0 * u{power: 1.0} = du/dx0{power: 1.0}") + +CASES = [ + ('ode', ODE_SOL0, 'u{power: 3.0}', False), # spurious u^3*u' -> PRUNE + ('ac', AC_TRUTH, 'd^2u/dx1^2', True), # weak diffusion -> KEEP +] + +# candidate cv formulas of the per-term quantities +FORMULAS = { + 'var0/g0^2': lambda q: q['var0'] / (q['C'] + 1e-30), + 'NCraw/g0^2': lambda q: q['nc_raw'] / (q['C'] + 1e-30), + 'NCdeb/g0^2': lambda q: q['nc_deb'] / (q['C'] + 1e-30), + 'NCdeb': lambda q: q['nc_deb'], + 'NCdeb/var0': lambda q: q['nc_deb'] / (q['var0'] + 1e-30), + 'NCdeb/(g0^2+var0)':lambda q: q['nc_deb'] / (q['C'] + q['var0'] + 1e-30), + 'sqrt(NCdeb)/|g0|': lambda q: np.sqrt(q['nc_deb']) / (abs(q['g0']) + _EPS), + 'NCdeb/contrib^2': lambda q: q['nc_deb'] / (q['contrib'] ** 2 + 1e-30), + 'NCdeb/contrib': lambda q: q['nc_deb'] / (q['contrib'] + _EPS), + 'NCdeb*var0/g0^4': lambda q: q['nc_deb'] * q['var0'] / (q['C'] ** 2 + 1e-30), + 'NCdeb/(g0^2*contrib)': lambda q: q['nc_deb'] / (q['C'] * q['contrib'] + 1e-30), +} + + +def analyse(system, eq_str, crit_sub): + cfg = load_config(system) + _set_seeds(0) + search = build_pool_only(cfg, pipeline_settings('new')) + sw = np.asarray(gv.grid_cache.g_func[gv.grid_cache.g_func_mask]).reshape(-1) + gshape = gv.grid_cache.inner_shape + + soeq = translate_equation(_normalize_grid_labels(eq_str), search.pool, + all_vars=['u']) + eq = soeq.vals['u'] + eq.main_var_to_explain = 'u' + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + _, target, feats = eq.evaluate(normalize=True, return_val=False) + feats = np.asarray(feats, dtype=float) + y = np.asarray(target, dtype=float).reshape(-1) + feat_terms = [t for i, t in enumerate(eq.structure) if i != eq.target_idx] + n = len(feat_terms) + + # rho exactly as PhysicsInformedLasso iter-1: X_aug = [feats, 1], rho = X^T(w y) + X_aug = np.hstack([feats, np.ones((feats.shape[0], 1))]) + X_T_y = X_aug.T @ (sw * y) + max_corr = float(np.max(np.abs(X_T_y))) + r = np.abs(X_T_y) / (max_corr + 1e-30) # relative correlation per col + + setup = VaryingCoefSetup(feats, y, sw, gshape, main_var='u', + fit_intercept=False) + sol = setup._solve_gammas(None) + gamma, var, B, mk = sol['gamma'], sol['var'], sol['B'], sol['mk'] + is_const = mk == 0 + nonconst = ~is_const + + out = [] + for i in range(n): + sl = slice(i * B, (i + 1) * B) + g = gamma[sl] + v = var[sl] + g0 = float(g[is_const][0]) + var0 = float(v[is_const][0]) + nc_g2 = g[nonconst] ** 2 + nc_raw = float(np.sum(nc_g2)) + nc_deb = float(np.sum(np.maximum(nc_g2 - v[nonconst], 0.0))) + contrib = float(np.sqrt(np.sum(sw * (g0 * feats[:, i]) ** 2))) + q = {'g0': g0, 'C': g0 * g0, 'var0': var0, 'nc_raw': nc_raw, + 'nc_deb': nc_deb, 'contrib': contrib} + out.append({'name': feat_terms[i].name, 'r': float(r[i]), 'q': q, + 'crit': crit_sub in feat_terms[i].name}) + return out + + +def main(): + crit = {} + for system, eq_str, crit_sub, want_keep in CASES: + rows = analyse(system, eq_str, crit_sub) + print(f"\n##### {system}: r=|rho|/max_corr per term (keep iff cv<=r)") + for row in rows: + star = ' *<-- CRITICAL' if row['crit'] else '' + print(f" {row['name'][:40]:40s} r={row['r']:.4g} " + f"g0={row['q']['g0']:.3g} NCdeb={row['q']['nc_deb']:.3g} " + f"contrib={row['q']['contrib']:.3g}{star}") + for row in rows: + if row['crit']: + crit[system] = {'r': row['r'], 'q': row['q'], + 'want_keep': want_keep} + + # the decisive table: for each formula, cv and verdict on both critical terms + print(f"\n{'='*78}\nFORMULA SCREEN (need ac KEEP: cv<=r AND ode PRUNE: cv>r)\n{'='*78}") + hdr = (f"{'formula':22s} | {'ac cv':>10} {'ac r':>9} {'keep?':>6} | " + f"{'ode cv':>10} {'ode r':>9} {'prune?':>7} | {'BOTH':>5}") + print(hdr) + print('-' * len(hdr)) + for fname, f in FORMULAS.items(): + ac = crit['ac'] + ode = crit['ode'] + ac_cv = float(f(ac['q'])) + ode_cv = float(f(ode['q'])) + ac_keep = ac_cv <= ac['r'] + ode_prune = ode_cv > ode['r'] + both = ac_keep and ode_prune + print(f"{fname:22s} | {ac_cv:>10.3g} {ac['r']:>9.3g} " + f"{('yes' if ac_keep else 'NO'):>6} | " + f"{ode_cv:>10.3g} {ode['r']:>9.3g} " + f"{('yes' if ode_prune else 'NO'):>7} | " + f"{('YES' if both else '-'):>5}") + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/projects/thesis/analyze_ode_cv.py b/projects/thesis/analyze_ode_cv.py new file mode 100644 index 00000000..c1a5bd97 --- /dev/null +++ b/projects/thesis/analyze_ode_cv.py @@ -0,0 +1,150 @@ +"""Analyze the 4 ode Pareto solutions to find a per-term CV that prunes the +tiny-coefficient spurious terms (the systematic hamming=1) while keeping the +true terms. + +Each discovered ode equation has the 3 TRUE terms (u, du/dx0*sin, x; coef O(1)) +plus 1-2 SPURIOUS terms with coef 1e-4..1e-2. The current CV var0/g0^2 (=1/t^2) +keeps the spurious ones because on clean data they are precisely estimated +(significant). We tabulate candidate per-term statistics, tag true/spurious, and +look for one where spurious >> true so it could drive pruning. + +Candidates (per term j; gamma_0 = const coef, var0 = Var(gamma_0), gamma_k = +basis modes, phi_j = feature column, y = target, w = sample weights): + cv_sig = var0 / gamma_0^2 significance 1/t^2 (current) + ncraw_g2 = sum gamma_k^2 / gamma_0^2 region-variation ratio + mad_med = mad_x(beta)/|median_x(beta)| robust region-variation + contrib = ||gamma_0 phi_j||_w the signal the term adds to y + frac = contrib / ||y||_w fraction of target explained + inv_frac = ||y||_w / contrib small-contribution -> large + relmag = max_k(gamma_0_k^2) / gamma_0_j^2 inverse relative magnitude^2 + sig_x_mag = cv_sig * relmag significance scaled by smallness +""" +from __future__ import annotations +import os +import sys + +_THIS = os.path.dirname(os.path.abspath(__file__)) +_ROOT = os.path.abspath(os.path.join(_THIS, '..', '..')) +for _p in (_ROOT, _THIS): + if _p not in sys.path: + sys.path.insert(0, _p) + +import numpy as np +import epde.globals as gv +from epde.interface.equation_translator import translate_equation +from epde.operators.common.stability import VaryingCoefSetup +from kdv_sindy_test import build_pool_only, _normalize_grid_labels +from thesis_runner import load_config, pipeline_settings, _set_seeds + +_EPS = 1e-12 + +TRUTH = ("-4.0 * u{power: 1.0} + -1.0 * du/dx0{power: 1.0} * " + "sin{power: 1.0, freq: 2.0, dim: 0.0} + " + "1.5 * x{power: 1.0, dim: 0.0} = d^2u/dx0^2{power: 1.0}") + +EQS = [ + "-0.9936633679622026 * du/dx0{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 0.0} + -3.9868949002903613 * u{power: 1.0} + 0.00011940105848764598 * u{power: 3.0} * du/dx0{power: 1.0} + 1.4973740096088297 * x{power: 1.0, dim: 0.0} + 0.0 = d^2u/dx0^2{power: 1.0}", + "-0.008456117004594182 * du/dx0{power: 1.0} * d^2u/dx0^2{power: 1.0} + -3.999889678771497 * u{power: 1.0} + 1.4981092221855905 * x{power: 1.0, dim: 0.0} + -1.0293913688973846 * du/dx0{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 0.0} + 0.0 = d^2u/dx0^2{power: 1.0}", + "-3.999955869355026 * u{power: 1.0} + 1.4980032313263567 * x{power: 1.0, dim: 0.0} + -0.9757548527888852 * du/dx0{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 0.0} + -0.0036890595533533733 * d^2u/dx0^2{power: 1.0} * du/dx0{power: 1.0} + 0.011783526161822731 * du/dx0{power: 2.0} + 0.0 = d^2u/dx0^2{power: 1.0}", + "-3.9994631249571406 * u{power: 1.0} + 1.4983573574999929 * x{power: 1.0, dim: 0.0} + -1.02610182727155 * du/dx0{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 0.0} + -0.00742408614998562 * du/dx0{power: 1.0} * d^2u/dx0^2{power: 1.0} + 0.0003709182161524908 * du/dx0{power: 1.0} * u{power: 2.0} + 0.0 = d^2u/dx0^2{power: 1.0}", +] + + +def main(): + cfg = load_config('ode') + _set_seeds(0) + search = build_pool_only(cfg, pipeline_settings('new')) + sw = np.asarray(gv.grid_cache.g_func[gv.grid_cache.g_func_mask]).reshape(-1) + gshape = gv.grid_cache.inner_shape + + tsoeq = translate_equation(_normalize_grid_labels(TRUTH), search.pool, + all_vars=['u']) + teq = tsoeq.vals['u'] + truth_names = {t.name for i, t in enumerate(teq.structure) + if i != teq.target_idx} + + for idx, eqs in enumerate(EQS): + soeq = translate_equation(_normalize_grid_labels(eqs), search.pool, + all_vars=['u']) + eq = soeq.vals['u'] + eq.main_var_to_explain = 'u' + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + _, target, feats = eq.evaluate(normalize=True, return_val=False) + feats = np.asarray(feats, dtype=float) + y = np.asarray(target, dtype=float).reshape(-1) + feat_terms = [t for i, t in enumerate(eq.structure) + if i != eq.target_idx] + n = len(feat_terms) + + setup = VaryingCoefSetup(feats, y, sw, gshape, main_var='u', + fit_intercept=False) + sol = setup._solve_gammas(None) + gamma, var, B, mk = sol['gamma'], sol['var'], sol['B'], sol['mk'] + Bvals = setup._Bvals + is_const = mk == 0 + nonconst = ~is_const + + yw = float(np.sqrt(np.sum(sw * y * y))) + rows = [] + g0s = np.empty(n) + for i in range(n): + sl = slice(i * B, (i + 1) * B) + g = gamma[sl] + v = var[sl] + g0 = float(g[is_const][0]) + g0s[i] = g0 + maxg2 = float(np.max(g0s ** 2)) + for i in range(n): + sl = slice(i * B, (i + 1) * B) + g = gamma[sl] + v = var[sl] + g0 = float(g[is_const][0]) + var0 = float(v[is_const][0]) + C = g0 * g0 + nc_raw = float(np.sum(g[nonconst] ** 2)) + beta = Bvals @ g + med = float(np.median(beta)) + mad = float(np.median(np.abs(beta - med))) + contrib = float(np.sqrt(np.sum(sw * (g0 * feats[:, i]) ** 2))) + frac = contrib / (yw + _EPS) + rows.append({ + 'name': feat_terms[i].name, + 'true': feat_terms[i].name in truth_names, + 'g0': g0, + 'cv_sig': var0 / (C + 1e-30), + 'ncraw_g2': nc_raw / (C + 1e-30), + 'mad_med': mad / (abs(med) + _EPS), + 'frac': frac, + 'inv_frac': 1.0 / (frac + _EPS), + 'relmag': maxg2 / (C + 1e-30), + 'sig_x_mag': (var0 / (C + 1e-30)) * (maxg2 / (C + 1e-30)), + }) + + print(f"\n=== Solution {idx} ===") + hdr = (f"{'term':34s} {'T?':3} {'g0':>10} {'cv_sig':>9} " + f"{'ncraw_g2':>9} {'mad_med':>8} {'frac':>9} {'inv_frac':>9} " + f"{'relmag':>9} {'sig_x_mag':>10}") + print(hdr) + print('-' * len(hdr)) + for r in rows: + print(f"{r['name'][:34]:34s} {'T' if r['true'] else 'SP':3} " + f"{r['g0']:>10.3g} {r['cv_sig']:>9.2g} {r['ncraw_g2']:>9.2g} " + f"{r['mad_med']:>8.2g} {r['frac']:>9.3g} {r['inv_frac']:>9.3g} " + f"{r['relmag']:>9.3g} {r['sig_x_mag']:>10.2g}") + # per-formula separation: min over SPUR - max over TRUE (want > 0) + tv = {k: [r[k] for r in rows if r['true']] for k in + ('cv_sig', 'ncraw_g2', 'mad_med', 'inv_frac', 'relmag', 'sig_x_mag')} + sv = {k: [r[k] for r in rows if not r['true']] for k in tv} + print(" separation (min_spur - max_true, want>0):") + for k in tv: + if sv[k] and tv[k]: + sep = min(sv[k]) - max(tv[k]) + print(f" {k:12s} {sep:>11.3g}" + f" {'OK' if sep > 0 else 'x'}") + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/projects/thesis/configs/wave.yaml b/projects/thesis/configs/wave.yaml index ccbbbb14..059e91ad 100644 --- a/projects/thesis/configs/wave.yaml +++ b/projects/thesis/configs/wave.yaml @@ -4,3 +4,11 @@ name: wave truth_equations: - "0.04 * d^2u/dx1^2{power: 1.0} = d^2u/dx0^2{power: 1.0}" + +# Alternative form: degenerate quadratic shadow of the wave equation. On +# u_tt = c^2*u_xx the terms u_xx^2, u_xx*u_tt, u_tt^2 are collinear, so this is +# an exact algebraic consequence (R^2 ~= 1; the c^2=0.04 wave-speed ratio is one +# root of its 1 = 48.68 r - 591.98 r^2). EPDE discovers it instead of the PDE +# form; it is algebraically correct for the dataset and so credits as valid. +truth_alternatives: + - - "48.67869238111131 * d^2u/dx1^2{power: 1.0} * d^2u/dx0^2{power: 1.0} + -591.9797844706457 * d^2u/dx0^2{power: 2.0} = d^2u/dx1^2{power: 2.0}" diff --git a/projects/thesis/kdv_sindy_test.py b/projects/thesis/kdv_sindy_test.py new file mode 100644 index 00000000..b17a1ee5 --- /dev/null +++ b/projects/thesis/kdv_sindy_test.py @@ -0,0 +1,334 @@ +"""KdV SINDy seeded-equation diagnostic. + +Bypasses MOEA/D search entirely. Seeds the truth KdV equation directly +into the EPDE pipeline, runs ``VWSRSparsity`` + ``L2LRFitness`` on it, +and prints which terms survive the sparsity step plus the per-term +fitness / coefficient-stability / AIC metrics. + +Also evaluates a "truncated truth" variant where one term has been +dropped, so we can see by how much the metrics worsen when the +nonlinearity is missing. + +Default config matches the thesis NEW pipeline (L2LRFitness + +VWSRSparsity) on the same kdv_sindy.mat data the production search +uses, so the metrics here are directly comparable to what RPS's +term-sweep would see during the search. + +Usage: + python projects/thesis/kdv_sindy_test.py [--gram-mode axis|vcoef] + [--seed N] +""" +from __future__ import annotations + +import argparse +import os +import re +import sys + +_THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +_REPO_ROOT = os.path.abspath(os.path.join(_THIS_DIR, '..', '..')) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) +if _THIS_DIR not in sys.path: + sys.path.insert(0, _THIS_DIR) + +import numpy as np + +import epde.globals as global_var +from epde import globals as gv +from epde.interface.equation_translator import translate_equation +from epde.operators.common.fitness import L2LRFitness +from epde.operators.common.sparsity import VWSRSparsity +from epde.operators.common.coeff_calculation import LinRegBasedCoeffsEquation +from epde.operators.utils.operator_mappers import map_operator_between_levels +from epde.operators.utils.default_parameter_loader import EvolutionaryParams +from epde.operators.common.stability import ( + VaryingCoefSetup, calculate_weights, +) + +from thesis_runner import ( # noqa: E402 + _boundary_for, _build_token_pool, _configure_preprocessor, + _construct_search, _set_seeds, load_config, pipeline_settings, +) + + +# Truth equations are read from configs/.yaml -> truth_equations. +# For multi-equation systems (lorenz, lv, ns) the list has one entry per +# variable; ``translate_equation`` consumes a dict {var: symbolic}. + +# Truth YAMLs label standalone grid tokens ``x_0{...}`` / ``x_1{...}`` (legacy +# per-axis names), but the pool registers a single ``x`` family with ``dim`` as +# an inner parameter. ``translate_equation`` matches the bare label against the +# pool, so collapse ``x_N{`` -> ``x{`` first (the ``dim:N`` param already +# disambiguates the axis). Without this, ode (``x_0``) and pde_divide (``x_1``) +# fail to seed. Same normalisation as ``kdv_sindy_sweep`` (kept local to avoid a +# circular import, since that module imports from this one). +_GRID_LABEL_RE = re.compile(r'\bx_\d+(?=\{)') + + +def _normalize_grid_labels(symbolic): + """Rewrite ``x_0{...}`` -> ``x{...}`` on a string or dict-of-strings.""" + if isinstance(symbolic, dict): + return {k: _GRID_LABEL_RE.sub('x', v) for k, v in symbolic.items()} + return _GRID_LABEL_RE.sub('x', symbolic) + + +def build_pool_only(cfg, pipeline_kwargs): + """Run the thesis builder up to (but not including) ``.fit``. + + Returns a primed ``EpdeSearch`` whose pool is built and + preprocessor is configured, so ``translate_equation`` can resolve + every token. The optimizer is never instantiated, so this is fast. + """ + coords, data, variable_names, dim = cfg.load_data() + additional_tokens = _build_token_pool(cfg, coords, dim) + search = _construct_search(cfg, coords, pipeline_kwargs) + _configure_preprocessor(search, cfg) + f = cfg.hparams['fit'] + max_deriv_order = f.get('max_deriv_order') + if max_deriv_order is None: + max_deriv_order = (2,) + (4,) * dim + else: + max_deriv_order = tuple(max_deriv_order) + search.create_pool( + data=data, + variable_names=variable_names, + max_deriv_order=max_deriv_order, + additional_tokens=additional_tokens, + data_fun_pow=f['data_fun_pow'], + deriv_fun_pow=f['deriv_fun_pow'], + fourier_layers=f['fourier_layers'], + ) + return search + + +def make_fit_operator(): + """Build a gene-level ``L2LRFitness`` + ``VWSRSparsity`` chain. + + We intentionally do NOT map to chromosome level here: the test + drives the gene-level operator per-equation by hand with + ``force_out_of_place=True`` so the sparsity sub-operator actually + runs (the chromosome wrapper short-circuits sparsity when + ``fitness_calculated`` is False, leaving the seeded weights + untouched). + """ + params = EvolutionaryParams() + op_params = params.get_default_params_for_operator('DiscrepancyBasedFitnessWithCV') + fit_op = L2LRFitness(list(op_params.keys())) + fit_op.params = op_params + fit_op.set_suboperators({ + 'sparsity': VWSRSparsity(), + 'coeff_calc': LinRegBasedCoeffsEquation(), + }) + return fit_op + + +def _per_term_cv_stats(weights_arr: np.ndarray) -> dict: + """Per-feature distribution stats over a (M, n_features) per-window + OLS weight stack. Returns arrays of shape (n_features,). + """ + if weights_arr.ndim == 1: + weights_arr = weights_arr[:, None] + mu = weights_arr.mean(axis=0) + std = weights_arr.std(axis=0, ddof=1) + median = np.median(weights_arr, axis=0) + mad = np.median(np.abs(weights_arr - median), axis=0) + with np.errstate(divide='ignore', invalid='ignore'): + cv_sqr_std = np.nan_to_num((std ** 2) / (mu ** 2)) + cv_lin_std = np.nan_to_num(std / np.abs(mu)) + cv_sqr_mad = np.nan_to_num((mad ** 2) / (median ** 2)) + return { + 'mu': mu, 'std': std, 'median': median, 'mad': mad, + 'cv_sqr_std': cv_sqr_std, + 'cv_lin_std': cv_lin_std, + 'cv_sqr_mad': cv_sqr_mad, + 'M': int(weights_arr.shape[0]), + } + + +def _gram_kwargs_for_current_mode() -> tuple: + """Diagnostic re-compute of per-window weights uses the axis backup + ``GramSetup`` (``calculate_weights``'s default), so this returns + ``(None, None)``. The vcoef default does not route through + ``calculate_weights`` -- it scores via ``VaryingCoefSetup`` directly. + """ + return (None, None) + + +def inspect_truth(symbolic, search, fit_op, all_vars=('u',)): + """Seed the truth equation; for every non-target term, print the + per-window OLS coefficient distribution, all three CV formulas, and + the LASSO threshold each term would face on iteration 1. Then run + sparsity to see which actually get dropped. + + The per-term table answers the user's question: "which term is + zeroed, and why" -- the "why" is whichever has CV * max_corr > rho. + + ``symbolic`` may be either a single string (single-equation system + like KdV) or a dict mapping variable name to symbolic string + (multi-equation systems like Lorenz). + """ + metaparams = { + ('sparsity', v): {'optimizable': False, 'value': 1e-6} + for v in all_vars + } + soeq = translate_equation(symbolic, search.pool, all_vars=list(all_vars)) + for v in all_vars: + eq = soeq.vals[v] + eq.main_var_to_explain = v + eq.metaparameters = metaparams + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + + print(f'\n{"=" * 84}\nTRUTH per-term diagnostic\n{"=" * 84}') + print(f' symbolic: {symbolic}') + + for v in all_vars: + eq = soeq.vals[v] + target_idx = eq.target_idx + feat_terms = [t for i, t in enumerate(eq.structure) if i != target_idx] + target_term = eq.structure[target_idx] + print(f'\n [{v}] target term: {target_term.name}') + + # --- step 1: dump per-window OLS weights via calculate_weights + # using the same Gram class the LASSO solver will use. + _, target, features = eq.evaluate(normalize=True, return_val=False) + g_fun_vals = global_var.grid_cache.g_func[global_var.grid_cache.g_func_mask] + data_shape = global_var.grid_cache.inner_shape + gram_cls, gram_kwargs = _gram_kwargs_for_current_mode() + weights = np.array(calculate_weights( + features, target, g_fun_vals, data_shape, True, + gram_cls=gram_cls, gram_kwargs=gram_kwargs, + )) # shape (M, n_features_aug) -- last column is intercept + stats = _per_term_cv_stats(weights) + M = stats['M'] + + # --- varying-coefficient scores on the SAME (features, target), + # aligned to [feat_terms..., ]. ``fit`` is the + # unbounded 1/significance^2 in-fit Lasso driver; ``report`` is + # the bounded [0,1] Pareto-objective form. + try: + _vc = VaryingCoefSetup(features, target, g_fun_vals, + data_shape, main_var=v) + vc_score = _vc.score(None) + except Exception as _e: + vc_score = np.full(features.shape[1] + 1, np.nan) + + # --- step 2: LASSO threshold context + # PhysicsInformedLasso uses: + # X_aug = [features, ones] + # X_T_y = X_aug.T @ target + # max_corr = max(|X_T_y[active]|) + # active_thresholds[j] = active_cv[j] * max_corr + # We pick the LIVE formula (sqr(mad/median)) for the threshold so + # the user sees exactly what the LASSO solver applies. The other + # two CVs are still printed alongside for comparison. + X_aug = np.hstack([features, np.ones((features.shape[0], 1))]) + # weighted X_T_y so the threshold lines up with the + # PhysicsInformedLasso interior weighting. + X_T_y = X_aug.T @ (g_fun_vals * target if g_fun_vals.shape == target.shape + else target) + max_corr = float(np.max(np.abs(X_T_y))) + # ``rho`` on iter 1 with all-zero active_coef and residual=y: + # rho_j = X_aug[:, j] @ (g_fun_vals * y - X_aug @ 0) = X_T_y[j]. + rho_iter1 = X_T_y + + # --- step 3: per-term table + names = [t.name for t in feat_terms] + [''] + n_print = len(names) + print(f' [{v}] per-window weight stack: M={M}, n_features_aug={n_print}') + print(f' [{v}] max_corr (=max(|X^T y| over all features)): {max_corr:.6g}') + print(f' [{v}] {"#":<3} {"term":<55} ' + f'{"mu":>10} {"median":>10} {"std":>10} {"mad":>10} ' + f'{"sqr(s/mu)":>10} {"lin(s/|mu|)":>11} {"sqr(mad/med)":>12} ' + f'{"vc(var0/C)":>11} ' + f'{"thr=cv*mc":>10} {"|rho1|":>10} {"killed?":>8}') + for j, name in enumerate(names): + cv_live = stats['cv_sqr_mad'][j] + threshold_live = cv_live * max_corr + killed = abs(rho_iter1[j]) < threshold_live + vcs = vc_score[j] if j < len(vc_score) else float('nan') + print( + f' [{v}] {j:<3} {name[:55]:<55} ' + f'{stats["mu"][j]:>10.3g} {stats["median"][j]:>10.3g} ' + f'{stats["std"][j]:>10.3g} {stats["mad"][j]:>10.3g} ' + f'{stats["cv_sqr_std"][j]:>10.3g} ' + f'{stats["cv_lin_std"][j]:>11.3g} ' + f'{stats["cv_sqr_mad"][j]:>12.3g} ' + f'{vcs:>11.3g} ' + f'{threshold_live:>10.3g} {abs(rho_iter1[j]):>10.3g} ' + f'{("YES" if killed else "no"):>8}' + ) + + # --- step 4: actually run sparsity to confirm the prediction. + sweep_fitness = fit_op.apply(eq, {}, force_out_of_place=True) + print(f'\n [{v}] after VWSRSparsity.apply:') + for j, name in enumerate(names): + if j < len(eq.weights_internal): + w_post = eq.weights_internal[j] + else: + continue + mark = ' <-- DROPPED' if w_post == 0.0 else '' + print(f' {j}: {name[:55]:<55} w_internal={w_post!r}{mark}') + print(f' [{v}] sparsity-pass fitness: {sweep_fitness!r}') + # Run the fitness pass so coef_stability gets updated for ref. + eq.fitness_calculated = False + eq.stability_calculated = False + fit_op.apply(eq, {}, force_out_of_place=False) + print(f' [{v}] weights_final: {eq.weights_final}') + print(f' [{v}] fitness_value: {eq.fitness_value!r}') + print(f' [{v}] coef_stability (live formula): {eq.coefficients_stability!r}') + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--system', default='kdv', + help="System name (loads configs/.yaml). " + "Truth equations come from cfg.truth_equations.") + p.add_argument('--gram-mode', default='vcoef', + choices=('axis', 'vcoef'), + help="Gram / stability strategy (default: vcoef = " + "varying-coefficient stability). 'axis' = legacy " + "axis-aligned sliding-window backup (var/mu^2 CV).") + p.add_argument('--seed', type=int, default=0) + args = p.parse_args(argv) + + gv.set_gram_config(args.gram_mode) + cfg = load_config(args.system) + cfg.hparams['moeadd']['early_stop_on_truth'] = False + pipeline_kwargs = pipeline_settings('new') + _set_seeds(args.seed) + + print(f'{args.system} SINDy seeded-equation test') + print(f' gram_mode={args.gram_mode}') + print(f' seed={args.seed}') + + search = build_pool_only(cfg, pipeline_kwargs) + fit_op = make_fit_operator() + coords, data, variable_names, dim = cfg.load_data() + + # ``SystemCfg`` only retains canonical-token truth; the raw symbolic + # strings live in the YAML. Re-read them here so ``translate_equation`` + # can build the seeded SoEq. Multi-equation systems (lorenz, lv, ns) + # need a dict {var: symbolic} so each variable gets the right gene. + import yaml as _yaml + cfg_path = os.path.join(_THIS_DIR, 'configs', f'{args.system}.yaml') + with open(cfg_path) as f: + truth_eqs = _yaml.safe_load(f).get('truth_equations') or [] + if len(variable_names) == 1: + seeded = truth_eqs[0] + else: + seeded = {var: truth_eqs[i] for i, var in enumerate(variable_names)} + # Collapse legacy ``x_N{`` grid labels to the pool's ``x{`` so seeding + # works for systems whose truth uses standalone coordinate tokens + # (ode: x_0, pde_divide: x_1). + seeded = _normalize_grid_labels(seeded) + + inspect_truth(seeded, search, fit_op, all_vars=tuple(variable_names)) + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/projects/thesis/profile_loop_stats.py b/projects/thesis/profile_loop_stats.py index 61f7ffa6..053ab22c 100644 --- a/projects/thesis/profile_loop_stats.py +++ b/projects/thesis/profile_loop_stats.py @@ -3,7 +3,10 @@ Sets ``EPDE_LOOP_STATS=1`` in the process environment **before** any epde import, then runs ``build_search`` for each system (defaults: lv + wave), prints the loop-stats table, and writes it to -``projects/thesis/profile_results/loop_stats_.txt``. +``projects/thesis/profile_results/loop_stats_.txt``. When +more than one system is profiled in a single invocation, also dumps +a side-by-side cross-system timer comparison to +``profile_results/timer_compare.txt``. Usage: python projects/thesis/profile_loop_stats.py [system ...] [--epochs N] [--seed N] @@ -11,6 +14,7 @@ from __future__ import annotations import argparse +import json import os import sys import time @@ -37,12 +41,25 @@ def profile_system(system_name: str, pipeline: str = 'new', seed: int = 0, - epochs: int | None = None) -> None: + epochs: int | None = None, + gram_mode: str = 'vcoef') -> tuple[float, dict]: + """Profile one (system, pipeline, seed) rep. + + Returns ``(wall, timers)`` where ``wall`` is the outer wall-clock + in seconds and ``timers`` is the snapshot returned by + ``_loop_stats.timers_snapshot()`` so the caller can build a + cross-system compare table. + """ cfg = load_config(system_name) cfg.hparams['moeadd']['early_stop_on_truth'] = False if epochs is not None: cfg.hparams['moeadd']['training_epochs'] = int(epochs) + # Pin the gram config before any build_search call so all operator + # paths see consistent settings for this rep. + from epde import globals as _gv # noqa: E402 + _gv.set_gram_config(gram_mode) + pipeline_kwargs = pipeline_settings(pipeline) _set_seeds(seed) @@ -57,13 +74,108 @@ def profile_system(system_name: str, pipeline: str = 'new', seed: int = 0, wall = time.time() - t0 print(f"\n[wall] build_search total: {wall:.2f}s\n") + # Suffix output filenames with the gram mode so vcoef vs axis runs can + # co-exist. The default (vcoef) gets no suffix. + cfg_tag = "" if gram_mode == 'vcoef' else f"_{gram_mode}" + out_dir = os.path.join(_THIS_DIR, 'profile_results') os.makedirs(out_dir, exist_ok=True) - out_path = os.path.join(out_dir, f"loop_stats_{system_name}.txt") + out_path = os.path.join(out_dir, f"loop_stats_{system_name}{cfg_tag}.txt") text = _loop_stats.report(path=out_path) print(text) print(f"\n[saved] {out_path}") + snapshot = _loop_stats.timers_snapshot() + + # JSON sidecar so parallel system runs can be aggregated post-hoc + # by ``--compare-only``. One file per (system, seed, epochs, cfg) + # tuple so concurrent processes never write the same path. + json_path = os.path.join( + out_dir, + f"loop_stats_{system_name}_seed{seed}_ep{cfg.hparams['moeadd']['training_epochs']}{cfg_tag}.json" + ) + with open(json_path, 'w') as f: + json.dump({'system': system_name, 'pipeline': pipeline, + 'seed': seed, + 'epochs': cfg.hparams['moeadd']['training_epochs'], + 'gram_mode': gram_mode, + 'wall': wall, 'timers': snapshot}, f, indent=2) + print(f"[saved] {json_path}") + + return wall, snapshot + + +def write_timer_compare(walls: dict, timers: dict, out_path: str) -> None: + """Emit a side-by-side total_s per (site, system) table. + + ``walls`` maps system_name -> outer wall-clock seconds. + ``timers`` maps system_name -> ``timers_snapshot()`` dict. + Rows are sorted by the largest per-row total_s across systems + (so the heaviest sites land at the top). + """ + systems = list(timers.keys()) + all_sites = set() + for system in systems: + all_sites.update(timers[system].keys()) + + def _row_max(site: str) -> float: + return max( + (timers[system].get(site, {}).get('total_s', 0.0) for system in systems), + default=0.0, + ) + + sorted_sites = sorted(all_sites, key=_row_max, reverse=True) + + lines = [] + lines.append('TIMER COMPARE (total_s per system; share% = total_s / outer wall)') + header_cells = [f"{'site':<35}"] + for system in systems: + header_cells.append(f"{system + ' s':>14}") + header_cells.append(f"{system + ' %':>8}") + header = ' '.join(header_cells) + lines.append(header) + lines.append('-' * len(header)) + lines.append(f"{'(outer wall)':<35} " + ' '.join( + f"{walls[system]:>13.2f}s {100.0:>7.1f}%" for system in systems + )) + lines.append('-' * len(header)) + for site in sorted_sites: + row_cells = [f"{site:<35}"] + for system in systems: + t = timers[system].get(site, {}).get('total_s', 0.0) + share = (100.0 * t / walls[system]) if walls[system] > 0 else 0.0 + row_cells.append(f"{t:>13.2f}s") + row_cells.append(f"{share:>7.1f}%") + lines.append(' '.join(row_cells)) + + text = '\n'.join(lines) + '\n' + with open(out_path, 'w') as f: + f.write(text) + print(text) + print(f"[saved] {out_path}") + + +def _load_sidecar_jsons(systems: list, seed: int, epochs: int) -> tuple[dict, dict]: + """Read ``loop_stats__seed_ep.json`` sidecars + written by prior parallel runs and return ``(walls, timers)``. + Missing sidecars are reported and skipped. + """ + out_dir = os.path.join(_THIS_DIR, 'profile_results') + walls: dict[str, float] = {} + timers: dict[str, dict] = {} + for system_name in systems: + json_path = os.path.join( + out_dir, f"loop_stats_{system_name}_seed{seed}_ep{epochs}.json" + ) + if not os.path.exists(json_path): + print(f"[compare-only] missing sidecar: {json_path}") + continue + with open(json_path) as f: + payload = json.load(f) + walls[system_name] = float(payload['wall']) + timers[system_name] = payload['timers'] + return walls, timers + def main(argv=None) -> int: p = argparse.ArgumentParser(description=__doc__, @@ -72,11 +184,40 @@ def main(argv=None) -> int: p.add_argument('--pipeline', default='new', choices=('legacy', 'new')) p.add_argument('--seed', type=int, default=0) p.add_argument('--epochs', type=int, default=5) + p.add_argument('--compare-only', action='store_true', + help="Skip runs; aggregate existing sidecar JSONs into timer_compare.txt.") + p.add_argument('--gram-mode', default='vcoef', + choices=('axis', 'vcoef'), + help="Gram / stability strategy (default: vcoef = " + "varying-coefficient stability). 'axis' = legacy " + "axis-aligned sliding-window backup (var/mu^2 CV).") args = p.parse_args(argv) + if args.compare_only: + walls, timers = _load_sidecar_jsons(args.systems, args.seed, args.epochs) + if not timers: + print("[compare-only] no sidecars found; nothing to compare.") + return 1 + out_dir = os.path.join(_THIS_DIR, 'profile_results') + os.makedirs(out_dir, exist_ok=True) + compare_path = os.path.join(out_dir, 'timer_compare.txt') + write_timer_compare(walls, timers, compare_path) + return 0 + + walls: dict[str, float] = {} + timers: dict[str, dict] = {} for system_name in args.systems: - profile_system(system_name, pipeline=args.pipeline, seed=args.seed, - epochs=args.epochs) + wall, snap = profile_system(system_name, pipeline=args.pipeline, + seed=args.seed, epochs=args.epochs, + gram_mode=args.gram_mode) + walls[system_name] = wall + timers[system_name] = snap + + if len(args.systems) > 1: + out_dir = os.path.join(_THIS_DIR, 'profile_results') + os.makedirs(out_dir, exist_ok=True) + compare_path = os.path.join(out_dir, 'timer_compare.txt') + write_timer_compare(walls, timers, compare_path) return 0 diff --git a/projects/thesis/run.py b/projects/thesis/run.py index 8c211e03..7d6281b0 100644 --- a/projects/thesis/run.py +++ b/projects/thesis/run.py @@ -20,6 +20,7 @@ import argparse import os +import re import sys _THIS_DIR = os.path.dirname(os.path.abspath(__file__)) @@ -72,6 +73,27 @@ def main(argv=None) -> int: help="overwrite existing per-rep JSONs instead of skipping") parser.add_argument('--seed-base', type=int, default=0, help="seed for rep 0 (rep i uses seed_base + i)") + parser.add_argument('--epochs', type=int, default=None, + help="override moeadd.training_epochs from the YAML.") + parser.add_argument('--gram-mode', default='vcoef', + choices=('axis', 'vcoef'), + help="Gram / stability strategy (default: vcoef = " + "varying-coefficient stability, patch-free, " + "data-driven basis modes). 'axis' = legacy " + "axis-aligned sliding-window backup (var/mu^2 CV).") + parser.add_argument('--noise-level', type=float, default=0.0, + help="Additive Gaussian noise applied to every data " + "array returned by ``cfg.load_data``. Convention " + "matches PySINDy's noisy benchmark: " + "``sigma = noise_level * 0.01 * std(data)``. " + "Noise re-applied per rep with seed " + "``seed_base + rep_idx`` so reps see " + "independent realisations.") + parser.add_argument('--vc-coord-penalty', type=float, default=None, + help="kappa weight for the vcoef coordinate-modulation " + "penalty (globals.vc_coord_penalty). 0 disables; " + "larger penalises coordinate-modulated spurious " + "terms harder. Default: leave globals' value.") args = parser.parse_args(argv) try: @@ -79,6 +101,45 @@ def main(argv=None) -> int: except FileNotFoundError as exc: parser.error(str(exc)) + # Pin the gram mode before the batch starts; applies to every rep in + # run_smoke since the setting is a process-level global. + from epde import globals as _gv + _gv.set_gram_config(args.gram_mode) + if args.vc_coord_penalty is not None: + _gv.vc_coord_penalty = float(args.vc_coord_penalty) + + # When ``--noise-level`` is set, monkey-patch ``cfg.load_data`` so + # every call inside ``build_search`` injects independent Gaussian + # noise sized at PySINDy's convention. ``_gv.noise_seed`` is + # advanced by run_smoke per rep so each rep sees a fresh draw. + if args.noise_level > 0: + import numpy as _np + orig_load = cfg.load_data + nl = float(args.noise_level) + def _noisy_load(): + coords, data, vars_, dim = orig_load() + seed = getattr(_gv, 'noise_seed', None) + if seed is None: + seed = args.seed_base + rng = _np.random.default_rng(int(seed)) + def _noisy(arr): + a = _np.asarray(arr, dtype=_np.float64) + sigma = nl * 0.01 * float(_np.std(a)) + if sigma > 0: + a = a + rng.normal(0.0, sigma, size=a.shape) + return a + if isinstance(data, _np.ndarray): + noisy = _noisy(data) + elif isinstance(data, (list, tuple)): + noisy = type(data)(_noisy(a) for a in data) + else: + noisy = data + return coords, noisy, vars_, dim + cfg.load_data = _noisy_load + + if args.epochs is not None: + cfg.hparams['moeadd']['training_epochs'] = int(args.epochs) + run_smoke( cfg, reps=args.reps, diff --git a/projects/thesis/thesis_ablation_aggregate.py b/projects/thesis/thesis_ablation_aggregate.py index 78480770..6dac9bf8 100644 --- a/projects/thesis/thesis_ablation_aggregate.py +++ b/projects/thesis/thesis_ablation_aggregate.py @@ -29,13 +29,64 @@ import glob import json import os +import re import statistics import sys from collections import defaultdict _THIS_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, _THIS_DIR) -from thesis_metrics import consistency_rate, wilson_ci # noqa: E402 +from thesis_metrics import ( # noqa: E402 + coefficient_error_best, + consistency_rate, + wilson_ci, +) +_CONFIGS_DIR = os.path.join(_THIS_DIR, 'configs') +_TRUTH_CACHE: dict = {} + + +def _load_truth_eq_alts(system: str): + """Same shape as :func:`thesis_aggregate._load_truth_eq_alts`: returns + ``[primary, *alternatives]`` lists of equation strings (or an empty + list when the system has no YAML truth).""" + if system in _TRUTH_CACHE: + return _TRUTH_CACHE[system] + import yaml + path = os.path.join(_CONFIGS_DIR, f'{system}.yaml') + if not os.path.exists(path): + _TRUTH_CACHE[system] = [] + return [] + with open(path, 'r', encoding='utf-8') as fh: + cfg = yaml.safe_load(fh) or {} + primary = list(cfg.get('truth_equations') or []) + if not primary: + _TRUTH_CACHE[system] = [] + return [] + alts = [list(eqs) for eqs in (cfg.get('truth_alternatives') or []) if eqs] + result = [primary] + alts + _TRUTH_CACHE[system] = result + return result + + +_EQ_PREFIX_RE = re.compile(r'^\s*[/\\|]\s+') + + +def _rep_coef_error(rec: dict, truth_alts: list): + """Mirror of :func:`thesis_aggregate._rep_coef_error`: filter the + trailing hparams string out of ``discovered_text`` and strip the + coupled-system equation prefix (``/``, ``|``, ``\\``) before calling + the coefficient-error metric.""" + if not rec.get('structural_success') or not truth_alts: + return None + raw = rec.get('discovered_text') or [] + disc_text = [_EQ_PREFIX_RE.sub('', s).strip() + for s in raw if isinstance(s, str) and '=' in s] + if not disc_text: + return None + err = coefficient_error_best(disc_text, truth_alts) + if err != err: + return None + return err # Ordered so the report reads from "all off" to "all on" along each axis. @@ -115,9 +166,10 @@ def _mean_std(values: list): return mean, std -def _summarize_cell(reps: list) -> dict: +def _summarize_cell(reps: list, system: str = '') -> dict: if not reps: return {'n': 0} + truth_alts = _load_truth_eq_alts(system) if system else [] successes = sum(1 for r in reps if r.get('structural_success')) hammings = [r['hamming'] for r in reps if r.get('hamming') is not None] runtimes = [r['runtime_sec'] for r in reps if 'runtime_sec' in r] @@ -129,12 +181,15 @@ def _summarize_cell(reps: list) -> dict: epochs_success = [r['discovery_epoch'] for r in reps if r.get('structural_success') and r.get('discovery_epoch') is not None] + coef_errs = [e for e in (_rep_coef_error(r, truth_alts) for r in reps) + if e is not None] rate = successes / len(reps) ci = wilson_ci(successes, len(reps)) - mean_h = statistics.fmean(hammings) if hammings else float('nan') + mean_h, std_h = _mean_std(hammings) mean_t, std_t = _mean_std(runtimes) mean_npar, std_npar = _mean_std(n_paretos) mean_ep, std_ep = _mean_std(epochs_success) + mean_ce, std_ce = _mean_std(coef_errs) discovered_tokens = [json.dumps(r.get('discovered_tokens', []), sort_keys=True) for r in reps] errors = sum(1 for r in reps if 'error' in r) return { @@ -144,6 +199,7 @@ def _summarize_cell(reps: list) -> dict: 'wilson_lo': ci[0], 'wilson_hi': ci[1], 'mean_hamming': mean_h, + 'std_hamming': std_h, 'consistency': consistency_rate(discovered_tokens), 'mean_runtime_sec': mean_t, 'std_runtime_sec': std_t, @@ -151,6 +207,9 @@ def _summarize_cell(reps: list) -> dict: 'std_n_pareto': std_npar, 'mean_epoch_identified': mean_ep, 'std_epoch_identified': std_ep, + 'mean_coef_error': mean_ce, + 'std_coef_error': std_ce, + 'n_coef_error': len(coef_errs), 'errors': errors, } @@ -167,10 +226,11 @@ def _cell_axes(cell: str) -> tuple: def _format_table(summary: dict) -> str: header = ( - '| System | Cell | W | I | R | n | Success | mean H | ' - 'runtime (mean±std) | unique cands (mean±std) | epoch identified (mean±std) |' + '| System | Cell | W | I | R | n | Success | H (mean±std) | ' + 'coef err (mean±std) | runtime (mean±std) | ' + 'unique cands (mean±std) | epoch identified (mean±std) |' ) - sep = '|---|---|---|---|---|---|---|---|---|---|---|' + sep = '|---|---|---|---|---|---|---|---|---|---|---|---|' rows = [header, sep] def _check(b: bool) -> str: @@ -179,10 +239,7 @@ def _check(b: bool) -> str: def _success(c): if c['n'] == 0: return '-' - return ( - f"{c['rate']*100:.0f}% [{c['wilson_lo']*100:.0f}-{c['wilson_hi']*100:.0f}%] " - f"({c['successes']}/{c['n']})" - ) + return f"{c['rate']*100:.0f}% ({c['successes']}/{c['n']})" def _num(c, key, fmt): if c['n'] == 0: @@ -210,7 +267,8 @@ def _ms(c, mean_key, std_key, fmt, suffix=''): rows.append( f"| {system} | {cell} | {_check(w)} | {_check(i)} | {_check(r)} | " f"{c['n']} | {_success(c)} | " - f"{_num(c, 'mean_hamming', '{:.1f}')} | " + f"{_ms(c, 'mean_hamming', 'std_hamming', '{:.1f}')} | " + f"{_ms(c, 'mean_coef_error', 'std_coef_error', '{:.3f}')} | " f"{_ms(c, 'mean_runtime_sec', 'std_runtime_sec', '{:.1f}', 's')} | " f"{_ms(c, 'mean_n_pareto', 'std_n_pareto', '{:.1f}')} | " f"{_ms(c, 'mean_epoch_identified', 'std_epoch_identified', '{:.1f}')} |" @@ -264,7 +322,8 @@ def aggregate(root: str = None) -> dict: root = root or DEFAULT_RESULTS_DIR records = _load_records(root) summary = { - system: {cell: _summarize_cell(reps) for cell, reps in by_cell.items()} + system: {cell: _summarize_cell(reps, system) + for cell, reps in by_cell.items()} for system, by_cell in records.items() } return summary diff --git a/projects/thesis/thesis_aggregate.py b/projects/thesis/thesis_aggregate.py index fce2b9e2..6c59ac21 100644 --- a/projects/thesis/thesis_aggregate.py +++ b/projects/thesis/thesis_aggregate.py @@ -21,17 +21,50 @@ import glob import json import os +import re import statistics import sys from collections import defaultdict _THIS_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, _THIS_DIR) -from thesis_metrics import consistency_rate, wilson_ci # noqa: E402 +from thesis_metrics import ( # noqa: E402 + coefficient_error_best, + consistency_rate, + wilson_ci, +) PIPELINES = ('legacy', 'new') DEFAULT_RESULTS_DIR = os.path.join(_THIS_DIR, 'results') +_CONFIGS_DIR = os.path.join(_THIS_DIR, 'configs') +_TRUTH_CACHE: dict = {} + + +def _load_truth_eq_alts(system: str): + """Return ``[primary, *alternatives]`` as a list of equation-string lists. + + Each element is itself a list of equation strings -- the YAML's + ``truth_equations`` (primary) followed by every ``truth_alternatives`` + entry (each an alternative analytical form). Empty list if the + config has no ``truth_equations``. Cached per process.""" + if system in _TRUTH_CACHE: + return _TRUTH_CACHE[system] + import yaml # local: only the aggregator needs PyYAML + path = os.path.join(_CONFIGS_DIR, f'{system}.yaml') + if not os.path.exists(path): + _TRUTH_CACHE[system] = [] + return [] + with open(path, 'r', encoding='utf-8') as fh: + cfg = yaml.safe_load(fh) or {} + primary = list(cfg.get('truth_equations') or []) + if not primary: + _TRUTH_CACHE[system] = [] + return [] + alts = [list(eqs) for eqs in (cfg.get('truth_alternatives') or []) if eqs] + result = [primary] + alts + _TRUTH_CACHE[system] = result + return result def _unique_history_count(rec: dict, rep_path: str) -> int | None: @@ -108,9 +141,38 @@ def _mean_std(values: list): return mean, std -def _summarize_cell(reps: list) -> dict: +_EQ_PREFIX_RE = re.compile(r'^\s*[/\\|]\s+') + + +def _rep_coef_error(rec: dict, truth_alts: list): + """Per-rep relative coefficient error against the best-matching truth + alternative (success-only). ``None`` if the rep didn't succeed + structurally or no usable ``discovered_text`` is stored. + + ``discovered_text`` is N equation strings followed by the rep's + hparams dict serialised as a trailing string. Coupled systems pre- + pend each equation with a single box-drawing char (``/``, ``|``, + ``\\``) which would otherwise be swallowed by the term-coef parser + as an unparseable prefix and drop the leading term's coefficient + silently. Filter to eq strings and strip the prefix before invoking + the metric.""" + if not rec.get('structural_success') or not truth_alts: + return None + raw = rec.get('discovered_text') or [] + disc_text = [_EQ_PREFIX_RE.sub('', s).strip() + for s in raw if isinstance(s, str) and '=' in s] + if not disc_text: + return None + err = coefficient_error_best(disc_text, truth_alts) + if err != err: # nan + return None + return err + + +def _summarize_cell(reps: list, system: str = '') -> dict: if not reps: return {'n': 0} + truth_alts = _load_truth_eq_alts(system) if system else [] successes = sum(1 for r in reps if r.get('structural_success')) hammings = [r['hamming'] for r in reps if r.get('hamming') is not None] runtimes = [r['runtime_sec'] for r in reps if 'runtime_sec' in r] @@ -126,12 +188,17 @@ def _summarize_cell(reps: list) -> dict: epochs_success = [r['discovery_epoch'] for r in reps if r.get('structural_success') and r.get('discovery_epoch') is not None] + # Coefficient error: success-only; nonsense on a structurally-wrong + # rep because there's no canonical partner to compare against. + coef_errs = [e for e in (_rep_coef_error(r, truth_alts) for r in reps) + if e is not None] rate = successes / len(reps) ci = wilson_ci(successes, len(reps)) - mean_h = statistics.fmean(hammings) if hammings else float('nan') + mean_h, std_h = _mean_std(hammings) mean_t, std_t = _mean_std(runtimes) mean_npar, std_npar = _mean_std(n_paretos) mean_ep, std_ep = _mean_std(epochs_success) + mean_ce, std_ce = _mean_std(coef_errs) discovered_tokens = [json.dumps(r.get('discovered_tokens', []), sort_keys=True) for r in reps] errors = sum(1 for r in reps if 'error' in r) return { @@ -141,6 +208,7 @@ def _summarize_cell(reps: list) -> dict: 'wilson_lo': ci[0], 'wilson_hi': ci[1], 'mean_hamming': mean_h, + 'std_hamming': std_h, 'consistency': consistency_rate(discovered_tokens), 'mean_runtime_sec': mean_t, 'std_runtime_sec': std_t, @@ -148,16 +216,20 @@ def _summarize_cell(reps: list) -> dict: 'std_n_pareto': std_npar, 'mean_epoch_identified': mean_ep, 'std_epoch_identified': std_ep, + 'mean_coef_error': mean_ce, + 'std_coef_error': std_ce, + 'n_coef_error': len(coef_errs), 'errors': errors, } def _format_table(summary: dict) -> str: header = ( - '| System | n | Legacy success | Legacy H | ' - 'NEW success | NEW H | runtime L (mean±std) | runtime N (mean±std) |' + '| System | n | Legacy success | Legacy H (mean±std) | Legacy coef err (mean±std) | ' + 'NEW success | NEW H (mean±std) | NEW coef err (mean±std) | ' + 'runtime L (mean±std) | runtime N (mean±std) |' ) - sep = '|---|---|---|---|---|---|---|---|' + sep = '|---|---|---|---|---|---|---|---|---|---|' rows = [header, sep] for system in sorted(summary.keys()): legacy = summary[system].get('legacy', {'n': 0}) @@ -166,10 +238,7 @@ def _format_table(summary: dict) -> str: def cell_success(c): if c['n'] == 0: return '-' - return ( - f"{c['rate']*100:.0f}% [{c['wilson_lo']*100:.0f}-{c['wilson_hi']*100:.0f}%] " - f"({c['successes']}/{c['n']})" - ) + return f"{c['rate']*100:.0f}% ({c['successes']}/{c['n']})" def cell_num(c, key, fmt): if c['n'] == 0: @@ -179,7 +248,7 @@ def cell_num(c, key, fmt): return '-' return fmt.format(v) - def cell_ms(c, mean_key, std_key, fmt): + def cell_ms(c, mean_key, std_key, fmt, suffix=''): if c['n'] == 0: return '-' m = c.get(mean_key) @@ -187,15 +256,17 @@ def cell_ms(c, mean_key, std_key, fmt): if m is None or (isinstance(m, float) and m != m): return '-' if s is None or (isinstance(s, float) and s != s): - return fmt.format(m) - return f"{fmt.format(m)}±{fmt.format(s)}" + return f"{fmt.format(m)}{suffix}" + return f"{fmt.format(m)}±{fmt.format(s)}{suffix}" rows.append( f"| {system} | {max(legacy['n'], new['n'])} | " - f"{cell_success(legacy)} | {cell_num(legacy, 'mean_hamming', '{:.1f}')} | " - f"{cell_success(new)} | {cell_num(new, 'mean_hamming', '{:.1f}')} | " - f"{cell_ms(legacy, 'mean_runtime_sec', 'std_runtime_sec', '{:.1f}')}s | " - f"{cell_ms(new, 'mean_runtime_sec', 'std_runtime_sec', '{:.1f}')}s |" + f"{cell_success(legacy)} | {cell_ms(legacy, 'mean_hamming', 'std_hamming', '{:.1f}')} | " + f"{cell_ms(legacy, 'mean_coef_error', 'std_coef_error', '{:.3f}')} | " + f"{cell_success(new)} | {cell_ms(new, 'mean_hamming', 'std_hamming', '{:.1f}')} | " + f"{cell_ms(new, 'mean_coef_error', 'std_coef_error', '{:.3f}')} | " + f"{cell_ms(legacy, 'mean_runtime_sec', 'std_runtime_sec', '{:.1f}', 's')} | " + f"{cell_ms(new, 'mean_runtime_sec', 'std_runtime_sec', '{:.1f}', 's')} |" ) return '\n'.join(rows) @@ -239,7 +310,8 @@ def aggregate(root: str = None) -> dict: root = root or DEFAULT_RESULTS_DIR records = _load_records(root) summary = { - system: {pipeline: _summarize_cell(reps) for pipeline, reps in by_pipeline.items()} + system: {pipeline: _summarize_cell(reps, system) + for pipeline, reps in by_pipeline.items()} for system, by_pipeline in records.items() } return summary diff --git a/projects/thesis/thesis_metrics.py b/projects/thesis/thesis_metrics.py index c386a97f..05c15f27 100644 --- a/projects/thesis/thesis_metrics.py +++ b/projects/thesis/thesis_metrics.py @@ -271,6 +271,180 @@ def wilson_ci(successes: int, n: int, z: float = 1.96): return (max(0.0, center - half), min(1.0, center + half)) +# --------------------------------------------------------------------------- +# Coefficient-error metric +# +# Companion to the structural Hamming metric: once a rep matches the truth +# structurally (canonical equality), how close are its numerical coefficients +# to the truth coefficients? Target-flip robust: both sides are written as +# ``sum(c_i * t_i) - target = 0`` and re-normalised so the truth's target +# term has coefficient 1, then per-term relative errors are averaged. +# --------------------------------------------------------------------------- + + +def _parse_term_with_coef(term_text: str): + """Parse one ``c * f1{...} * f2{...}`` term into ``(term_canonical, coef)``. + + ``term_canonical`` is the same ``frozenset(factors)`` :func:`_parse_term` + produces. Returns ``None`` for pure-constant terms (``0.0``, + ``-0.5``) and zero-coef terms, mirroring the structural metric's + drop rule so coefficient comparison stays aligned with structure. + """ + pieces = [p.strip() for p in term_text.split('*')] + factors = [] + coef = 1.0 + coef_seen = False + for piece in pieces: + if not piece: + continue + factor = _parse_factor(piece) + if factor is None: + try: + coef *= float(piece) + coef_seen = True + except ValueError: + continue + else: + factors.append(factor) + if not factors: + return None + if coef_seen and abs(coef) < 1e-12: + return None + if not coef_seen: + coef = 1.0 + return (frozenset(factors), coef) + + +def _equation_term_coefs(eq_text: str): + """Parse ``sum_terms = target`` into ``(coef_by_term, target_key)``. + + Equation is rewritten as ``sum_terms - target = 0``; ``coef_by_term`` + holds the signed coefficient of every canonical term in that form + (target term gets ``-target_coef`` so the dict is in ``Σ c_i t_i = 0`` + form). Same factor / param canonicalisation as :func:`_parse_term`. + """ + if '=' not in eq_text: + return None + lhs, rhs = eq_text.split('=', 1) + target = _parse_term_with_coef(rhs) + if target is None: + return None + target_key, target_coef = target + coef_by_term: dict = {} + for term_text in lhs.split('+'): + parsed = _parse_term_with_coef(term_text) + if parsed is None: + continue + key, coef = parsed + coef_by_term[key] = coef_by_term.get(key, 0.0) + coef + coef_by_term[target_key] = coef_by_term.get(target_key, 0.0) - target_coef + return coef_by_term, target_key + + +def _equation_relative_coef_error(disc_eq_text: str, truth_eq_text: str) -> float: + """Mean per-term relative coefficient error between two equations. + + Both equations are written in ``Σ c_i t_i = 0`` form, anchored at the + truth's target term so both sides have anchor coef 1, then matched + term-by-term on the canonical factor set. The relative error for a + matched term ``t_i`` is ``|c_disc - c_truth| / |c_truth|``; missing + terms (in discovered or truth) contribute 1.0 each. Returns + ``float('nan')`` if either equation fails to parse or the truth's + anchor term is absent from / has zero coefficient in the discovered + equation (target-flip unresolvable). + """ + disc = _equation_term_coefs(disc_eq_text) + truth = _equation_term_coefs(truth_eq_text) + if disc is None or truth is None: + return float('nan') + disc_coefs, _ = disc + truth_coefs, anchor = truth + truth_anchor_coef = truth_coefs.get(anchor, 0.0) + disc_anchor_coef = disc_coefs.get(anchor, 0.0) + if abs(truth_anchor_coef) < 1e-12 or abs(disc_anchor_coef) < 1e-12: + return float('nan') + truth_norm = {k: v / truth_anchor_coef for k, v in truth_coefs.items()} + disc_norm = {k: v / disc_anchor_coef for k, v in disc_coefs.items()} + keys = set(truth_norm) | set(disc_norm) + errors: List[float] = [] + for k in keys: + if k == anchor: + continue # both 1.0 by construction + tc = truth_norm.get(k) + dc = disc_norm.get(k) + if tc is None: + errors.append(1.0) # extra term in discovered + continue + if dc is None: + errors.append(1.0) # missing term in discovered + continue + if abs(tc) < 1e-12: + errors.append(0.0 if abs(dc) < 1e-12 else 1.0) + continue + errors.append(abs(dc - tc) / abs(tc)) + if not errors: + return 0.0 + return sum(errors) / len(errors) + + +def _system_coef_error(discovered_eq_texts: Sequence[str], + truth_eq_texts: Sequence[str]) -> float: + """Bipartite coef-error matching between two equation systems. + + Pads the shorter side with empty equations (each empty pair scores + 1.0) and brute-forces over permutations to minimise the average + per-equation :func:`_equation_relative_coef_error`. Returns + ``float('nan')`` if every permutation contains an unparseable pair. + """ + disc = [s for s in discovered_eq_texts if isinstance(s, str) and s.strip()] + truth = [s for s in truth_eq_texts if isinstance(s, str) and s.strip()] + if not disc or not truth: + return float('nan') + n = max(len(disc), len(truth)) + pad_d = list(disc) + [''] * (n - len(disc)) + pad_t = list(truth) + [''] * (n - len(truth)) + best = float('nan') + for perm in permutations(range(n)): + total = 0.0 + valid = True + for i in range(n): + de, te = pad_d[i], pad_t[perm[i]] + if not de or not te: + total += 1.0 # unmatched eq counts as fully-wrong + continue + err = _equation_relative_coef_error(de, te) + if err != err: # nan + valid = False + break + total += err + if not valid: + continue + avg = total / n + if best != best or avg < best: + best = avg + return best + + +def coefficient_error_best(discovered_eq_texts: Sequence[str], + truth_alternatives_text_lists) -> float: + """Lowest mean coef error across all declared truth alternatives. + + ``truth_alternatives_text_lists`` is an iterable of equation-string + lists -- the primary truth followed by each alternative. The minimum + is taken across alternatives so a target-flipped / identity-based + discovery is scored against the closest valid analytical form (same + convention as :func:`hamming_best`). + """ + best = float('nan') + for truth_alt in truth_alternatives_text_lists: + err = _system_coef_error(discovered_eq_texts, truth_alt) + if err != err: + continue + if best != best or err < best: + best = err + return best + + if __name__ == '__main__': # Quick self-check: round-trip the Lorenz triple and confirm Hamming == 0 # against itself, then perturb one term and confirm Hamming == 2. @@ -372,4 +546,49 @@ def wilson_ci(successes: int, n: int, z: float = 1.96): '(expected', expected, '— missing duplicate of', same_eq[:30], ')') assert h_dup == expected, f"expected {expected}, got {h_dup}" + # Coefficient-error metric: identical equations -> 0. + truth_eq = ('10.0 * v{power: 1.0} + -10.0 * u{power: 1.0} ' + '= du/dx0{power: 1.0}') + ce_zero = _equation_relative_coef_error(truth_eq, truth_eq) + print('coef_err(identical) =', ce_zero) + assert ce_zero == 0.0, f"expected 0, got {ce_zero}" + + # 10% perturbation on one term -> 10% / 2 matched non-anchor terms = 5%. + perturbed_eq = ('11.0 * v{power: 1.0} + -10.0 * u{power: 1.0} ' + '= du/dx0{power: 1.0}') + ce_perturb = _equation_relative_coef_error(perturbed_eq, truth_eq) + print('coef_err(+10% on v term) =', ce_perturb) + assert abs(ce_perturb - 0.05) < 1e-9, f"expected 0.05, got {ce_perturb}" + + # Target-flip robustness: wave-style flip with scale factor. + wave_truth_eq = '0.04 * d^2u/dx1^2{power: 1.0} = d^2u/dx0^2{power: 1.0}' + wave_flipped_eq = '25.0 * d^2u/dx0^2{power: 1.0} = d^2u/dx1^2{power: 1.0}' + ce_wave = _equation_relative_coef_error(wave_flipped_eq, wave_truth_eq) + print('coef_err(wave 25 vs 0.04 flipped) =', ce_wave) + assert abs(ce_wave) < 1e-9, f"expected 0, got {ce_wave}" + + # System-level bipartite pairing for LV: swap the two equations' + # order on the discovered side; result must be unchanged. + lv_truth = [ + ('20.0 * u{power: 1.0} + -20.0 * u{power: 1.0} * v{power: 1.0} ' + '= du/dx0{power: 1.0}'), + ('20.0 * u{power: 1.0} * v{power: 1.0} + -20.0 * v{power: 1.0} ' + '= dv/dx0{power: 1.0}'), + ] + lv_disc_swapped = list(reversed(lv_truth)) + ce_lv = _system_coef_error(lv_disc_swapped, lv_truth) + print('coef_err(LV swapped order) =', ce_lv) + assert abs(ce_lv) < 1e-9, f"expected 0, got {ce_lv}" + + # coefficient_error_best: picks the matching alternative. + burgers_truth_alts = [ + ['-1.0 * u{power: 1.0} * du/dx1{power: 1.0} = du/dx0{power: 1.0}'], + ['1.0 * u{power: 1.0} = x{power: 1.0, dim: 1.0} * du/dx1{power: 1.0}'], + ] + # Discovered the similarity-solution identity with 5% coef drift. + disc_burgers = ['1.05 * u{power: 1.0} = x{power: 1.0, dim: 1.0} * du/dx1{power: 1.0}'] + ce_best = coefficient_error_best(disc_burgers, burgers_truth_alts) + print('coef_err_best(burgers alt) =', ce_best) + assert abs(ce_best - 0.05) < 1e-9, f"expected 0.05, got {ce_best}" + print('thesis_metrics self-check OK') diff --git a/projects/thesis/thesis_runner.py b/projects/thesis/thesis_runner.py index 5ce7deac..3d9c6cde 100644 --- a/projects/thesis/thesis_runner.py +++ b/projects/thesis/thesis_runner.py @@ -693,9 +693,13 @@ def run_smoke( """ out_root = _resolve_out_root(system_cfg, outdir) os.makedirs(out_root, exist_ok=True) + from epde import globals as _gv for pipeline in pipelines: for rep in range(reps): seed = seed_base + rep + # Per-rep noise seed for the ``--noise-level`` injection + # path (no-op when noise is disabled). + _gv.noise_seed = seed out_path = os.path.join(out_root, f"{pipeline}_rep{rep:02d}.json") if resume and os.path.exists(out_path): try: diff --git a/projects/thesis/vcoef_margin.py b/projects/thesis/vcoef_margin.py new file mode 100644 index 00000000..d090b517 --- /dev/null +++ b/projects/thesis/vcoef_margin.py @@ -0,0 +1,197 @@ +"""True-vs-spurious stability MARGIN harness for ``gram_mode='vcoef'``. + +Bypasses MOEA/D. For each system: seed the truth equation, pull its true +feature columns + target, then inject *spurious* feature columns built as +``true_feature * coordinate_ramp`` (a legitimate library product -- feature +times a coordinate token -- whose best constant coefficient is +region-dependent, exactly the kind of term a search wrongly proposes). A +good stability estimator scores the TRUE terms low and the SPURIOUS terms +high. + +Reported per system: +* ``vc_true_max`` -- max varying-coefficient score over the true terms. +* ``vc_spur_min`` -- min vc score over the injected spurious terms. +* ``vc_margin`` -- ``vc_spur_min - vc_true_max`` (want > 0 on all systems). +* the patch-CV counterpart (``mad_median``, the current default) for A/B. + +The margin being POSITIVE on every system -- with a single fixed config and +no per-dataset tuning -- is the dataset-independence the design targets. + +Usage: + python projects/thesis/vcoef_margin.py [--systems ode,lv,wave,kdv,...] + [--axis 0] [--ramp-strength 1.0] +""" +from __future__ import annotations + +import argparse +import os +import sys + +_THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +_REPO_ROOT = os.path.abspath(os.path.join(_THIS_DIR, '..', '..')) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) +if _THIS_DIR not in sys.path: + sys.path.insert(0, _THIS_DIR) + +import numpy as np +import yaml + +import epde.globals as global_var +from epde.interface.equation_translator import translate_equation +from epde.operators.common.stability import VaryingCoefSetup, calculate_weights +from kdv_sindy_test import build_pool_only, make_fit_operator # noqa: E402 +from kdv_sindy_sweep import _normalize_grid_labels # noqa: E402 +from thesis_runner import _set_seeds, load_config, pipeline_settings # noqa: E402 + +_ALL = ['ode', 'lv', 'vdp', 'lorenz', 'kdv', 'kdv_cossin', 'wave', + 'burgers_viscous', 'burgers_inviscid', 'pde_divide', 'pde_compound', + 'ac', 'ks', 'ns'] + + +def _coord_ramps(grid_shape, axis): + """Normalised coordinate ramp(s) in [-1, 1] flattened over the grid. + + Returns one ramp per requested axis (``axis=-1`` -> every axis), each a + length-N vector matching the C-order flatten of the feature columns. + """ + D = len(grid_shape) + axes = range(D) if axis < 0 else [axis] + ramps = {} + for d in axes: + if d >= D: + continue + n_d = grid_shape[d] + line = np.linspace(-1.0, 1.0, n_d) + shp = [1] * D + shp[d] = n_d + ramps[d] = np.broadcast_to(line.reshape(shp), grid_shape).reshape(-1) + return ramps + + +def _patch_cv_scores(features, target, sw, grid_shape): + """Per-feature patch-CV (mad_median, the current default) for A/B.""" + try: + weights = np.array(calculate_weights(features, target, sw, grid_shape, + True)) + center = np.median(weights, axis=0) + mad = np.median(np.abs(weights - center), axis=0) + with np.errstate(divide='ignore', invalid='ignore'): + cv = np.nan_to_num((mad ** 2) / (center ** 2)) + return cv[:-1] # drop intercept column + except Exception: + return None + + +def margin_for_system(system, axis, ramp_strength): + cfg = load_config(system) + pipeline_kwargs = pipeline_settings('new') + _set_seeds(0) + search = build_pool_only(cfg, pipeline_kwargs) + + cfg_path = os.path.join(_THIS_DIR, 'configs', f'{system}.yaml') + with open(cfg_path) as fh: + truth_eqs = yaml.safe_load(fh).get('truth_equations') or [] + _, _, variable_names, _ = cfg.load_data() + all_vars = list(variable_names) + if len(all_vars) == 1: + seeded = truth_eqs[0] + else: + seeded = {var: truth_eqs[i] for i, var in enumerate(all_vars)} + seeded = _normalize_grid_labels(seeded) + metaparams = {('sparsity', v): {'optimizable': False, 'value': 1e-6} + for v in all_vars} + soeq = translate_equation(seeded, search.pool, all_vars=all_vars) + + grid_shape = global_var.grid_cache.inner_shape + sw = global_var.grid_cache.g_func[global_var.grid_cache.g_func_mask] + ramps = _coord_ramps(grid_shape, axis) + + rows = [] + for v in all_vars: + eq = soeq.vals[v] + eq.main_var_to_explain = v + eq.metaparameters = metaparams + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + _, target, features = eq.evaluate(normalize=True, return_val=False) + if features is None or features.ndim != 2 or features.shape[1] == 0: + continue + n_true = features.shape[1] + + # Inject spurious columns: each true feature * each coordinate ramp. + spur_cols = [] + for d, ramp in ramps.items(): + spur_cols.append(features * (ramp_strength * ramp)[:, None]) + spur = np.hstack(spur_cols) if spur_cols else np.zeros((len(target), 0)) + n_spur = spur.shape[1] + aug = np.hstack([features, spur]) + + vc = VaryingCoefSetup(aug, target, sw, grid_shape, + main_var=v).score(None) + vc_true = vc[:n_true] + vc_spur = vc[n_true:n_true + n_spur] + + pcv = _patch_cv_scores(aug, target, sw, grid_shape) + if pcv is not None and len(pcv) >= n_true + n_spur: + pcv_true, pcv_spur = pcv[:n_true], pcv[n_true:n_true + n_spur] + else: + pcv_true = pcv_spur = None + + rows.append(dict( + var=v, n_true=n_true, n_spur=n_spur, + vc_true_max=float(np.max(vc_true)), + vc_spur_min=float(np.min(vc_spur)) if n_spur else float('nan'), + vc_margin=(float(np.min(vc_spur) - np.max(vc_true)) + if n_spur else float('nan')), + pcv_true_max=(float(np.max(pcv_true)) + if pcv_true is not None else float('nan')), + pcv_spur_min=(float(np.min(pcv_spur)) + if pcv_spur is not None and len(pcv_spur) + else float('nan')), + pcv_margin=(float(np.min(pcv_spur) - np.max(pcv_true)) + if pcv_true is not None and pcv_spur is not None + and len(pcv_spur) else float('nan')), + )) + return rows + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--systems', default=','.join(_ALL), + help="Comma-separated system list (default: all 14).") + p.add_argument('--axis', type=int, default=-1, + help="Coordinate axis for the spurious ramp; -1 = every axis.") + p.add_argument('--ramp-strength', type=float, default=1.0) + args = p.parse_args(argv) + + systems = [s.strip() for s in args.systems.split(',') if s.strip()] + hdr = (f"{'system':<16} {'var':<4} {'nT':>3} {'nS':>3} " + f"{'vc_trueMax':>11} {'vc_spurMin':>11} {'vc_MARGIN':>11} " + f"{'pcv_trueMax':>11} {'pcv_spurMin':>11} {'pcv_MARGIN':>11}") + print(hdr) + print('-' * len(hdr)) + vc_pos = vc_tot = 0 + for s in systems: + try: + rows = margin_for_system(s, args.axis, args.ramp_strength) + except Exception as e: + print(f"{s:<16} ERROR: {type(e).__name__}: {e}") + continue + for r in rows: + vc_tot += 1 + vc_pos += int(r['vc_margin'] > 0) + print(f"{s:<16} {r['var']:<4} {r['n_true']:>3} {r['n_spur']:>3} " + f"{r['vc_true_max']:>11.4g} {r['vc_spur_min']:>11.4g} " + f"{r['vc_margin']:>11.4g} " + f"{r['pcv_true_max']:>11.4g} {r['pcv_spur_min']:>11.4g} " + f"{r['pcv_margin']:>11.4g}") + print('-' * len(hdr)) + print(f"vcoef positive margin: {vc_pos}/{vc_tot} equations") + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/projects/thesis/vcoef_stat_compare.py b/projects/thesis/vcoef_stat_compare.py new file mode 100644 index 00000000..5a6ec1ed --- /dev/null +++ b/projects/thesis/vcoef_stat_compare.py @@ -0,0 +1,714 @@ +"""Compare candidate vcoef CV statistics on the 14 systems: TRUE vs SPURIOUS. + +We are collapsing the vcoef estimator to a SINGLE per-term statistic ``S`` that +drives both the in-fit Lasso pruning (threshold ``S_j * max_corr``) and the +MOEA/D stability objective (now ``sum_j S_j``). This script decides which ``S`` +to use by measuring, on seeded truth across all 14 systems, whether each +candidate scores TRUE terms LOW and SPURIOUS confuser terms HIGH. + +Per system x variable: + * seed the truth equation, evaluate -> truth feature columns + target; + * fit ONE ``VaryingCoefSetup`` on the truth-only design -> per-true-term + statistics (``sum_true`` = the objective the true equation gets; + ``max_true`` = worst single true term); + * for each curated SPURIOUS confuser, fit ``[truth | one spurious]`` (mirrors + RPS adding one term at a time) and read the spurious column's statistics + (``min_spur`` = best-separated confuser). + +Candidate statistics (all from the SAME gammas, directly comparable; this script +builds ``VaryingCoefSetup`` DIRECTLY so ``_Bvals`` is present and the robust MAD +forms -- which the live ``from_full`` path cannot currently compute -- ARE +available here): + + var0_C = Var(g0)/g0^2 significance (current 'fit') + NCdeb_C = sum(max(g_k^2 - var_k, 0))/g0^2 noise-debiased region-variation + NCraw_C = sum(g_k^2)/g0^2 raw region-variation (==(std/mu)^2) + mad_med = mad_x(beta)/|median_x(beta)| robust region-variation + mad_med_sq = (mad/median)^2 + +A GOOD single ``S`` gives LOW values to true terms and HIGH to spurious, so +``separation = min_spur - max_true`` should be > 0. The discriminators are the +collinear systems (lorenz, ac) and the soliton system (kdv). +""" +from __future__ import annotations +import argparse +import os +import sys + +_THIS = os.path.dirname(os.path.abspath(__file__)) +_ROOT = os.path.abspath(os.path.join(_THIS, '..', '..')) +for _p in (_ROOT, _THIS): + if _p not in sys.path: + sys.path.insert(0, _p) + +import numpy as np +import yaml + +import epde.globals as gv +from epde.interface.equation_translator import translate_equation, parse_factor +from epde.structure.main_structures import Term +from epde.operators.common.stability import VaryingCoefSetup, resolve_vc_modes_from_input +from kdv_sindy_test import build_pool_only, _normalize_grid_labels +from thesis_runner import load_config, pipeline_settings, _set_seeds + +_ALL = ['ode', 'lv', 'vdp', 'lorenz', 'kdv', 'kdv_cossin', 'wave', + 'burgers_viscous', 'burgers_inviscid', 'pde_divide', 'pde_compound', + 'ac', 'ks', 'ns'] + +_EPS = 1e-12 + +# Candidate statistics, in the order tabulated. ``std_mu`` is printed once as a +# redundancy check that it matches sqrt(NCraw_C) (Parseval). +STAT_KEYS = ['var0_C', 'NCdeb_C', 'NCraw_C', 'mad_med', 'mad_med_sq'] + +# Curated SPURIOUS confuser terms per system, keyed by the equation's variable. +# Uses the EXACT token syntax from configs/.yaml (grid labels already +# collapsed x_N{ -> x{ by _normalize_grid_labels). Terms that fail to build, +# evaluate degenerately, or coincide with a truth term are skipped at runtime, +# so the lists are deliberately generous. +_SPURIOUS = { + 'ode': {'u': [ + 'u{power: 2.0}', + 'du/dx0{power: 1.0}', + 'x{power: 2.0, dim: 0.0}', + 'u{power: 1.0} * du/dx0{power: 1.0}', + 'sin{power: 1.0, freq: 2.0, dim: 0.0}', + ]}, + 'lv': { + 'u': ['v{power: 1.0}', 'u{power: 2.0}', 'v{power: 2.0}'], + 'v': ['u{power: 1.0}', 'u{power: 2.0}', 'v{power: 2.0}'], + }, + 'vdp': {'u': [ + 'u{power: 2.0}', 'u{power: 3.0}', 'du/dx0{power: 2.0}', + 'u{power: 1.0} * du/dx0{power: 1.0}', + ]}, + 'lorenz': { + 'u': ['w{power: 1.0}', 'u{power: 1.0} * v{power: 1.0}', + 'u{power: 1.0} * w{power: 1.0}'], + 'v': ['w{power: 1.0}', 'v{power: 1.0} * w{power: 1.0}', + 'u{power: 1.0} * v{power: 1.0}'], + 'w': ['u{power: 1.0}', 'v{power: 1.0}', + 'u{power: 1.0} * w{power: 1.0}'], + }, + 'kdv': {'u': [ + 'u{power: 1.0}', 'du/dx1{power: 1.0}', 'd^2u/dx1^2{power: 1.0}', + 'u{power: 2.0}', 'du/dx1{power: 2.0}', + ]}, + 'kdv_cossin': {'u': [ + 'u{power: 1.0}', 'du/dx1{power: 1.0}', 'd^2u/dx1^2{power: 1.0}', + 'u{power: 2.0}', + ]}, + 'wave': {'u': [ + 'u{power: 1.0}', 'du/dx1{power: 1.0}', 'du/dx0{power: 1.0}', + ]}, + 'burgers_viscous': {'u': [ + 'u{power: 1.0}', 'du/dx1{power: 1.0}', 'du/dx1{power: 2.0}', + 'u{power: 2.0}', 'd^3u/dx1^3{power: 1.0}', + ]}, + 'burgers_inviscid': {'u': [ + 'u{power: 1.0}', 'du/dx1{power: 1.0}', 'd^2u/dx1^2{power: 1.0}', + 'u{power: 2.0}', + ]}, + 'pde_divide': {'u': [ + 'u{power: 1.0}', 'd^2u/dx1^2{power: 1.0}', + 'du/dx1{power: 1.0} * x{power: 1.0, dim: 1.0}', + 'u{power: 1.0} * x{power: 1.0, dim: 1.0}', + ]}, + 'pde_compound': {'u': [ + 'u{power: 1.0}', 'du/dx1{power: 1.0}', 'd^2u/dx1^2{power: 1.0}', + 'u{power: 2.0}', + ]}, + 'ac': {'u': [ + 'du/dx1{power: 1.0}', 'u{power: 2.0}', 'du/dx1{power: 2.0}', + 'd^3u/dx1^3{power: 1.0}', + ]}, + 'ks': {'u': [ + 'u{power: 1.0}', 'du/dx1{power: 1.0}', 'd^3u/dx1^3{power: 1.0}', + 'u{power: 2.0}', + ]}, + 'ns': { + 'u': ['u{power: 1.0}', 'du/dx1{power: 1.0}', 'du/dx2{power: 1.0}', + 'v{power: 1.0} * du/dx2{power: 1.0}'], + 'v': ['v{power: 1.0}', 'dv/dx1{power: 1.0}', 'dv/dx2{power: 1.0}', + 'u{power: 1.0} * dv/dx1{power: 1.0}'], + 'p': ['du/dx1{power: 1.0}', 'dv/dx2{power: 1.0}', 'u{power: 1.0}', + 'v{power: 1.0}'], + }, +} + +# Coordinate-MODULATED confusers -- the terms that actually survive in the +# 30x14 benchmark's coordinate-degeneracy failures. Kept SEPARATE from +# ``_SPURIOUS`` (plain confusers, caught by the significance A-term) because +# these need a non-constant ``beta(x)`` to be flagged at all: they are the +# B-term / mode-resolution question the K-sweep exists to answer. Axis +# convention: x0=t=dim0, x1=x=dim1. Focus systems only. +_COORD_SPURIOUS = { + 'ode': {'u': [ + 'd^2u/dx0^2{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 0.0}', + ]}, + 'ac': {'u': [ + 'd^2u/dx1^2{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 1.0}', + 'd^4u/dx1^4{power: 1.0} * cos{power: 1.0, freq: 2.0, dim: 1.0}', + 'd^2u/dx1^2{power: 1.0} * x{power: 1.0, dim: 1.0}', + ]}, + 'wave': {'u': [ + 'd^2u/dx1^2{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 1.0}', + ]}, + 'burgers_inviscid': {'u': [ + 'du/dx1{power: 1.0} * x{power: 1.0, dim: 1.0}', + 'du/dx1{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 1.0}', + ]}, +} + +# REPLACEMENT specs: (distinctive label of the truth term to REMOVE, the +# coordinate-modulated term that replaces it). This is the LOAD-BEARING +# coordinate-degeneracy test -- the actual benchmark failure mode. Unlike the +# _COORD_SPURIOUS 'add' arm (confuser appended to COMPLETE truth, where it is +# redundant so the significance A-term flags it trivially), here the true term +# is GONE and its modulated twin must carry the signal (beta ~ 1/modulation). +# Whether that modulated term reads as unstable IS the K-sensitive question. +# The label is matched (space-insensitive substring, must be unique) against +# the truth feature-term names. +_REPLACE = { + 'ac': {'u': [ + ('d^2u/dx1^2', 'd^2u/dx1^2{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 1.0}'), + ('d^2u/dx1^2', 'd^2u/dx1^2{power: 1.0} * cos{power: 1.0, freq: 2.0, dim: 1.0}'), + ('d^2u/dx1^2', 'd^2u/dx1^2{power: 1.0} * x{power: 1.0, dim: 1.0}'), + ]}, + 'wave': {'u': [ + ('d^2u/dx1^2', 'd^2u/dx1^2{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 1.0}'), + ]}, + 'burgers_inviscid': {'u': [ + ('du/dx1', 'du/dx1{power: 1.0} * x{power: 1.0, dim: 1.0}'), + ('du/dx1', 'du/dx1{power: 1.0} * sin{power: 1.0, freq: 2.0, dim: 1.0}'), + ]}, +} + +# K (modes-per-axis incl. constant slot; K=2 -> 1 cosine, K=16 -> 15 cosines) +# swept by ``--k-sweep``. Spans the microscale's [2,6] band and well beyond. +_KSWEEP_DEFAULT = [2, 3, 4, 6, 10, 16] + + +def candidate_stats(setup: VaryingCoefSetup) -> dict: + """All candidate statistics for every active feature of ``setup``, from a + SINGLE gamma solve. Returns dict {stat_key: array(nf)} where nf includes the + internal intercept column (last).""" + sol = setup._solve_gammas(None) + if sol is None: + return None + gamma, var = sol['gamma'], sol['var'] + nf, B, mk = sol['nf'], sol['B'], sol['mk'] + Bvals = getattr(setup, '_Bvals', None) + is_const = mk == 0 + nonconst = ~is_const + out = {k: np.full(nf, np.nan) for k in + ['var0_C', 'NCdeb_C', 'NCraw_C', 'mad_med', 'mad_med_sq', 'std_mu']} + for i in range(nf): + sl = slice(i * B, (i + 1) * B) + g = gamma[sl] + v = var[sl] + g0 = float(g[is_const][0]) if np.any(is_const) else 0.0 + var0 = float(v[is_const][0]) if np.any(is_const) else 0.0 + C = g0 * g0 + nc_raw = float(np.sum(g[nonconst] ** 2)) + nc_deb = float(np.sum(np.maximum(g[nonconst] ** 2 - v[nonconst], 0.0))) + out['var0_C'][i] = var0 / (C + 1e-30) + out['NCdeb_C'][i] = nc_deb / (C + 1e-30) + out['NCraw_C'][i] = nc_raw / (C + 1e-30) + if Bvals is not None: + beta = Bvals @ g + mu = float(np.mean(beta)) + sd = float(np.std(beta)) + med = float(np.median(beta)) + mad = float(np.median(np.abs(beta - med))) + out['std_mu'][i] = sd / (abs(mu) + _EPS) + out['mad_med'][i] = mad / (abs(med) + _EPS) + out['mad_med_sq'][i] = (mad / (abs(med) + _EPS)) ** 2 + return {k: np.nan_to_num(v, nan=0.0, posinf=1e30) for k, v in out.items()} + + +def build_term_values(term_str, pool, all_vars): + """Build a single Term from a symbolic string and evaluate it on the grid. + Returns (term, flat_values) or raises.""" + factors = [parse_factor(f.strip(), pool, all_vars) + for f in term_str.split(' * ')] + term = Term(pool, passed_term=factors, collapse_powers=False) + vals = np.asarray(term.evaluate(False), dtype=float).reshape(-1) + return term, vals + + +def analyse_system(system): + """Return {var: {'truth': stats_dict, 'spur': [stats_dict, ...], + 'n_truth': int}} for one system, or raise on a load failure.""" + cfg = load_config(system) + _set_seeds(0) + search = build_pool_only(cfg, pipeline_settings('new')) + coords, data, variable_names, dim = cfg.load_data() + truth = yaml.safe_load(open(os.path.join(_THIS, 'configs', f'{system}.yaml'))) + truth_eqs = truth.get('truth_equations') or [] + all_vars = list(variable_names) + seeded = (truth_eqs[0] if len(all_vars) == 1 + else {v: truth_eqs[i] for i, v in enumerate(all_vars)}) + seeded = _normalize_grid_labels(seeded) + soeq = translate_equation(seeded, search.pool, all_vars=all_vars) + + sw = gv.grid_cache.g_func[gv.grid_cache.g_func_mask] + gshape = gv.grid_cache.inner_shape + spur_cfg = _SPURIOUS.get(system, {}) + + result = {} + for v in all_vars: + eq = soeq.vals[v] + eq.main_var_to_explain = v + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + _, target, features = eq.evaluate(normalize=True, return_val=False) + if features is None or np.asarray(features).ndim != 2: + continue + features = np.asarray(features, dtype=float) + target = np.asarray(target, dtype=float).reshape(-1) + n_truth = features.shape[1] + truth_names = {t.name for i, t in enumerate(eq.structure) + if i != eq.target_idx} + + # truth-only fit -> per-true-term statistics (cols [0:n_truth]; the + # last column is the internal intercept and is ignored). + base = candidate_stats( + VaryingCoefSetup(features, target, sw, gshape, main_var=v)) + if base is None: + continue + truth_stats = {k: base[k][:n_truth] for k in base} + + # each curated spurious confuser -> [truth | spurious] fit; read the + # spurious column (index n_truth). + spur_stats = [] + for s in spur_cfg.get(v, []): + try: + term, vals = build_term_values(s, search.pool, all_vars) + except Exception: + continue + if term.name in truth_names: + continue + if vals.shape[0] != target.shape[0] or not np.any(np.abs(vals) > 0): + continue + aug = np.hstack([features, vals[:, None]]) + st = candidate_stats( + VaryingCoefSetup(aug, target, sw, gshape, main_var=v)) + if st is None: + continue + spur_stats.append({'name': term.name, + **{k: float(st[k][n_truth]) for k in base}}) + result[v] = {'truth': truth_stats, 'spur': spur_stats, + 'n_truth': n_truth} + return result + + +def _agg_system(res): + """Aggregate per-equation results into per-stat (sum_true, max_true, + min_spur, sep) for one system. sep is the WORST (min) per-equation + separation.""" + agg = {} + for k in STAT_KEYS: + sum_true_eq, max_true_eq, sep_eq, min_spur_all = [], [], [], [] + for v, d in res.items(): + tv = d['truth'][k] + if tv.size == 0: + continue + sum_true_eq.append(float(np.sum(tv))) + mt = float(np.max(tv)) + max_true_eq.append(mt) + spur_vals = [s[k] for s in d['spur']] + if spur_vals: + ms = float(np.min(spur_vals)) + min_spur_all.append(ms) + sep_eq.append(ms - mt) + agg[k] = { + 'sum_true': float(np.sum(sum_true_eq)) if sum_true_eq else np.nan, + 'max_true': float(np.max(max_true_eq)) if max_true_eq else np.nan, + 'min_spur': float(np.min(min_spur_all)) if min_spur_all else np.nan, + 'sep': float(np.min(sep_eq)) if sep_eq else np.nan, + } + return agg + + +def dump_system(system): + """Raw per-true-term beta(x) diagnostics: does gamma_0 recover the true + constant coefficient, and how wildly does beta(x) swing? Verifies the beta + reconstruction (mad/median) on a clean (ode) vs collinear (lorenz) system.""" + cfg = load_config(system) + _set_seeds(0) + search = build_pool_only(cfg, pipeline_settings('new')) + coords, data, variable_names, dim = cfg.load_data() + truth = yaml.safe_load(open(os.path.join(_THIS, 'configs', f'{system}.yaml'))) + truth_eqs = truth.get('truth_equations') or [] + all_vars = list(variable_names) + seeded = (truth_eqs[0] if len(all_vars) == 1 + else {v: truth_eqs[i] for i, v in enumerate(all_vars)}) + seeded = _normalize_grid_labels(seeded) + soeq = translate_equation(seeded, search.pool, all_vars=all_vars) + sw = gv.grid_cache.g_func[gv.grid_cache.g_func_mask] + gshape = gv.grid_cache.inner_shape + print(f'\n##### {system} grid_shape={gshape} truth: {seeded}') + for v in all_vars: + eq = soeq.vals[v] + eq.main_var_to_explain = v + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + _, target, features = eq.evaluate(normalize=True, return_val=False) + if features is None or np.asarray(features).ndim != 2: + continue + features = np.asarray(features, dtype=float) + target = np.asarray(target, dtype=float).reshape(-1) + n_truth = features.shape[1] + names = [t.name for i, t in enumerate(eq.structure) + if i != eq.target_idx] + setup = VaryingCoefSetup(features, target, sw, gshape, main_var=v) + sol = setup._solve_gammas(None) + gamma, varr, B, mk = sol['gamma'], sol['var'], sol['B'], sol['mk'] + Bvals = setup._Bvals + is_const = mk == 0 + nonconst = ~is_const + # Constant-only weighted OLS (the standard SINDy coefficient): the + # const columns of the super-Gram are at i*B for each feature i. + const_global = np.arange(setup.n_features) * B + Ac = setup.G[np.ix_(const_global, const_global)] + bc = setup.Phiy[const_global] + try: + coef_c = np.linalg.solve(Ac + 1e-12 * np.eye(len(const_global)), bc) + except np.linalg.LinAlgError: + coef_c = np.full(len(const_global), np.nan) + # conditioning of the constant block (features-only OLS): + try: + cond_c = float(np.linalg.cond(Ac)) + except np.linalg.LinAlgError: + cond_c = float('inf') + print(f' [{v}] B={B} modes(k present)={sorted(set(mk.tolist()))} ' + f'N={Bvals.shape[0]} target|mean|={np.mean(np.abs(target)):.3g} ' + f'cond(const-block)={cond_c:.3g}') + for i in range(n_truth): + sl = slice(i * B, (i + 1) * B) + g = gamma[sl] + g0 = float(g[is_const][0]) + var0 = float(varr[sl][is_const][0]) + beta = Bvals @ g + med = float(np.median(beta)) + mad = float(np.median(np.abs(beta - med))) + vc = var0 / (g0 ** 2 + 1e-30) + print(f' {names[i][:30]:30s} g0={g0:>9.4g} var0={var0:>9.3g} ' + f'vc=var0/g0^2={vc:>9.4g} ' + f'std={np.std(beta):>8.3g} NCraw={float(np.sum(g[nonconst]**2)):.3g}') + print(f' gamma_k(nonconst)={np.round(g[nonconst], 3).tolist()}') + + +def k_sweep_system(system, K_list): + """Sweep the per-axis basis resolution ``K`` (overriding the Taylor + microscale via ``modes=(K,)*D``) and measure, per K, whether the + coordinate-modulated confusers become separable from the true terms. + + Returns ``{var: {'rows': [...], 'kstar': tuple|None, 'n_coord': int, + 'n_plain': int, 'grid': tuple}}``. Each row is one K with the B-term + (region-variation ``NCdeb_C``) and full-score (``sum_asym`` = A+B) + separations, plus the worst (min-B) coordinate confuser's beta-field + spread ``std/|mu|`` -- the visual of whether its coefficient field starts + to vary as K grows. + """ + cfg = load_config(system) + _set_seeds(0) + search = build_pool_only(cfg, pipeline_settings('new')) + coords, data, variable_names, dim = cfg.load_data() + truth = yaml.safe_load(open(os.path.join(_THIS, 'configs', f'{system}.yaml'))) + truth_eqs = truth.get('truth_equations') or [] + all_vars = list(variable_names) + seeded = (truth_eqs[0] if len(all_vars) == 1 + else {v: truth_eqs[i] for i, v in enumerate(all_vars)}) + seeded = _normalize_grid_labels(seeded) + soeq = translate_equation(seeded, search.pool, all_vars=all_vars) + + sw = gv.grid_cache.g_func[gv.grid_cache.g_func_mask] + gshape = gv.grid_cache.inner_shape + D = len(gshape) + plain_cfg = _SPURIOUS.get(system, {}) + coord_cfg = _COORD_SPURIOUS.get(system, {}) + + result = {} + for v in all_vars: + eq = soeq.vals[v] + eq.main_var_to_explain = v + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + _, target, features = eq.evaluate(normalize=True, return_val=False) + if features is None or np.asarray(features).ndim != 2: + continue + features = np.asarray(features, dtype=float) + target = np.asarray(target, dtype=float).reshape(-1) + n_truth = features.shape[1] + truth_names = {t.name for i, t in enumerate(eq.structure) + if i != eq.target_idx} + names = [t.name for i, t in enumerate(eq.structure) + if i != eq.target_idx] # ordered, aligned to features columns + + # K* the microscale would resolve (production reference line). + try: + kstar = resolve_vc_modes_from_input(gshape, main_var=v, k_max=6) + except Exception: + kstar = None + + # Build confuser value columns ONCE (independent of K); skip the ones + # that fail to build / are degenerate / collide with a truth term. + def build_cols(strs): + cols = [] + for s in strs: + try: + term, vals = build_term_values(s, search.pool, all_vars) + except Exception: + continue + if term.name in truth_names: + continue + if vals.shape[0] != target.shape[0] or not np.any(np.abs(vals) > 0): + continue + cols.append((term.name, vals)) + return cols + coord_cols = build_cols(coord_cfg.get(v, [])) + plain_cols = build_cols(plain_cfg.get(v, [])) + + def spur_rows(cols, modes): + res = [] + for name, vals in cols: + aug = np.hstack([features, vals[:, None]]) + st = candidate_stats(VaryingCoefSetup( + aug, target, sw, gshape, main_var=v, modes=modes)) + if st is None: + continue + res.append({ + 'name': name, + 'B': float(st['NCdeb_C'][n_truth]), + 'S': float(st['var0_C'][n_truth] + st['NCdeb_C'][n_truth]), + 'stdmu': float(st['std_mu'][n_truth]), + }) + return res + + rows = [] + for K in K_list: + modes = (int(K),) * D + base = candidate_stats(VaryingCoefSetup( + features, target, sw, gshape, main_var=v, modes=modes)) + if base is None: + continue + true_B = base['NCdeb_C'][:n_truth] + true_S = base['var0_C'][:n_truth] + base['NCdeb_C'][:n_truth] + max_true_B = float(np.max(true_B)) if true_B.size else float('nan') + max_true_S = float(np.max(true_S)) if true_S.size else float('nan') + + cr = spur_rows(coord_cols, modes) + pr = spur_rows(plain_cols, modes) + min_coord = min(cr, key=lambda r: r['B']) if cr else None + all_S = [r['S'] for r in cr + pr] + + rows.append({ + 'K': int(K), + 'max_true_B': max_true_B, + 'min_spur_B': (min_coord['B'] if min_coord else float('nan')), + 'sep_B': ((min_coord['B'] - max_true_B) if min_coord + else float('nan')), + 'max_true_S': max_true_S, + 'min_spur_S': (min(all_S) if all_S else float('nan')), + 'sep_S': ((min(all_S) - max_true_S) if all_S + else float('nan')), + 'worst_coord': (min_coord['name'] if min_coord else '-'), + 'worst_stdmu': (min_coord['stdmu'] if min_coord + else float('nan')), + }) + # --- REPLACEMENT arm: drop a true term, let its modulated twin carry + # the load, and score the twin vs the remaining true terms per K. + def _norm(s): + return ''.join(s.split()) + + repl_out = [] + for label, confuser_str in _REPLACE.get(system, {}).get(v, []): + nl = _norm(label) + hits = [j for j, nm in enumerate(names) if nl in _norm(nm)] + if len(hits) != 1: + continue # ambiguous / absent -> skip + j = hits[0] + try: + cterm, cvals = build_term_values(confuser_str, search.pool, all_vars) + except Exception: + continue + if cvals.shape[0] != target.shape[0] or not np.any(np.abs(cvals) > 0): + continue + feats_rm = np.delete(features, j, axis=1) + nrem = feats_rm.shape[1] + rrows = [] + for K in K_list: + modes = (int(K),) * D + aug = np.hstack([feats_rm, cvals[:, None]]) + st = candidate_stats(VaryingCoefSetup( + aug, target, sw, gshape, main_var=v, modes=modes)) + if st is None: + continue + true_S = st['var0_C'][:nrem] + st['NCdeb_C'][:nrem] + true_B = st['NCdeb_C'][:nrem] + spur_S = float(st['var0_C'][nrem] + st['NCdeb_C'][nrem]) + spur_B = float(st['NCdeb_C'][nrem]) + mts = float(np.max(true_S)) if true_S.size else float('nan') + mtb = float(np.max(true_B)) if true_B.size else float('nan') + rrows.append({ + 'K': int(K), 'true_S': mts, 'spur_S': spur_S, + 'sep_S': spur_S - mts, 'true_B': mtb, 'spur_B': spur_B, + 'sep_B': spur_B - mtb, 'stdmu': float(st['std_mu'][nrem]), + }) + repl_out.append({'removed': label, 'confuser': cterm.name, + 'rows': rrows}) + + result[v] = {'rows': rows, 'kstar': kstar, 'grid': tuple(gshape), + 'n_coord': len(coord_cols), 'n_plain': len(plain_cols), + 'repl': repl_out} + return result + + +def _print_k_sweep(system, res): + """Per-system K-vs-separation table. ``sep_B`` is the coordinate-confuser + region-variation gap (>0 => flaggable); ``sep_S`` the full sum_asym gap.""" + for v, d in res.items(): + ks = d['kstar'] + print(f"\n##### {system} [{v}] grid={d['grid']} K*(microscale)={ks}" + f" n_coord={d['n_coord']} n_plain={d['n_plain']}") + if not d['rows']: + print(' (no evaluable rows)') + continue + if d['n_coord'] == 0: + print(' (no coordinate confusers for this system -- sep_B is NA;' + ' read sep_S only)') + hdr = (f" {'K':>3} {'max_true_B':>11} {'min_spur_B':>11} {'sep_B':>10}" + f" | {'max_true_S':>11} {'min_spur_S':>11} {'sep_S':>10}" + f" | {'beta_spur std/mu':>16} worst_coord") + print(hdr) + print(' ' + '-' * (len(hdr) - 3)) + for r in d['rows']: + fB = ' <0' if (r['sep_B'] == r['sep_B'] and r['sep_B'] < 0) else '' + print(f" {r['K']:>3} {r['max_true_B']:>11.3g} " + f"{r['min_spur_B']:>11.3g} {r['sep_B']:>10.3g}{fB:<3}" + f" | {r['max_true_S']:>11.3g} {r['min_spur_S']:>11.3g} " + f"{r['sep_S']:>10.3g} | {r['worst_stdmu']:>16.3g} " + f"{r['worst_coord'][:34]}") + + # Replacement arm: the load-bearing coordinate-degeneracy test. + for rp in d.get('repl', []): + print(f" -- REPLACE drop [{rp['removed']}] -> load-bearing " + f"{rp['confuser'][:48]} (sep>0 => flagged)") + rh = (f" {'K':>3} {'true_S':>11} {'spur_S':>11} {'sep_S':>10}" + f" | {'true_B':>11} {'spur_B':>11} {'sep_B':>10}" + f" | {'beta std/mu':>12}") + print(rh) + for r in rp['rows']: + fS = ' <0' if (r['sep_S'] == r['sep_S'] and r['sep_S'] < 0) else '' + print(f" {r['K']:>3} {r['true_S']:>11.3g} {r['spur_S']:>11.3g} " + f"{r['sep_S']:>10.3g}{fS:<3} | {r['true_B']:>11.3g} " + f"{r['spur_B']:>11.3g} {r['sep_B']:>10.3g} | " + f"{r['stdmu']:>12.3g}") + + +def main(argv=None): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--systems', default=','.join(_ALL)) + p.add_argument('--dump', action='store_true', + help='Raw per-true-term beta(x) diagnostics instead of the ' + 'comparison tables.') + p.add_argument('--k-sweep', action='store_true', + help='Sweep basis modes K per axis (overriding the Taylor ' + 'microscale via modes=(K,)*D) and tabulate ' + 'true-vs-coordinate-confuser separation per K.') + p.add_argument('--k-list', default=','.join(str(k) for k in _KSWEEP_DEFAULT), + help='Comma-separated K values for --k-sweep.') + args = p.parse_args(argv) + systems = [s.strip() for s in args.systems.split(',') if s.strip()] + + if args.k_sweep: + K_list = [int(x) for x in args.k_list.split(',') if x.strip()] + print(f'K-sweep (modes per axis overriding microscale): {K_list}') + for sysn in systems: + try: + res = k_sweep_system(sysn, K_list) + except Exception as e: + print(f'{sysn:16s} ERROR {type(e).__name__}: {str(e)[:80]}') + continue + if not res: + print(f'{sysn:16s} (no evaluable equations)') + continue + _print_k_sweep(sysn, res) + return 0 + + if args.dump: + for sysn in systems: + try: + dump_system(sysn) + except Exception as e: + print(f'{sysn:16s} ERROR {type(e).__name__}: {str(e)[:80]}') + return 0 + + per_system = {} + redundancy = [] # (system, var, max |std_mu - sqrt(NCraw_C)| over true) + for sysn in systems: + try: + res = analyse_system(sysn) + except Exception as e: + print(f'{sysn:16s} ERROR {type(e).__name__}: {str(e)[:60]}') + continue + if not res: + print(f'{sysn:16s} (no evaluable equations)') + continue + per_system[sysn] = _agg_system(res) + for v, d in res.items(): + sm = d['truth']['std_mu'] + nc = np.sqrt(np.maximum(d['truth']['NCraw_C'], 0.0)) + if sm.size: + redundancy.append((sysn, v, float(np.max(np.abs(sm - nc))))) + + # ---- per-statistic tables over systems -------------------------------- + for k in STAT_KEYS: + print(f'\n=== {k} === (sum_true,max_true LOW good; min_spur HIGH good; ' + f'sep=min_spur-max_true >0 good)') + hdr = (f"{'system':16s} {'sum_true':>11} {'max_true':>11} " + f"{'min_spur':>11} {'sep':>11}") + print(hdr) + print('-' * len(hdr)) + for sysn in systems: + if sysn not in per_system: + continue + a = per_system[sysn][k] + flag = '' if not np.isfinite(a['sep']) else (' <-- NEG' if a['sep'] < 0 else '') + print(f"{sysn:16s} {a['sum_true']:>11.3g} {a['max_true']:>11.3g} " + f"{a['min_spur']:>11.3g} {a['sep']:>11.3g}{flag}") + + # ---- decision summary: rank statistics -------------------------------- + print(f'\n{"=" * 72}\nDECISION SUMMARY (across {len(per_system)} systems)\n{"=" * 72}') + hdr = (f"{'statistic':12s} {'#sep>0':>7} {'min_sep':>11} {'worst_sum_true':>15} " + f"{'worst_max_true':>15}") + print(hdr) + print('-' * len(hdr)) + for k in STAT_KEYS: + seps = [per_system[s][k]['sep'] for s in per_system + if np.isfinite(per_system[s][k]['sep'])] + sums = [per_system[s][k]['sum_true'] for s in per_system + if np.isfinite(per_system[s][k]['sum_true'])] + maxs = [per_system[s][k]['max_true'] for s in per_system + if np.isfinite(per_system[s][k]['max_true'])] + n_pos = sum(1 for x in seps if x > 0) + print(f"{k:12s} {n_pos:>3}/{len(seps):<3} " + f"{(min(seps) if seps else float('nan')):>11.3g} " + f"{(max(sums) if sums else float('nan')):>15.3g} " + f"{(max(maxs) if maxs else float('nan')):>15.3g}") + + if redundancy: + worst = max(redundancy, key=lambda r: r[2]) + print(f'\nParseval check max|std_mu - sqrt(NCraw_C)| over true terms = ' + f'{worst[2]:.2e} ({worst[0]}/{worst[1]}) [should be ~0]') + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/projects/thesis/viz_ac_instability.py b/projects/thesis/viz_ac_instability.py new file mode 100644 index 00000000..4923ba4f --- /dev/null +++ b/projects/thesis/viz_ac_instability.py @@ -0,0 +1,195 @@ +"""Visualize vcoef *instability* on Allen-Cahn for three structural variants: + + 1. missing diffusion (the true ``u_xx`` term dropped), + 2. the true equation, + 3. spurious (an extra ``u_x * u_xx`` term added). + +For each variant we fit the varying-coefficient model and reconstruct, per RHS +term, the coefficient field ``beta_j(x,t) = sum_b gamma_{j,b} B_b(x,t)`` over the +grid. A TRUE term fits a near-CONSTANT coefficient (flat field -> instability +score ~ 0); a misspecified equation forces some coefficient(s) to VARY across +the domain (curved field -> large score). Each panel shows the relative +deviation ``(beta - beta0)/|beta0|`` (flat/white = stable, colored = unstable) +annotated with the per-term instability score s = (Var(gamma0)+NC_deb)/gamma0^2. + +Writes three PNGs to projects/thesis/plots/. +""" +from __future__ import annotations +import os +import sys + +_THIS = os.path.dirname(os.path.abspath(__file__)) +_ROOT = os.path.abspath(os.path.join(_THIS, "..", "..")) +for _p in (_ROOT, _THIS): + if _p not in sys.path: + sys.path.insert(0, _p) + +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.colors import Normalize + +import epde.globals as gv +from epde.interface.equation_translator import translate_equation +from epde.operators.common.stability import VaryingCoefSetup +from kdv_sindy_test import build_pool_only, _normalize_grid_labels +from thesis_runner import load_config, pipeline_settings, _set_seeds + +OUT = os.path.join(_THIS, "plots") +os.makedirs(OUT, exist_ok=True) + +_UXX = "d^2u/dx1^2{power: 1.0}" +_U3 = "u{power: 3.0}" +_U = "u{power: 1.0}" +_UX_UXX = "du/dx1{power: 1.0} * d^2u/dx1^2{power: 1.0}" +_TARGET = "du/dx0{power: 1.0}" + +VARIANTS = [ + ("1_missing_uxx", "Missing diffusion (no $u_{xx}$)", + f"-5.0 * {_U3} + 5.0 * {_U} = {_TARGET}"), + ("2_true", "True Allen-Cahn", + f"0.0001 * {_UXX} + -5.0 * {_U3} + 5.0 * {_U} = {_TARGET}"), + ("3_spurious_ux_uxx", "Spurious ($+\\,u_x\\,u_{xx}$)", + f"0.0001 * {_UXX} + -5.0 * {_U3} + 5.0 * {_U} + 0.0001 * {_UX_UXX} = {_TARGET}"), +] + +_BASE = {"d^2u/dx1^2": "u_{xx}", "du/dx1": "u_x", "du/dx0": "u_t", "u": "u"} + + +def pretty(name: str) -> str: + parts = [] + for fac in name.split(" * "): + base = fac.split("{")[0].strip() + power = 1.0 + if "power:" in fac: + try: + power = float(fac.split("power:")[1].split(",")[0].split("}")[0]) + except ValueError: + power = 1.0 + b = _BASE.get(base, base) + if base == "u" and power != 1.0: + b = f"{b}^{int(power)}" + parts.append(b) + return "$" + r" \cdot ".join(parts) + "$" + + +def setup_pool(): + cfg = load_config("ac") + _set_seeds(0) + search = build_pool_only(cfg, pipeline_settings("new")) + sw = np.asarray(gv.grid_cache.g_func[gv.grid_cache.g_func_mask]).reshape(-1) + gshape = tuple(int(n) for n in gv.grid_cache.inner_shape) + return search, sw, gshape + + +def analyze(search, sw, gshape, eq_str): + soeq = translate_equation(_normalize_grid_labels(eq_str), search.pool, + all_vars=["u"]) + eq = soeq.vals["u"] + eq.main_var_to_explain = "u" + eq.weights_internal = np.ones(len(eq.structure) - 1) + eq.weights_internal_evald = True + eq.weights_final_evald = True + # raw (un-normalized) features so gamma0 is in physical coefficient units; + # the score / relative-deviation field are scale-invariant either way. + _, target, feats = eq.evaluate(normalize=False, return_val=False) + feats = np.asarray(feats, dtype=float) + y = np.asarray(target, dtype=float).reshape(-1) + feat_terms = [t for i, t in enumerate(eq.structure) if i != eq.target_idx] + + setup = VaryingCoefSetup(feats, y, sw, gshape, main_var="u", + fit_intercept=False) + sol = setup._solve_gammas(None) + gamma, B, mk = sol["gamma"], sol["B"], sol["mk"] + is_const = mk == 0 + Bvals = setup._Bvals + scores = np.asarray(setup.score(None), dtype=float) + + rows = [] + for i, term in enumerate(feat_terms): + block = gamma[i * B:(i + 1) * B] + g0 = float(block[is_const][0]) + beta = (Bvals @ block).reshape(gshape) + rows.append({"label": pretty(term.name), "g0": g0, + "beta": beta, "score": float(scores[i])}) + return rows + + +def plot_variant(title, rows, path): + n = len(rows) + fig, axes = plt.subplots(1, n, figsize=(3.7 * n, 4.6), squeeze=False) + axes = axes[0] + norm = Normalize(vmin=-1.0, vmax=1.0) + total = float(np.sum([r["score"] for r in rows])) + im = None + for ax, r in zip(axes, rows): + rel = (r["beta"] - r["g0"]) / (abs(r["g0"]) + 1e-30) + im = ax.imshow(rel, origin="lower", aspect="auto", cmap="coolwarm", + norm=norm) + ax.set_title(f"{r['label']}\n$\\beta_0$={r['g0']:.2g}, " + f"score={r['score']:.1e}", fontsize=10) + ax.set_xlabel("$x$"); ax.set_ylabel("$t$") + ax.set_xticks([]); ax.set_yticks([]) + # explicit layout so the 2-line panel titles never collide with suptitle + fig.subplots_adjust(left=0.06, right=0.87, top=0.72, bottom=0.10, wspace=0.28) + cax = fig.add_axes([0.89, 0.12, 0.015, 0.58]) + cbar = fig.colorbar(im, cax=cax) + cbar.set_label(r"relative deviation $(\beta-\beta_0)/|\beta_0|$") + fig.suptitle(f"{title} total instability $\\Sigma$ = {total:.2e}", + fontsize=13, y=0.93) + fig.savefig(path, dpi=140) + plt.close(fig) + + +def plot_sum_field(title, rows, path): + """One heatmap per case: the SUM over terms of each term's coefficient-field + relative deviation ``(beta_j - beta0_j)/|beta0_j|`` over the (x,t) grid (each + term clipped to +-1 first, so a near-zero-coefficient spurious term cannot + blow up the scale). Flat/white = every coefficient is ~constant (stable); + colored = some coefficient is forced to vary there. The title carries the + total instability Sigma = sum_j score_j.""" + gshape = rows[0]["beta"].shape + field = np.zeros(gshape, dtype=float) + for r in rows: + rel = (r["beta"] - r["g0"]) / (abs(r["g0"]) + 1e-30) + field += np.clip(rel, -1.0, 1.0) + total = float(np.sum([r["score"] for r in rows])) + terms = " + ".join(r["label"] for r in rows) + fig, ax = plt.subplots(figsize=(5.4, 4.9)) + im = ax.imshow(field, origin="lower", aspect="auto", cmap="coolwarm", + norm=Normalize(vmin=-1.5, vmax=1.5)) + ax.set_xlabel("$x$"); ax.set_ylabel("$t$") + ax.set_xticks([]); ax.set_yticks([]) + cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + cbar.set_label(r"$\sum_j$ clip$\,[(\beta_j-\beta_{0j})/|\beta_{0j}|,\ \pm1]$") + ax.set_title(f"{title}\n{terms}\ntotal instability $\\Sigma$ = {total:.2e}", + fontsize=11) + fig.tight_layout() + fig.savefig(path, dpi=150) + plt.close(fig) + + +def main(): + search, sw, gshape = setup_pool() + print(f"AC grid (inner) = {gshape}, N = {int(np.prod(gshape))} points\n") + for tag, title, eq_str in VARIANTS: + try: + rows = analyze(search, sw, gshape, eq_str) + except Exception as e: + import traceback + print(f"[{tag}] FAILED: {type(e).__name__}: {e}") + traceback.print_exc() + continue + path = os.path.join(OUT, f"ac_instability_{tag}.png") + plot_sum_field(title, rows, path) + tot = float(np.sum([r["score"] for r in rows])) + print(f"[{tag}] {title} (Sigma={tot:.3e})") + for r in rows: + print(f" {r['label']:18s} beta0={r['g0']:+.3e} score={r['score']:.3e}") + print(f" saved -> {path}\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main())