Skip to content

feat: add entropic regression for equation discovery (#8) - #61

Open
ParitKansal wants to merge 1 commit into
Center-For-Complex-Systems-Science:mainfrom
ParitKansal:feature/entropic-regression
Open

ParitKansal wants to merge 1 commit into
Center-For-Complex-Systems-Science:mainfrom
ParitKansal:feature/entropic-regression

Conversation

@ParitKansal

@ParitKansal ParitKansal commented Oct 23, 2026 •

Copy link
Copy Markdown

Addresses #8 (equation discovery).

Summary

Adds entropic regression (ER), from AlMomani, Sun and Bollt, "How entropic regression beats the outliers problem in nonlinear system identification", Chaos 30, 013107 (2020), arXiv:1905.08061. ER finds a sparse equation for each state variable in a library of basis functions. It uses the same forward selection and backward elimination as oCSE, but scores each candidate term by the conditional mutual information of its least-squares model output.

New module causationentropy/core/estimators.py:

Name What it does
EntropicRegression fit(X, dt=...) or fit(X, X_dot=...), then equations(), predict(), coef_, support_
entropic_regression(f, Phi) ER for one target (Algorithm 1 of the paper)
polynomial_library(X, degree) Monomial basis with term names
central_difference(X, dt) Second-order derivative estimates
from causationentropy.core.estimators import EntropicRegression
model = EntropicRegression(degree=2, random_state=0).fit(X, dt=0.01, feature_names=["x", "y", "z"])
print("\n".join(model.equations()))
# x' = -9.982 x + 9.982 y
# y' = 27.812 x - 0.967 y - 0.995 x z
# z' = -2.660 z + 0.997 x y

The code reuses the library's conditional_mutual_information (default information="knn", k = 2 as in the paper) and shuffle_test. It was written from the paper only; no code from the GPL-licensed ERFit MATLAB package was used.

Design choices (documented in the docstrings)

  1. Tolerance: dynamic shuffle test at every forward and backward step (paper SI.3, Algorithm 2). The authors' ERFit package defaults to a static tolerance (0.99 quantile of I(f; f_π), 200 permutations), which the docstring describes.
  2. Constant term: mutual information does not change when a constant is added, so a column of ones can never pass the information criterion. A constant offset is handled with fit_intercept=True instead.
  3. Exact fits: when the conditioning model output already fits f exactly, the CMI is set to 0 (its true value). Without this, the k-NN estimate on nearly identical vectors is positive and terms with zero coefficients get selected on exact data.
  4. Metric: for information="knn", the metric defaults to Chebyshev in every step. In the current library, knn_conditional_mutual_information uses Minkowski when Z is None and Chebyshev otherwise (fixed by fix: break ties in the k-NN estimators and use Chebyshev when Z is None (follow-up to #14) #53); without this default, the first ER step would use a different metric from the later ones.

Checks

Check Result
New tests (tests/core/test_estimators.py, 17 tests + 3 doctests) Pass
Lorenz, exact derivatives Exact terms and coefficients (σ = 10, ρ = 28, β = 8/3)
Lorenz, state noise 0.01, central differences Exact terms, coefficients within 5%
Logistic map (X_dot = next state) x' = 3.9 x − 3.9 x²
Outlier case of the paper's Fig. 3 (ε₁ = 1e-5, ε₂ = 0.2, p = 0.2, 1500 samples, dt = 0.005) ER: exact 7 terms, relative coefficient error 0.9%. Plain least squares: 11 terms with abs(a) > 0.05, error 6.4%
Full test suite 447 passed. 10 tests in test_data_integration.py fail the same way on the symbolic-cmi-estimator branch, and test_standard_knn also fails on main; this PR does not touch discover_network or the estimators those tests use
Docs (sphinx-build) Builds; no new warnings (370 on main and on this branch)
black, isort, flake8 Clean

Independent re-implementation

A separate session, given only the ER paper (no access to this code), re-implemented ER, wrote its own tests, and froze its code with SHA-256 hashes. Then it compared both versions on identical inputs:

  • Polynomial library, central differences, least-squares fits, and k-NN MI and CMI (same metric): identical (difference 0).
  • Lorenz (exact and noisy), sparse linear model, pure noise: same selected terms.
  • The only bug found was the metric inconsistency above (in the existing estimator, not in this PR); it is handled by choice 4.
  • It also read the ERFit MATLAB code (without running it) to compare defaults; those differences are listed under choice 1 and in the docstring.

Notes for review

  • Issue Dynamical Systems Discovery with Causation Entropy #8 mentions "notebooks"; I couldn't find these in the repository. If there are reference notebooks, I'm happy to compare against them or add a tutorial notebook.
  • Cost: the k-NN estimator builds full distance matrices, and the dynamic test runs a shuffle test at every step, so 1500 samples take about 40 s for the 3 Lorenz equations.

…stems-Science#8)

New module causationentropy.core.estimators with EntropicRegression,
entropic_regression, polynomial_library and central_difference, following
AlMomani, Sun and Bollt (Chaos 30, 013107, 2020): greedy forward selection
and backward elimination of basis functions scored by the conditional
mutual information of their least-squares model output, with shuffle tests.
The k-NN estimator uses the Chebyshev metric in every step. Adds tests, an
API docs page and a changelog entry.
@ParitKansal
ParitKansal force-pushed the feature/entropic-regression branch from 2d8ca47 to 3e7ee38 Compare October 23, 2026 05:39

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant