Skip to content

Commit 1a45000

Browse files
apoorvalalapoorvalal
authored
Merge pull request #3 from py-econometrics/fe_linear
add FEs and tests
2 parents 8e04850 + cf8b592 commit 1a45000

File tree

10 files changed

+804
-6
lines changed

10 files changed

+804
-6
lines changed

‎.github/workflows/comprehensive-tests.yml

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
name: Comprehensive Tests
2+
3+
on:
4+
push:
5+
branches: [ "*" ]
6+
schedule:
7+
# Run tests weekly on Sundays at 2 AM UTC
8+
- cron: '0 2 * * 0'
9+
10+
jobs:
11+
test-matrix:
12+
name: Test Python ${{ matrix.python-version }} on ${{ matrix.os }}
13+
runs-on: ${{ matrix.os }}
14+
timeout-minutes: 45
15+
strategy:
16+
fail-fast: false
17+
matrix:
18+
os: [ubuntu-latest, macos-latest]
19+
python-version: ["3.10", "3.11", "3.12"]
20+
exclude:
21+
# Skip some combinations to reduce CI time
22+
- os: macos-latest
23+
python-version: "3.10"
24+
25+
steps:
26+
- name: Checkout code
27+
uses: actions/checkout@v4
28+
29+
- name: Set up Python ${{ matrix.python-version }}
30+
uses: actions/setup-python@v4
31+
with:
32+
python-version: ${{ matrix.python-version }}
33+
34+
- name: Cache pip dependencies
35+
uses: actions/cache@v3
36+
with:
37+
path: ~/.cache/pip
38+
key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}
39+
restore-keys: |
40+
${{ runner.os }}-pip-${{ matrix.python-version }}-
41+
${{ runner.os }}-pip-
42+
43+
- name: Install system dependencies (Ubuntu)
44+
if: matrix.os == 'ubuntu-latest'
45+
run: |
46+
sudo apt-get update
47+
sudo apt-get install -y build-essential
48+
49+
- name: Install Python dependencies
50+
run: |
51+
python -m pip install --upgrade pip wheel setuptools
52+
pip install pytest pytest-cov pytest-xdist
53+
54+
- name: Install jaxonometrics
55+
run: |
56+
pip install -e .
57+
58+
- name: Install test dependencies
59+
run: |
60+
pip install pyfixest pandas scikit-learn
61+
62+
- name: Verify installation
63+
run: |
64+
python -c "import jaxonometrics; print('jaxonometrics version:', jaxonometrics.__version__)"
65+
python -c "import jax; print('JAX version:', jax.__version__)"
66+
python -c "import pyfixest; print('pyfixest available')"
67+
68+
- name: Run basic tests
69+
run: |
70+
python -m pytest tests/test_linear.py -v
71+
72+
- name: Run fixed effects tests
73+
run: |
74+
python -m pytest tests/test_fe.py -v
75+
76+
- name: Run all tests with coverage
77+
run: |
78+
python -m pytest tests/ -v --cov=jaxonometrics --cov-report=xml --cov-report=term
79+
80+
- name: Upload coverage to Codecov
81+
if: matrix.python-version == '3.11' && matrix.os == 'ubuntu-latest'
82+
uses: codecov/codecov-action@v3
83+
with:
84+
file: ./coverage.xml
85+
flags: unittests
86+
name: codecov-umbrella
87+

‎.github/workflows/tests.yml

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
name: Tests
2+
3+
on:
4+
push:
5+
branches: [ "*" ] # Run on pushes to all branches
6+
jobs:
7+
test:
8+
runs-on: ubuntu-latest
9+
timeout-minutes: 30
10+
strategy:
11+
matrix:
12+
python-version: ["3.10", "3.11", "3.12"]
13+
14+
steps:
15+
- uses: actions/checkout@v4
16+
17+
- name: Set up Python ${{ matrix.python-version }}
18+
uses: actions/setup-python@v4
19+
with:
20+
python-version: ${{ matrix.python-version }}
21+
22+
- name: Install dependencies
23+
run: |
24+
python -m pip install --upgrade pip
25+
pip install pytest
26+
pip install -e .
27+
28+
- name: Install test dependencies
29+
run: |
30+
pip install pyfixest pandas scikit-learn
31+
32+
- name: Run tests
33+
run: |
34+
python -m pytest tests/ -v --tb=short
35+
36+
- name: Run specific fixed effects tests
37+
run: |
38+
python -m pytest tests/test_fe.py -v
39+
40+
test-fast:
41+
# Quick test job for faster feedback
42+
runs-on: ubuntu-latest
43+
timeout-minutes: 15
44+
steps:
45+
- uses: actions/checkout@v4
46+
47+
- name: Set up Python 3.11
48+
uses: actions/setup-python@v4
49+
with:
50+
python-version: "3.11"
51+
52+
- name: Install dependencies
53+
run: |
54+
python -m pip install --upgrade pip
55+
pip install pytest
56+
pip install -e .
57+
pip install pyfixest pandas scikit-learn
58+
59+
- name: Run minimal test suite
60+
run: |
61+
python -m pytest tests/test_linear.py tests/test_fe.py::test_minimal_example -v

‎.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,3 +28,4 @@
2828
*.out
2929
*.synctex.gz
3030
nb/scratch.ipynb
31+
tmp/

‎README.md

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,66 @@
11
# `jaxonometrics`: Econometrics in jax
22

3-
Simple library that provides performant implementations of standard econometrics routines in the JAX ecosystem.
3+
![Tests](https://github.com/py-econometrics/jaxonometrics/workflows/Tests/badge.svg)
4+
![Fixed Effects Tests](https://github.com/py-econometrics/jaxonometrics/workflows/Fixed%20Effects%20Tests/badge.svg)
5+
6+
Simple library that provides performant implementations of standard econometrics routines in the JAX ecosystem.
47

58
- `jax` arrays everywhere
69
- `lineax` for solving linear systems
710
- `jaxopt` and `optax` for numerical optimization (Levenberg–Marquardt for NNLS-type problems and SGD for larger problems)
11+
12+
## Features
13+
14+
- **Linear Regression** with multiple solver backends (lineax, JAX, numpy)
15+
- **Fixed Effects Regression** with JAX-accelerated alternating projections
16+
- **GMM and IV Estimation**
17+
- **Causal Inference** (IPW, AIPW, Entropy Balancing)
18+
- **Maximum Likelihood Estimation** (Logistic, Poisson)
19+
20+
### Fixed Effects
21+
22+
jaxonometrics supports high-performance fixed effects regression with multiple FE variables:
23+
24+
```python
25+
from jaxonometrics import LinearRegression
26+
import jax.numpy as jnp
27+
28+
# Your data
29+
X = jnp.asarray(data) # (n_obs, n_features)
30+
y = jnp.asarray(target) # (n_obs,)
31+
firm_ids = jnp.asarray(firm_identifiers, dtype=jnp.int32)
32+
year_ids = jnp.asarray(year_identifiers, dtype=jnp.int32)
33+
34+
# Two-way fixed effects
35+
model = LinearRegression(solver="lineax")
36+
model.fit(X, y, fe=[firm_ids, year_ids])
37+
coefficients = model.params["coef"]
38+
```
39+
40+
## Installation and Development
41+
42+
### Install
43+
44+
```bash
45+
uv pip install git+https://github.com/py-econometrics/jaxonometrics
46+
```
47+
48+
or clone the repository and install in editable mode.
49+
50+
51+
### Testing
52+
53+
Run the full test suite:
54+
```bash
55+
pytest tests/ -v
56+
```
57+
58+
Run only fixed effects tests:
59+
```bash
60+
pytest tests/ -m fe -v
61+
```
62+
63+
Run tests excluding slow ones:
64+
```bash
65+
pytest tests/ -m "not slow" -v
66+
```

‎jaxonometrics/__init__.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from .gmm import GMM, LinearIVGMM, TwoStepGMM
1010
from .linear import LinearRegression
1111
from .mle import LogisticRegression, PoissonRegression, MaximumLikelihoodEstimator # Added MLE models
12+
from .demean import demean_jax, prepare_fixed_effects
1213

1314
__all__ = [
1415
"BaseEstimator",
@@ -22,4 +23,6 @@
2223
"MaximumLikelihoodEstimator",
2324
"LogisticRegression",
2425
"PoissonRegression",
26+
"demean_jax",
27+
"prepare_fixed_effects",
2528
]

0 commit comments

Comments
 (0)