marginaleffects 0.6.0__tar.gz → 0.6.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (142) hide show
  1. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/PKG-INFO +2 -4
  2. marginaleffects-0.6.2/benchmarks/benchmark_inference.py +39 -0
  3. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/__init__.py +5 -9
  4. marginaleffects-0.6.2/marginaleffects/by.py +297 -0
  5. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/classes/__init__.py +7 -5
  6. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/classes/model.py +60 -34
  7. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/classes/result.py +22 -12
  8. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/comparisons.py +197 -65
  9. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/datagrid.py +31 -27
  10. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/datasets.py +4 -4
  11. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/docstrings/params.py +14 -7
  12. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/docstrings/qmd.py +2 -1
  13. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/estimands.py +6 -5
  14. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/formula.py +0 -1
  15. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/hypothesis_compile.py +115 -4
  16. marginaleffects-0.6.2/marginaleffects/inference/__init__.py +30 -0
  17. marginaleffects-0.6.2/marginaleffects/inference/analytic.py +258 -0
  18. marginaleffects-0.6.2/marginaleffects/inference/delta.py +64 -0
  19. marginaleffects-0.6.2/marginaleffects/inference/gradients.py +85 -0
  20. marginaleffects-0.6.2/marginaleffects/inference/statistics.py +37 -0
  21. marginaleffects-0.6.2/marginaleffects/inference/unconditional.py +733 -0
  22. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/linearmodels/__init__.py +0 -1
  23. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/linearmodels/model.py +16 -10
  24. marginaleffects-0.6.2/marginaleffects/plan.py +7 -0
  25. marginaleffects-0.6.2/marginaleffects/planning/__init__.py +43 -0
  26. marginaleffects-0.6.2/marginaleffects/planning/core.py +436 -0
  27. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/plot/__init__.py +0 -1
  28. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/plot/common.py +16 -16
  29. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/plot/comparisons.py +7 -6
  30. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/plot/predictions.py +6 -5
  31. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/plot/slopes.py +8 -7
  32. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/predictions.py +59 -52
  33. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/pyfixest/model.py +13 -6
  34. marginaleffects-0.6.2/marginaleffects/sanitize/__init__.py +49 -0
  35. marginaleffects-0.6.2/marginaleffects/sanitize/by.py +47 -0
  36. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/categorical.py +2 -2
  37. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/comparison.py +13 -11
  38. marginaleffects-0.6.0/marginaleffects/sanitize/sanitize_model.py → marginaleffects-0.6.2/marginaleffects/sanitize/model.py +1 -4
  39. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/newdata.py +8 -15
  40. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/utils.py +19 -8
  41. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/validation.py +9 -6
  42. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/variables.py +55 -17
  43. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/vcov.py +9 -3
  44. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sklearn/__init__.py +0 -1
  45. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sklearn/model.py +6 -4
  46. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/slopes.py +5 -5
  47. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/statsmodels/__init__.py +0 -1
  48. marginaleffects-0.6.2/marginaleffects/statsmodels/model.py +332 -0
  49. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/test/equivalence.py +1 -3
  50. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/test/formula.py +3 -3
  51. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/test/joint.py +3 -2
  52. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/test/main.py +5 -7
  53. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/transform.py +1 -2
  54. marginaleffects-0.6.2/marginaleffects/uncertainty.py +8 -0
  55. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/utils.py +42 -20
  56. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects.egg-info/PKG-INFO +2 -4
  57. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects.egg-info/SOURCES.txt +18 -12
  58. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects.egg-info/requires.txt +0 -2
  59. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/pyproject.toml +2 -5
  60. marginaleffects-0.6.2/tests/test_adapter_contract.py +38 -0
  61. marginaleffects-0.6.2/tests/test_aggregation_parity.py +498 -0
  62. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_comparison_plan.py +2 -2
  63. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_comparisons.py +28 -2
  64. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_comparisons_interaction.py +1 -1
  65. marginaleffects-0.6.2/tests/test_gradients.py +77 -0
  66. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_hypotheses.py +4 -2
  67. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_hypothesis_compile.py +32 -2
  68. marginaleffects-0.6.2/tests/test_inference_analytic.py +280 -0
  69. marginaleffects-0.6.2/tests/test_jacobian_method.py +101 -0
  70. marginaleffects-0.6.2/tests/test_package_imports.py +53 -0
  71. marginaleffects-0.6.2/tests/test_planning_estimands.py +45 -0
  72. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_prediction_plan.py +9 -3
  73. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_pyfixest.py +29 -0
  74. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_sklearn.py +43 -0
  75. marginaleffects-0.6.2/tests/test_statsmodels_design.py +393 -0
  76. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_mixedlm.py +1 -1
  77. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_mnlogit.py +1 -1
  78. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_ordinal.py +11 -2
  79. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_vcov.py +5 -1
  80. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_utils.py +17 -7
  81. marginaleffects-0.6.0/benchmarks/benchmark_autodiff.py +0 -111
  82. marginaleffects-0.6.0/marginaleffects/autodiff/__init__.py +0 -60
  83. marginaleffects-0.6.0/marginaleffects/autodiff/glm/__init__.py +0 -1
  84. marginaleffects-0.6.0/marginaleffects/autodiff/glm/families.py +0 -108
  85. marginaleffects-0.6.0/marginaleffects/autodiff/lower.py +0 -219
  86. marginaleffects-0.6.0/marginaleffects/autodiff/ops.py +0 -60
  87. marginaleffects-0.6.0/marginaleffects/autodiff/pipeline.py +0 -237
  88. marginaleffects-0.6.0/marginaleffects/by.py +0 -122
  89. marginaleffects-0.6.0/marginaleffects/plan.py +0 -170
  90. marginaleffects-0.6.0/marginaleffects/sanitize/__init__.py +0 -27
  91. marginaleffects-0.6.0/marginaleffects/sanitize/by.py +0 -15
  92. marginaleffects-0.6.0/marginaleffects/settings.py +0 -119
  93. marginaleffects-0.6.0/marginaleffects/statsmodels/model.py +0 -272
  94. marginaleffects-0.6.0/marginaleffects/uncertainty.py +0 -97
  95. marginaleffects-0.6.0/tests/test_autodiff.py +0 -546
  96. marginaleffects-0.6.0/tests/test_autodiff_lower.py +0 -180
  97. marginaleffects-0.6.0/tests/test_autodiff_pipeline.py +0 -240
  98. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/README.md +0 -0
  99. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/benchmarks/__init__.py +0 -0
  100. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/docstrings/__init__.py +0 -0
  101. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/pyfixest/__init__.py +0 -0
  102. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/deprecated.py +0 -0
  103. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/sanitize/hypothesis_null.py +0 -0
  104. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/test/__init__.py +10 -10
  105. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects/test/core.py +0 -0
  106. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects.egg-info/dependency_links.txt +0 -0
  107. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/marginaleffects.egg-info/top_level.txt +0 -0
  108. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/setup.cfg +0 -0
  109. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/__init__.py +0 -0
  110. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/helpers.py +0 -0
  111. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_analytic.py +0 -0
  112. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_bugfix.py +0 -0
  113. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_by.py +0 -0
  114. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_categorical.py +0 -0
  115. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_categorical_validation.py +0 -0
  116. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_datagrid_01.py +0 -0
  117. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_datagrid_02.py +0 -0
  118. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_equivalence.py +0 -0
  119. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_formula.py +0 -0
  120. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_formulaic_utils.py +0 -0
  121. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_hypotheses_joint.py +0 -0
  122. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_hypothesis.py +0 -0
  123. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_jss.py +0 -0
  124. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_linearmodels_panelols.py +0 -0
  125. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_missing.py +0 -0
  126. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_newdata.py +0 -0
  127. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_plan.py +0 -0
  128. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_plot_comparisons.py +0 -0
  129. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_plot_predictions.py +0 -0
  130. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_plot_slopes.py +0 -0
  131. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_predictions.py +0 -0
  132. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_slopes.py +0 -0
  133. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels.py +0 -0
  134. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_logit.py +0 -0
  135. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_negativebinomial.py +0 -0
  136. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_ols.py +0 -0
  137. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_poisson.py +0 -0
  138. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_probit.py +0 -0
  139. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_quantreg.py +0 -0
  140. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_statsmodels_wls.py +0 -0
  141. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/test_typical.py +0 -0
  142. {marginaleffects-0.6.0 → marginaleffects-0.6.2}/tests/utilities.py +0 -0
@@ -1,12 +1,11 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: marginaleffects
3
- Version: 0.6.0
3
+ Version: 0.6.2
4
4
  Summary: Predictions, counterfactual comparisons, slopes, and hypothesis tests for statistical models.
5
5
  License-Expression: GPL-3.0-or-later
6
- Requires-Python: >=3.10
6
+ Requires-Python: >=3.12
7
7
  Description-Content-Type: text/markdown
8
8
  Requires-Dist: formulaic>=1.0.2
9
- Requires-Dist: jax>=0.4.0
10
9
  Requires-Dist: narwhals>=1.34.0
11
10
  Requires-Dist: numpy>=2.0.0
12
11
  Requires-Dist: patsy>=1.0.1
@@ -22,7 +21,6 @@ Requires-Dist: linearmodels>=6.1; extra == "test"
22
21
  Requires-Dist: pandas>=2.2.2; extra == "test"
23
22
  Requires-Dist: pre-commit>=4.2.0; extra == "test"
24
23
  Requires-Dist: numba; extra == "test"
25
- Requires-Dist: pyarrow>=17.0.0; extra == "test"
26
24
  Requires-Dist: pyfixest>=0.28.0; extra == "test"
27
25
  Requires-Dist: statsmodels>=0.14.0; extra == "test"
28
26
  Requires-Dist: scikit-learn; extra == "test"
@@ -0,0 +1,39 @@
1
+ """End-to-end inference benchmark across observations and coefficient counts.
2
+
3
+ Run with: ``uv run --all-extras python benchmarks/benchmark_inference.py``.
4
+ """
5
+
6
+ import time
7
+
8
+ import numpy as np
9
+ import polars as pl
10
+ import statsmodels.api as sm
11
+
12
+ from marginaleffects import avg_comparisons, avg_predictions, comparisons, predictions
13
+
14
+
15
+ def make_model(n_obs, n_coef, seed=42):
16
+ rng = np.random.default_rng(seed)
17
+ X = rng.normal(size=(n_obs, n_coef))
18
+ beta = rng.normal(size=n_coef)
19
+ data = pl.DataFrame({f"x{i}": X[:, i] for i in range(n_coef)})
20
+ data = data.with_columns(y=pl.Series(X @ beta + rng.normal(size=n_obs)))
21
+ formula = "y ~ " + " + ".join(f"x{i}" for i in range(n_coef))
22
+ return sm.OLS.from_formula(formula, data.to_pandas()).fit()
23
+
24
+
25
+ def elapsed(call):
26
+ start = time.perf_counter()
27
+ call()
28
+ return time.perf_counter() - start
29
+
30
+
31
+ if __name__ == "__main__":
32
+ cases = (predictions, avg_predictions, comparisons, avg_comparisons)
33
+ for n_obs in (100, 1_000, 10_000):
34
+ for n_coef in (5, 25, 100):
35
+ model = make_model(n_obs, n_coef)
36
+ for func in cases:
37
+ kwargs = {"variables": "x0"} if "comparisons" in func.__name__ else {}
38
+ duration = elapsed(lambda: func(model, **kwargs))
39
+ print(f"{func.__name__:18} n={n_obs:5} p={n_coef:3}: {duration:.4f}s")
@@ -1,7 +1,5 @@
1
- from importlib import import_module
2
- import importlib.util
3
1
  import sys
4
-
2
+ from importlib import import_module
5
3
 
6
4
  _EXPORTS = {
7
5
  "avg_comparisons": ("marginaleffects.comparisons", "avg_comparisons"),
@@ -20,14 +18,12 @@ _EXPORTS = {
20
18
  "fit_linearmodels": ("marginaleffects.linearmodels.model", "fit_linearmodels"),
21
19
  "get_dataset": ("marginaleffects.datasets", "get_dataset"),
22
20
  "MarginaleffectsResult": ("marginaleffects.classes", "MarginaleffectsResult"),
23
- "autodiff": ("marginaleffects.settings", "autodiff"),
24
- "set_autodiff": ("marginaleffects.settings", "set_autodiff"),
25
- "get_autodiff": ("marginaleffects.settings", "get_autodiff"),
21
+ "vcovUnconditional": (
22
+ "marginaleffects.inference.unconditional",
23
+ "vcovUnconditional",
24
+ ),
26
25
  }
27
26
 
28
- if importlib.util.find_spec("jax") is not None:
29
- _EXPORTS["autodiff_module"] = ("marginaleffects.autodiff", None)
30
-
31
27
  __all__ = list(_EXPORTS.keys())
32
28
 
33
29
 
@@ -0,0 +1,297 @@
1
+ import warnings
2
+
3
+ import polars as pl
4
+
5
+ from .planning import AggGroup
6
+ from .sanitize.by import by_frame_keys, by_is_frame
7
+
8
+ BY_FRAME_COVERAGE_WARNING = (
9
+ "The `by` data frame does not cover all combinations of response levels "
10
+ "and/or predictors. Some estimates will not be included in the aggregation."
11
+ )
12
+
13
+
14
+ def _mean_expr() -> pl.Expr:
15
+ """Unweighted group mean, dropping missing estimates as R's `na.rm` does.
16
+
17
+ Polars propagates NaN through `mean()` while skipping nulls. Treating the
18
+ two alike keeps this displayed estimate identical to the one the plan
19
+ replay recomputes, which the replay baseline check requires.
20
+ """
21
+ return pl.col("estimate").fill_nan(None).mean().alias("estimate")
22
+
23
+
24
+ def _weighted_mean_expr(wts: str) -> pl.Expr:
25
+ """Weighted group mean under R's missing-value rules.
26
+
27
+ A missing estimate zeroes both itself and its weight, so it leaves the
28
+ aggregate untouched. A zero weight blanks the estimate it multiplies,
29
+ because `0 * Inf` is NaN and would otherwise poison the whole group. A
30
+ missing weight propagates: dropping the row it belongs to would silently
31
+ report a different estimand than the one requested.
32
+ """
33
+ est = pl.col("estimate")
34
+ w = pl.col(wts).cast(pl.Float64).fill_null(float("nan"))
35
+ missing = est.is_null() | est.is_nan()
36
+ numerator = pl.when(missing | (w == 0)).then(0.0).otherwise(est * w).sum()
37
+ denominator = pl.when(missing).then(0.0).otherwise(w).sum()
38
+ return (numerator / denominator).alias("estimate")
39
+
40
+
41
+ def harmonize_by_types(estimates: pl.DataFrame, by: pl.DataFrame) -> pl.DataFrame:
42
+ """Reconcile join-key dtypes between the estimates and a `by` data frame.
43
+
44
+ A label table is usually typed by hand, so its keys routinely disagree with
45
+ the estimates on numeric-versus-string. Joining on mismatched dtypes finds
46
+ nothing and would silently drop every row, so the keys are cast to the
47
+ estimates' side before the join rather than after it fails.
48
+ """
49
+ casts = []
50
+ for column in by.columns:
51
+ if column == "by" or column not in estimates.columns:
52
+ continue
53
+ target = estimates[column].dtype
54
+ if by[column].dtype == target:
55
+ continue
56
+ if target.is_numeric() or target in (pl.String, pl.Utf8):
57
+ casts.append(pl.col(column).cast(target, strict=False))
58
+ else:
59
+ casts.append(pl.col(column).cast(pl.String).cast(target, strict=False))
60
+ return by.with_columns(casts) if casts else by
61
+
62
+
63
+ def attach_by_labels(estimates: pl.DataFrame, by: pl.DataFrame, verbose: bool = True):
64
+ """Left-join a `by` data frame's labels onto the estimates.
65
+
66
+ Returns the labelled estimates and a boolean mask of the rows the label
67
+ table covers. Uncovered rows are dropped rather than aggregated under a
68
+ missing label, matching R, and the caller is warned because a partial
69
+ label table is far more often a typo than an intent.
70
+ """
71
+ keys = [x for x in by_frame_keys(by) if x in estimates.columns]
72
+ if not keys:
73
+ raise ValueError(
74
+ "None of the columns in the `by` data frame match the estimates. "
75
+ f"Available columns: {', '.join(estimates.columns)}."
76
+ )
77
+ labels = harmonize_by_types(estimates, by.select([*keys, "by"])).unique(subset=keys)
78
+ if "by" in estimates.columns:
79
+ estimates = estimates.drop("by")
80
+ out = estimates.join(labels, on=keys, how="left")
81
+ keep = out["by"].is_not_null()
82
+ if verbose and not keep.all():
83
+ warnings.warn(BY_FRAME_COVERAGE_WARNING, UserWarning, stacklevel=2)
84
+ return out, keep
85
+
86
+
87
+ def get_by(model, estimand, newdata, by=None, wts=None):
88
+ out, _ = _get_by_internal(
89
+ model=model,
90
+ estimand=estimand,
91
+ newdata=newdata,
92
+ by=by,
93
+ wts=wts,
94
+ )
95
+ return out
96
+
97
+
98
+ def get_by_plan(model, estimand, newdata, by=None, wts=None):
99
+ out, agg_groups = _get_by_internal(
100
+ model=model,
101
+ estimand=estimand,
102
+ newdata=newdata,
103
+ by=by,
104
+ wts=wts,
105
+ return_plan_groups=True,
106
+ )
107
+ return out, agg_groups
108
+
109
+
110
+ def _get_by_internal(
111
+ model,
112
+ estimand,
113
+ newdata,
114
+ by=None,
115
+ wts=None,
116
+ return_plan_groups: bool = False,
117
+ ) -> tuple[pl.DataFrame, list[AggGroup] | None]:
118
+ # for predictions
119
+ if (
120
+ isinstance(by, list)
121
+ and len(by) == 1
122
+ and by[0] == "group"
123
+ and "group" not in estimand.columns
124
+ ):
125
+ by = True
126
+
127
+ agg_groups: list[AggGroup] | None = None
128
+
129
+ if by_is_frame(by):
130
+ return _get_by_frame(
131
+ estimand=estimand,
132
+ newdata=newdata,
133
+ by=by,
134
+ wts=wts,
135
+ return_plan_groups=return_plan_groups,
136
+ )
137
+
138
+ if by is True:
139
+ # A grand mean must respect user weights exactly like the named-group
140
+ # path below does; ignoring `wts` here would silently report a
141
+ # different estimand than the one requested, with a Jacobian to match.
142
+ w_arr = None
143
+ if wts is not None:
144
+ tmp = estimand
145
+ if (
146
+ wts not in tmp.columns
147
+ and "rowid" in tmp.columns
148
+ and "rowid" in newdata.columns
149
+ ):
150
+ # The recorded weights must line up with the Jacobian rows,
151
+ # which follow the estimand's row order. A join does not
152
+ # promise to preserve that order, so carry an explicit
153
+ # position column through it and restore the order afterward.
154
+ tmp = tmp.with_columns(
155
+ pl.Series(
156
+ "_marginaleffects_wpos",
157
+ range(tmp.height),
158
+ dtype=pl.Int32,
159
+ )
160
+ )
161
+ tmp = tmp.join(newdata.select(["rowid", wts]), on="rowid", how="left")
162
+ tmp = tmp.sort("_marginaleffects_wpos").drop("_marginaleffects_wpos")
163
+ if wts not in tmp.columns:
164
+ raise ValueError(f"Weights column '{wts}' not found for aggregation.")
165
+ if tmp.height != estimand.height:
166
+ raise ValueError(
167
+ "Joining the weights column changed the number of rows; "
168
+ "`newdata` has duplicated `rowid` values."
169
+ )
170
+ if tmp[wts].null_count() > 0:
171
+ raise ValueError(
172
+ f"Weights column '{wts}' has missing values after joining "
173
+ "on `rowid`."
174
+ )
175
+ result = tmp.select(_weighted_mean_expr(wts))
176
+ w_arr = tmp[wts].cast(pl.Float64).to_numpy()
177
+ else:
178
+ result = estimand.select(_mean_expr())
179
+ if return_plan_groups:
180
+ agg_groups = [
181
+ AggGroup(idx=pl.Series(range(estimand.height)).to_numpy(), w=w_arr)
182
+ ]
183
+ return result, agg_groups
184
+ elif by is False:
185
+ return estimand, agg_groups
186
+
187
+ if "group" in estimand.columns:
188
+ by = ["group"] + by
189
+
190
+ if return_plan_groups:
191
+ estimand = estimand.with_columns(
192
+ pl.Series("_marginaleffects_pos", range(estimand.height), dtype=pl.Int32)
193
+ )
194
+
195
+ if "rowid" in estimand.columns and "rowid" in newdata.columns:
196
+ out = estimand.join(newdata, on="rowid", how="left")
197
+ else:
198
+ out = pl.DataFrame({"estimate": estimand["estimate"]})
199
+
200
+ by = list(dict.fromkeys(x for x in by if x in out.columns))
201
+
202
+ if isinstance(by, list) and len(by) == 0:
203
+ return out, agg_groups
204
+
205
+ agg_exprs: list[pl.Expr] = []
206
+ if wts is None:
207
+ agg_exprs.append(_mean_expr())
208
+ else:
209
+ agg_exprs.append(_weighted_mean_expr(wts))
210
+ if return_plan_groups and "_marginaleffects_pos" in out.columns:
211
+ agg_exprs.append(pl.col("_marginaleffects_pos").alias("_positions"))
212
+ if wts is not None:
213
+ agg_exprs.append(pl.col(wts).cast(pl.Float64).alias("_plan_weights"))
214
+
215
+ out = out.group_by(by, maintain_order=True).agg(agg_exprs)
216
+
217
+ # Sort by 'by' columns ONLY if they are Enum type to ensure consistent categorical ordering
218
+ # For Enum columns, sort() respects the enum order (not lexical order)
219
+ # For other types (strings, numbers), maintain the group_by order to preserve existing behavior
220
+ if isinstance(by, str):
221
+ by_cols = [by]
222
+ else:
223
+ by_cols = list(by)
224
+
225
+ should_sort = any(
226
+ out[col].dtype == pl.Enum for col in by_cols if col in out.columns
227
+ )
228
+ if should_sort:
229
+ out = out.sort(by)
230
+
231
+ if return_plan_groups and "_positions" in out.columns:
232
+ weights = (
233
+ out["_plan_weights"].to_list()
234
+ if wts is not None and "_plan_weights" in out.columns
235
+ else None
236
+ )
237
+ agg_groups = []
238
+ for i, positions in enumerate(out["_positions"].to_list()):
239
+ w = None if weights is None else pl.Series(weights[i]).to_numpy()
240
+ agg_groups.append(AggGroup(idx=pl.Series(positions).to_numpy(), w=w))
241
+ drop_cols = ["_positions"]
242
+ if "_plan_weights" in out.columns:
243
+ drop_cols.append("_plan_weights")
244
+ out = out.drop(drop_cols)
245
+
246
+ return out, agg_groups
247
+
248
+
249
+ def _get_by_frame(estimand, newdata, by, wts, return_plan_groups: bool):
250
+ """Aggregate estimates under labels supplied as a `by` data frame.
251
+
252
+ The label table can map several predictor combinations onto one label, so
253
+ grouping happens on the joined label rather than on the estimates' own
254
+ columns. Row positions are recorded before the join because the plan's
255
+ prediction vector is indexed by the *unfiltered* estimand order, which
256
+ dropping uncovered rows would otherwise shift.
257
+ """
258
+ position = "_marginaleffects_pos"
259
+ estimand = estimand.with_columns(
260
+ pl.Series(position, range(estimand.height), dtype=pl.Int64)
261
+ )
262
+
263
+ keys = by_frame_keys(by)
264
+ missing = [x for x in keys if x not in estimand.columns]
265
+ if missing and "rowid" in estimand.columns and "rowid" in newdata.columns:
266
+ payload = [x for x in missing if x in newdata.columns]
267
+ if payload:
268
+ estimand = estimand.join(
269
+ newdata.select(["rowid", *payload]), on="rowid", how="left"
270
+ )
271
+
272
+ labelled, keep = attach_by_labels(estimand, by)
273
+ labelled = labelled.filter(keep)
274
+
275
+ group_cols = [x for x in ("term", "by") if x in labelled.columns]
276
+ weighted = wts is not None and wts in labelled.columns
277
+ agg_exprs = [_weighted_mean_expr(wts) if weighted else _mean_expr()]
278
+ if return_plan_groups:
279
+ agg_exprs.append(pl.col(position).alias("_positions"))
280
+ if weighted:
281
+ agg_exprs.append(pl.col(wts).cast(pl.Float64).alias("_plan_weights"))
282
+
283
+ out = labelled.group_by(group_cols, maintain_order=True).agg(agg_exprs)
284
+
285
+ agg_groups = None
286
+ if return_plan_groups:
287
+ weights = out["_plan_weights"].to_list() if weighted else None
288
+ agg_groups = [
289
+ AggGroup(
290
+ idx=pl.Series(positions).to_numpy(),
291
+ w=None if weights is None else pl.Series(weights[i]).to_numpy(),
292
+ )
293
+ for i, positions in enumerate(out["_positions"].to_list())
294
+ ]
295
+ out = out.drop([x for x in ("_positions", "_plan_weights") if x in out.columns])
296
+
297
+ return out.select([*group_cols, "estimate"]), agg_groups
@@ -1,18 +1,20 @@
1
- from .result import MarginaleffectsResult, MarginaleffectsDataFrame
1
+ from .result import MarginaleffectsDataFrame, MarginaleffectsResult
2
2
 
3
3
  __all__ = [
4
- "MarginaleffectsResult",
5
4
  "MarginaleffectsDataFrame",
5
+ "MarginaleffectsResult",
6
6
  "ModelAbstract",
7
+ "ModelAdapter",
7
8
  "ModelVault",
8
9
  ]
9
10
 
10
11
 
11
12
  def __getattr__(name):
12
- if name in ("ModelAbstract", "ModelVault"):
13
- from .model import ModelAbstract, ModelVault
13
+ if name in ("ModelAbstract", "ModelAdapter", "ModelVault"):
14
+ from .model import ModelAbstract, ModelAdapter, ModelVault
14
15
 
15
16
  globals()["ModelAbstract"] = ModelAbstract
17
+ globals()["ModelAdapter"] = ModelAdapter
16
18
  globals()["ModelVault"] = ModelVault
17
- return ModelAbstract if name == "ModelAbstract" else ModelVault
19
+ return globals()[name]
18
20
  raise AttributeError(f"module 'marginaleffects.classes' has no attribute {name!r}")
@@ -1,36 +1,50 @@
1
1
  from abc import ABC, abstractmethod
2
2
  from dataclasses import dataclass, field
3
- from typing import Any, Dict, List, Optional
3
+ from typing import Any, Protocol, runtime_checkable
4
4
 
5
5
  import numpy as np
6
6
  import polars as pl
7
7
 
8
- from ..sanitize.validation import ModelValidation
9
8
  from .. import formula as fml
9
+ from ..sanitize.validation import ModelValidation
10
10
 
11
11
 
12
12
  @dataclass
13
13
  class ModelVault:
14
14
  """Typed container for model metadata shared across all adapters."""
15
15
 
16
- coef: Optional[np.ndarray] = None
17
- coefnames: Optional[np.ndarray] = None
18
- formula: Optional[str] = None
16
+ coef: np.ndarray | None = None
17
+ coefnames: np.ndarray | None = None
18
+ formula: str | None = None
19
19
  formula_engine: str = "formulaic"
20
- modeldata: Optional[pl.DataFrame] = None
20
+ modeldata: pl.DataFrame | None = None
21
21
  package: str = "unknown"
22
- vcov: Optional[np.ndarray] = None
23
- variables_type: Dict[str, str] = field(default_factory=dict)
24
- variable_names: Optional[List[str]] = None
25
- engine_running: Optional[Any] = None
22
+ vcov: np.ndarray | None = None
23
+ variables_type: dict[str, str] = field(default_factory=dict)
24
+ variable_names: list[str] | None = None
25
+ engine_running: Any | None = None
26
26
  # statsmodels-specific
27
- design_info_patsy: Optional[Any] = None
28
- pandas_categorical_orders: Dict[str, list] = field(default_factory=dict)
27
+ design_info_patsy: Any | None = None
28
+ pandas_categorical_orders: dict[str, list] = field(default_factory=dict)
29
29
  # sklearn-specific
30
- model_spec: Optional[Any] = None
31
- original_columns: Optional[List[str]] = None
30
+ model_spec: Any | None = None
31
+ original_columns: list[str] | None = None
32
32
  # linearmodels-specific
33
- multiindex: Optional[List[str]] = None
33
+ multiindex: list[str] | None = None
34
+
35
+
36
+ @runtime_checkable
37
+ class ModelAdapter(Protocol):
38
+ """Explicit contract consumed by marginaleffects estimation code."""
39
+
40
+ def get_modeldata(self) -> pl.DataFrame: ...
41
+ def get_vcov(self, vcov=True) -> np.ndarray | None: ...
42
+ def get_coef(self) -> np.ndarray: ...
43
+ def get_coefnames(self) -> np.ndarray: ...
44
+ def get_formula(self): ...
45
+ def get_exog(self, newdata: pl.DataFrame): ...
46
+ def get_predict(self, params: np.ndarray, newdata) -> pl.DataFrame: ...
47
+ def get_df(self) -> float: ...
34
48
 
35
49
 
36
50
  class ModelAbstract(ModelValidation, ABC):
@@ -39,22 +53,22 @@ class ModelAbstract(ModelValidation, ABC):
39
53
  self.vault = vault
40
54
  self.validation()
41
55
 
42
- def get_modeldata(self) -> Optional[pl.DataFrame]:
56
+ def get_modeldata(self) -> pl.DataFrame | None:
43
57
  return self.vault.modeldata
44
58
 
45
- def get_vcov(self, vcov=False) -> Optional[np.ndarray]:
59
+ def get_vcov(self, vcov=False) -> np.ndarray | None:
46
60
  return self.vault.vcov
47
61
 
48
- def get_coef(self) -> Optional[np.ndarray]:
62
+ def get_coef(self) -> np.ndarray | None:
49
63
  return self.vault.coef
50
64
 
51
- def get_coefnames(self) -> Optional[np.ndarray]:
65
+ def get_coefnames(self) -> np.ndarray | None:
52
66
  return self.vault.coefnames
53
67
 
54
- def get_engine_running(self) -> Optional[Any]:
68
+ def get_engine_running(self) -> Any | None:
55
69
  return self.vault.engine_running
56
70
 
57
- def get_formula(self) -> Optional[str]:
71
+ def get_formula(self) -> str | None:
58
72
  return self.vault.formula
59
73
 
60
74
  def get_formula_engine(self) -> str:
@@ -63,7 +77,7 @@ class ModelAbstract(ModelValidation, ABC):
63
77
  def get_package(self) -> str:
64
78
  return self.vault.package
65
79
 
66
- def get_variable_type(self, name=None) -> Dict[str, str]:
80
+ def get_variable_type(self, name=None) -> dict[str, str]:
67
81
  variables = self.vault.variables_type
68
82
  if isinstance(name, str) and name in variables:
69
83
  return variables[name]
@@ -73,7 +87,7 @@ class ModelAbstract(ModelValidation, ABC):
73
87
  def set_variable_type(self, name, value):
74
88
  self.vault.variables_type[name] = value
75
89
 
76
- def find_variables(self) -> Optional[List[str]]:
90
+ def find_variables(self) -> list[str] | None:
77
91
  if self.vault.variable_names is not None:
78
92
  return self.vault.variable_names
79
93
 
@@ -87,14 +101,14 @@ class ModelAbstract(ModelValidation, ABC):
87
101
 
88
102
  return out
89
103
 
90
- def find_response(self) -> Optional[str]:
104
+ def find_response(self) -> str | None:
91
105
  vars = self.find_variables()
92
106
  if vars is None:
93
107
  return None
94
108
  else:
95
109
  return vars[0]
96
110
 
97
- def find_predictors(self) -> Optional[List[str]]:
111
+ def find_predictors(self) -> list[str] | None:
98
112
  vars = self.find_variables()
99
113
  if vars is None:
100
114
  return None
@@ -127,13 +141,25 @@ class ModelAbstract(ModelValidation, ABC):
127
141
  def get_predict(self, params: np.ndarray, newdata) -> pl.DataFrame:
128
142
  pass
129
143
 
130
- def get_autodiff_args(self):
131
- return None
144
+ def get_analytic_args(self):
145
+ """Return adapter metadata needed by the analytic Jacobian engine."""
146
+ return
147
+
148
+ def get_link_functions(self):
149
+ """Return inverse-link and derivative callables for response predictions."""
150
+ return
151
+
152
+ def get_exog_names(self, value):
153
+ """Return design-column names, or ``None`` when they are unavailable."""
154
+ design_info = getattr(value, "design_info", None)
155
+ columns = getattr(design_info, "column_names", None)
156
+ if columns is None:
157
+ columns = getattr(value, "columns", None)
158
+ return None if columns is None else list(columns)
159
+
160
+ def get_df(self) -> float:
161
+ return np.inf
132
162
 
133
- def __getattr__(self, name: str) -> Any:
134
- """Forward attribute access to the underlying fitted model."""
135
- try:
136
- return object.__getattribute__(self, name)
137
- except AttributeError:
138
- # Forward to the wrapped model
139
- return getattr(self.model, name)
163
+ def get_fitted_model(self) -> Any:
164
+ """Return the wrapped engine object for adapter-specific integrations."""
165
+ return self.model
@@ -1,8 +1,9 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  import warnings
4
+ from collections.abc import Iterable, Iterator
4
5
  from dataclasses import dataclass
5
- from typing import Any, Dict, Iterable, Iterator, Optional
6
+ from typing import Any
6
7
 
7
8
  import polars as pl
8
9
 
@@ -10,7 +11,7 @@ import polars as pl
10
11
  def _default_mapping(
11
12
  data: pl.DataFrame,
12
13
  datagrid_explicit: Iterable[str],
13
- ) -> Dict[str, str]:
14
+ ) -> dict[str, str]:
14
15
  datagrid_explicit = list(datagrid_explicit)
15
16
  datagrid_map = dict(zip(datagrid_explicit, datagrid_explicit))
16
17
  contrast_columns = {
@@ -18,7 +19,7 @@ def _default_mapping(
18
19
  for col in data.columns
19
20
  if col.startswith("contrast_")
20
21
  }
21
- mapping: Dict[str, str] = {
22
+ mapping: dict[str, str] = {
22
23
  "term": "Term",
23
24
  "group": "Group",
24
25
  **datagrid_map,
@@ -40,12 +41,13 @@ def _default_mapping(
40
41
 
41
42
  @dataclass
42
43
  class _Metadata:
43
- by: Optional[Any]
44
+ by: Any | None
44
45
  conf_level: float
45
- jacobian: Optional[Any]
46
+ jacobian: Any | None
47
+ jacobian_method: str | None
46
48
  datagrid_explicit: Iterable[str]
47
49
  print_head: str
48
- mapping: Dict[str, str]
50
+ mapping: dict[str, str]
49
51
 
50
52
 
51
53
  class MarginaleffectsResult:
@@ -58,11 +60,12 @@ class MarginaleffectsResult:
58
60
  self,
59
61
  data: pl.DataFrame,
60
62
  *,
61
- by: Optional[Any] = None,
63
+ by: Any | None = None,
62
64
  conf_level: float = 0.95,
63
- jacobian: Optional[Any] = None,
64
- newdata: Optional[Any] = None,
65
- mapping: Optional[Dict[str, str]] = None,
65
+ jacobian: Any | None = None,
66
+ jacobian_method: str | None = None,
67
+ newdata: Any | None = None,
68
+ mapping: dict[str, str] | None = None,
66
69
  print_head: str = "",
67
70
  ) -> None:
68
71
  if not isinstance(data, pl.DataFrame):
@@ -82,6 +85,7 @@ class MarginaleffectsResult:
82
85
  by=by,
83
86
  conf_level=conf_level,
84
87
  jacobian=jacobian,
88
+ jacobian_method=jacobian_method,
85
89
  datagrid_explicit=list(datagrid_explicit),
86
90
  print_head=print_head,
87
91
  mapping=final_mapping,
@@ -96,7 +100,7 @@ class MarginaleffectsResult:
96
100
  return self._data
97
101
 
98
102
  @property
99
- def by(self) -> Optional[Any]:
103
+ def by(self) -> Any | None:
100
104
  return self._meta.by
101
105
 
102
106
  @property
@@ -104,9 +108,14 @@ class MarginaleffectsResult:
104
108
  return self._meta.conf_level
105
109
 
106
110
  @property
107
- def jacobian(self) -> Optional[Any]:
111
+ def jacobian(self) -> Any | None:
108
112
  return self._meta.jacobian
109
113
 
114
+ @property
115
+ def jacobian_method(self) -> str | None:
116
+ """Method used to differentiate the coefficient-to-estimand map."""
117
+ return self._meta.jacobian_method
118
+
110
119
  @property
111
120
  def datagrid_explicit(self) -> Iterable[str]:
112
121
  return self._meta.datagrid_explicit
@@ -205,6 +214,7 @@ class MarginaleffectsResult:
205
214
  by=self.by,
206
215
  conf_level=self.conf_level,
207
216
  jacobian=self.jacobian,
217
+ jacobian_method=self.jacobian_method,
208
218
  mapping=self._meta.mapping.copy(),
209
219
  print_head=self._meta.print_head,
210
220
  )