feat: Implement MLE and IPW/AIPW #2
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Adds Maximum Likelihood Estimators: - Base `MaximumLikelihoodEstimator` class. - `Logit` regression. - `PoissonRegression`. Adds new Causal Inference methods: - `IPW` (Inverse Propensity Weighting) using Logit for propensity scores. - `AIPW` (Augmented Inverse Propensity Weighting) using configurable outcome and propensity score models (defaults to LinearRegression and Logit). Includes: - Performance enhancement by JIT-compiling relevant functions in existing `LinearRegression` and `EntropyBalancing` classes. - Unit tests for all new estimators (`tests/test_mle.py`, `tests/test_causal.py`) using synthetic data and comparison with sklearn where applicable. - Docstrings and API refinements for the new components.
This resolves a TypeError when calling the JIT-compiled standard error calculation function with a string argument (se_type) that was not marked as static. `n` and `k` are also marked static for robustness as they can influence computations and array shapes known at compile time.
Refactors MaximumLikelihoodEstimator and its subclasses (Logit, PoissonRegression) to use Optax optimizers directly, replacing the deprecated jaxopt.OptaxSolver. - MLE `fit` method now uses a standard Optax optimization loop. - MLE `__init__` accepts an Optax GradientTransformation object. Updates causal estimators: - `IPW` constructor now takes `propensity_optimizer` (Optax object) and `propensity_maxiter` for its internal Logit model. - `AIPW` remains compatible, as its default Logit propensity model now uses Optax, and a custom-configured Logit (with Optax) can be passed. Test files (`test_mle.py`, `test_causal.py`) have been updated or confirmed to be compatible with these changes.
apoorvalal
changed the title
apoorvalal
force-pushed
the
feat/mle-causal-methods
branch
from
7692512 to
04e5a90
Compare
There was a problem hiding this comment.
Pull Request Overview
This PR introduces new maximum likelihood estimators (Logistic and Poisson regression) along with causal inference methods (IPW and AIPW) while also enhancing performance via JIT-compilation and expanding unit test coverage.
- Adds MLE models and causal estimators with corresponding unit tests.
- Refactors naming conventions (e.g. using "coef" consistently) and integrates new methods in the package's public API.
Reviewed Changes
Copilot reviewed 8 out of 9 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/test_mle.py | New unit tests added for logistic and Poisson regression models. |
| tests/test_linear.py | Updated to reflect the naming change from "beta" to "coef". |
| tests/test_causal.py | New tests for IPW and AIPW estimators with synthetic data. |
| nb/linmod.ipynb | Minor notebook adjustments (timestamps and execution counts). |
| jaxonometrics/mle.py | Added the MaximumLikelihoodEstimator base and its subclasses. |
| jaxonometrics/linear.py | Updated VCOV computation using JIT and consistent parameter naming. |
| jaxonometrics/causal.py | Added IPW and AIPW implementations and enhanced documentation. |
| jaxonometrics/init.py | Updated public API to include new estimators and models. |
Comments suppressed due to low confidence (2)
tests/test_linear.py:22
- Ensure that the renaming from 'beta' to 'coef' is consistently reflected in all related tests and modules.
jax_coef = jax_model.params["coef"][1:]
jaxonometrics/causal.py:221
- [nitpick] When instantiating new outcome models via class(), ensure that any necessary constructor arguments are preserved to avoid losing configuration from the original template.
model1 = self.outcome_model_template.__class__() # Create a new instance of the same type
| self.params["se"] = se_values | ||
| else: | ||
| # This case should ideally not be reached if fit() is called first. | ||
| print("Coefficients not available for SE calculation.") |
There was a problem hiding this comment.
[nitpick] Consider raising an error or using a more robust error-handling approach instead of printing a warning when coefficients are unavailable during SE calculation.
Suggested change
| print("Coefficients not available for SE calculation.") | |
| raise ValueError("Coefficients are not available. Ensure `fit()` is called before `_vcov()`.") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Adds Maximum Likelihood Estimators:
MaximumLikelihoodEstimatorclass.Logitregression.PoissonRegression.Adds new Causal Inference methods:
IPW(Inverse Propensity Weighting) using Logit for propensity scores.AIPW(Augmented Inverse Propensity Weighting) using configurable outcome and propensity score models (defaults to LinearRegression and Logit).Includes:
LinearRegressionandEntropyBalancingclasses.tests/test_mle.py,tests/test_causal.py) using synthetic data and comparison with sklearn where applicable.