# jimla
A Python package for Bayesian linear regression using variational inference, inspired by R's `lm()` function and the `broom` package.
## Features
- **Formula parsing**: Uses [fiasto-py](https://github.com/alexhallam/fiasto-py) for formula parsing and [wayne-trade](https://github.com/alexhallam/wayne) for model matrix generation
- **Bayesian inference**: Uses [blackjax](https://github.com/blackjax-devs/blackjax) pathfinder for variational inference
- **Automatic prior scaling**: Stan/brms-style autoscaling makes models robust to data scales
- **Complete Broom API**: Full equivalents of `tidy()`, `augment()`, and `glance()` functions
- **Enhanced display**: Results displayed using tidy-viewer for beautiful, formatted output
- **Progress tracking**: Rich progress bars show variational inference progress
- **Polars integration**: Works seamlessly with Polars DataFrames
- **Scale invariant**: Works with any data scale (dollars, inches, milliseconds) without manual tuning
## Installation
```bash
pip install jimla
```
## Quick Start
```python
import polars as pl
import numpy as np
from jimla import lm, tidy, augment, glance
# Create sample data
np.random.seed(42)
n = 100
x1 = np.random.normal(0, 1, n)
x2 = np.random.normal(0, 1, n)
y = 2 + 1.5 * x1 + 0.8 * x2 + np.random.normal(0, 0.5, n)
df = pl.DataFrame({
"y": y,
"x1": x1,
"x2": x2
})
# Fit regression model
result = lm(df, "y ~ x1 + x2")
# Get tidy output (coefficients and statistics)
tidy_result = tidy(result)
# Augment original data with model information
augmented_data = augment(result, df)
# Get one-row model summary
model_summary = glance(result)
```
## API Reference
### `lm(df: pl.DataFrame, formula: str, **kwargs) -> RegressionResult`
Fit a Bayesian linear regression model using blackjax pathfinder.
**Parameters:**
- `df`: Polars DataFrame containing the data
- `formula`: Wilkinson's formula string (e.g., "y ~ x1 + x2")
- `**kwargs`: Additional arguments passed to blackjax pathfinder
**Returns:**
- `RegressionResult`: Object containing coefficients, R-squared, and model information
### `tidy(result: RegressionResult, display: bool = True, title: str = None, color_theme: str = "default") -> pl.DataFrame`
Create a tidy summary of regression results, similar to `broom::tidy()`.
Uses tidy-viewer for enhanced display by default.
**Parameters:**
- `result`: RegressionResult from `lm()`
- `display`: Whether to display the results using tidy-viewer (default: True)
- `title`: Optional title for the display
- `color_theme`: Color theme for display ("default", "dracula", etc.)
**Returns:**
- `pl.DataFrame`: DataFrame with columns: term, estimate, std.error, statistic, p.value, conf.low, conf.high
### `augment(result: RegressionResult, data: pl.DataFrame, display: bool = True, title: str = None, color_theme: str = "default") -> pl.DataFrame`
Add model information to the original data, similar to `broom::augment()`.
**Parameters:**
- `result`: RegressionResult from `lm()`
- `data`: Original Polars DataFrame
- `display`: Whether to display results using tidy-viewer (default: True)
- `title`: Optional title for the display
- `color_theme`: Color theme for display ("default", "dracula", etc.)
**Returns:**
- `pl.DataFrame`: Original data plus model columns: .fitted, .resid, .fitted_std, .fitted_low, .fitted_high, .hat, .std.resid, .sigma
### `glance(result: RegressionResult, display: bool = True, title: str = None, color_theme: str = "default") -> pl.DataFrame`
Create a one-row model summary, similar to `broom::glance()`.
**Parameters:**
- `result`: RegressionResult from `lm()`
- `display`: Whether to display results using tidy-viewer (default: True)
- `title`: Optional title for the display
- `color_theme`: Color theme for display ("default", "dracula", etc.)
**Returns:**
- `pl.DataFrame`: One-row DataFrame with: r_squared, adj_r_squared, sigma, sigma_std, n_obs, n_params, df_residual, n_samples, n_eff, formula, method
## Supported Formula Syntax
jimla supports Wilkinson's notation through fiasto-py and wayne-trade:
- **Basic formulas**: `y ~ x1 + x2`
- **Interactions**: `y ~ x1 * x2`
- **Polynomials**: `y ~ poly(x1, 2)`
- **Intercept control**: `y ~ x1 + x2 - 1` (no intercept)
- **Complex interactions**: `y ~ x1 + x2*x3 + poly(x1, 2)`
## Example Output
```
Formula: y ~ x1 + x2
R-squared: 0.8951
Number of observations: 100
Number of parameters: 3
Coefficients:
(Intercept): 2.0370
x1: 1.6048
x2: 0.7877
Tidy output (with tidy-viewer):
```
Regression Results: y ~ x1 + x2
tv dim: 3 x 7
term estimate std.error statistic p.value conf.low conf.high
<str> <f64> <f64> <f64> <f64> <f64> <f64>
1 (Intercept) 2.04 0.0463 43.9 0 1.94 2.12
2 x1 1.60 0.0592 27.1 0 1.49 1.72
3 x2 0.788 0.0654 12.0 0 0.662 0.912
Model Summary:
Formula: y ~ x1 + x2
R-squared: 0.8951
Observations: 100
Parameters: 3
```
```
## Dependencies
- [blackjax](https://github.com/blackjax-devs/blackjax) - Bayesian inference
- [fiasto-py](https://github.com/alexhallam/fiasto-py) - Formula parsing
- [wayne-trade](https://github.com/alexhallam/wayne) - Model matrix generation
- [polars](https://github.com/pola-rs/polars) - Data manipulation
- [jax](https://github.com/google/jax) - Numerical computing
- [tidy-viewer-py](https://github.com/alexhallam/tv/tree/main/tidy-viewer-py) - Enhanced data display
- [rich](https://github.com/Textualize/rich) - Progress bars and terminal formatting
## License
MIT License
Raw data
{
"_id": null,
"home_page": null,
"name": "jimla",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "bayesian, jax, machine-learning, polars, regression, statistics",
"author": null,
"author_email": "Alex Hallam <alex.hallam@cfacorp.com>",
"download_url": "https://files.pythonhosted.org/packages/94/11/32728ea04018e0a94c925cad473a9864d1d5567d79af275105eaf1c76ad6/jimla-0.1.0.tar.gz",
"platform": null,
"description": "# jimla\n\nA Python package for Bayesian linear regression using variational inference, inspired by R's `lm()` function and the `broom` package.\n\n## Features\n\n- **Formula parsing**: Uses [fiasto-py](https://github.com/alexhallam/fiasto-py) for formula parsing and [wayne-trade](https://github.com/alexhallam/wayne) for model matrix generation\n- **Bayesian inference**: Uses [blackjax](https://github.com/blackjax-devs/blackjax) pathfinder for variational inference\n- **Automatic prior scaling**: Stan/brms-style autoscaling makes models robust to data scales\n- **Complete Broom API**: Full equivalents of `tidy()`, `augment()`, and `glance()` functions\n- **Enhanced display**: Results displayed using tidy-viewer for beautiful, formatted output\n- **Progress tracking**: Rich progress bars show variational inference progress\n- **Polars integration**: Works seamlessly with Polars DataFrames\n- **Scale invariant**: Works with any data scale (dollars, inches, milliseconds) without manual tuning\n\n## Installation\n\n```bash\npip install jimla\n```\n\n## Quick Start\n\n```python\nimport polars as pl\nimport numpy as np\nfrom jimla import lm, tidy, augment, glance\n\n# Create sample data\nnp.random.seed(42)\nn = 100\nx1 = np.random.normal(0, 1, n)\nx2 = np.random.normal(0, 1, n)\ny = 2 + 1.5 * x1 + 0.8 * x2 + np.random.normal(0, 0.5, n)\n\ndf = pl.DataFrame({\n \"y\": y,\n \"x1\": x1,\n \"x2\": x2\n})\n\n# Fit regression model\nresult = lm(df, \"y ~ x1 + x2\")\n\n# Get tidy output (coefficients and statistics)\ntidy_result = tidy(result)\n\n# Augment original data with model information\naugmented_data = augment(result, df)\n\n# Get one-row model summary\nmodel_summary = glance(result)\n```\n\n## API Reference\n\n### `lm(df: pl.DataFrame, formula: str, **kwargs) -> RegressionResult`\n\nFit a Bayesian linear regression model using blackjax pathfinder.\n\n**Parameters:**\n- `df`: Polars DataFrame containing the data\n- `formula`: Wilkinson's formula string (e.g., \"y ~ x1 + x2\")\n- `**kwargs`: Additional arguments passed to blackjax pathfinder\n\n**Returns:**\n- `RegressionResult`: Object containing coefficients, R-squared, and model information\n\n### `tidy(result: RegressionResult, display: bool = True, title: str = None, color_theme: str = \"default\") -> pl.DataFrame`\n\nCreate a tidy summary of regression results, similar to `broom::tidy()`.\nUses tidy-viewer for enhanced display by default.\n\n**Parameters:**\n- `result`: RegressionResult from `lm()`\n- `display`: Whether to display the results using tidy-viewer (default: True)\n- `title`: Optional title for the display\n- `color_theme`: Color theme for display (\"default\", \"dracula\", etc.)\n\n**Returns:**\n- `pl.DataFrame`: DataFrame with columns: term, estimate, std.error, statistic, p.value, conf.low, conf.high\n\n### `augment(result: RegressionResult, data: pl.DataFrame, display: bool = True, title: str = None, color_theme: str = \"default\") -> pl.DataFrame`\n\nAdd model information to the original data, similar to `broom::augment()`.\n\n**Parameters:**\n- `result`: RegressionResult from `lm()`\n- `data`: Original Polars DataFrame\n- `display`: Whether to display results using tidy-viewer (default: True)\n- `title`: Optional title for the display\n- `color_theme`: Color theme for display (\"default\", \"dracula\", etc.)\n\n**Returns:**\n- `pl.DataFrame`: Original data plus model columns: .fitted, .resid, .fitted_std, .fitted_low, .fitted_high, .hat, .std.resid, .sigma\n\n### `glance(result: RegressionResult, display: bool = True, title: str = None, color_theme: str = \"default\") -> pl.DataFrame`\n\nCreate a one-row model summary, similar to `broom::glance()`.\n\n**Parameters:**\n- `result`: RegressionResult from `lm()`\n- `display`: Whether to display results using tidy-viewer (default: True)\n- `title`: Optional title for the display\n- `color_theme`: Color theme for display (\"default\", \"dracula\", etc.)\n\n**Returns:**\n- `pl.DataFrame`: One-row DataFrame with: r_squared, adj_r_squared, sigma, sigma_std, n_obs, n_params, df_residual, n_samples, n_eff, formula, method\n\n## Supported Formula Syntax\n\njimla supports Wilkinson's notation through fiasto-py and wayne-trade:\n\n- **Basic formulas**: `y ~ x1 + x2`\n- **Interactions**: `y ~ x1 * x2`\n- **Polynomials**: `y ~ poly(x1, 2)`\n- **Intercept control**: `y ~ x1 + x2 - 1` (no intercept)\n- **Complex interactions**: `y ~ x1 + x2*x3 + poly(x1, 2)`\n\n## Example Output\n\n```\nFormula: y ~ x1 + x2\nR-squared: 0.8951\nNumber of observations: 100\nNumber of parameters: 3\n\nCoefficients:\n (Intercept): 2.0370\n x1: 1.6048\n x2: 0.7877\n\nTidy output (with tidy-viewer):\n```\nRegression Results: y ~ x1 + x2\n\n tv dim: 3 x 7\n term estimate std.error statistic p.value conf.low conf.high \n <str> <f64> <f64> <f64> <f64> <f64> <f64> \n 1 (Intercept) 2.04 0.0463 43.9 0 1.94 2.12 \n 2 x1 1.60 0.0592 27.1 0 1.49 1.72 \n 3 x2 0.788 0.0654 12.0 0 0.662 0.912 \n\nModel Summary:\n Formula: y ~ x1 + x2\n R-squared: 0.8951\n Observations: 100\n Parameters: 3\n```\n```\n\n## Dependencies\n\n- [blackjax](https://github.com/blackjax-devs/blackjax) - Bayesian inference\n- [fiasto-py](https://github.com/alexhallam/fiasto-py) - Formula parsing\n- [wayne-trade](https://github.com/alexhallam/wayne) - Model matrix generation\n- [polars](https://github.com/pola-rs/polars) - Data manipulation\n- [jax](https://github.com/google/jax) - Numerical computing\n- [tidy-viewer-py](https://github.com/alexhallam/tv/tree/main/tidy-viewer-py) - Enhanced data display\n- [rich](https://github.com/Textualize/rich) - Progress bars and terminal formatting\n\n## License\n\nMIT License\n",
"bugtrack_url": null,
"license": "MIT",
"summary": "Bayesian linear regression with variational inference, inspired by R's lm() and broom",
"version": "0.1.0",
"project_urls": {
"Bug Tracker": "https://github.com/alexhallam/jimla/issues",
"Documentation": "https://github.com/alexhallam/jimla#readme",
"Homepage": "https://github.com/alexhallam/jimla",
"Repository": "https://github.com/alexhallam/jimla"
},
"split_keywords": [
"bayesian",
" jax",
" machine-learning",
" polars",
" regression",
" statistics"
],
"urls": [
{
"comment_text": null,
"digests": {
"blake2b_256": "f95aaa6d43f7511cb239bba1441d0684aa700bc2ecb2b7a2936ebf8745a43924",
"md5": "3c244e1c0576954050621c9e31c2ccff",
"sha256": "b71b0c05d0e4d762437f626fc01b4d41d86f2ca04cdd3f67bb430420d2ed1d3b"
},
"downloads": -1,
"filename": "jimla-0.1.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "3c244e1c0576954050621c9e31c2ccff",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 11021,
"upload_time": "2025-09-14T20:21:33",
"upload_time_iso_8601": "2025-09-14T20:21:33.266704Z",
"url": "https://files.pythonhosted.org/packages/f9/5a/aa6d43f7511cb239bba1441d0684aa700bc2ecb2b7a2936ebf8745a43924/jimla-0.1.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": null,
"digests": {
"blake2b_256": "941132728ea04018e0a94c925cad473a9864d1d5567d79af275105eaf1c76ad6",
"md5": "a117cbb0ce469b71bc7b62d62c1e2c74",
"sha256": "8368dba53bc0990d9beae9d40eba5145e33bc012b3ad775af11e2514f24d8d42"
},
"downloads": -1,
"filename": "jimla-0.1.0.tar.gz",
"has_sig": false,
"md5_digest": "a117cbb0ce469b71bc7b62d62c1e2c74",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 55744,
"upload_time": "2025-09-14T20:21:34",
"upload_time_iso_8601": "2025-09-14T20:21:34.681405Z",
"url": "https://files.pythonhosted.org/packages/94/11/32728ea04018e0a94c925cad473a9864d1d5567d79af275105eaf1c76ad6/jimla-0.1.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-09-14 20:21:34",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "alexhallam",
"github_project": "jimla",
"travis_ci": false,
"coveralls": false,
"github_actions": false,
"lcname": "jimla"
}
JSON
