Update pixi.toml and pyproject.toml (#1209)

Remove upper version limit for scipy and JAX, remove jaxlib as separate dependency. Close #1131

The previous restriction scipy<1.16 is not necessary when lifting the restriction jax = ">=0.4.38, <0.8".

Author: leostimpfle

pixi.lock

@@ -2,7 +2,7 @@
 authors = [
     "Alexander Fischer <alexander-fischer1801@t-online.de>",
     "Styfen Schär"
-    ]
+]
 description = "Fast high dimensional fixed effect estimation following syntax of the fixest R package."
 channels = ["conda-forge"]
 name = "pyfixest"
@@ -13,16 +13,16 @@ version = "0.50.0a1"
 
 
 [dependencies]
-scipy=">=1.6,<1.16"
-formulaic=">=1.1.0"
-pandas=">=1.1.0"
-numba=">=0.58.0"
-seaborn=">=0.13.2"
-tabulate=">=0.9.0"
-tqdm=">=4.0.0"
-numpy=">=1.25.2"
-narwhals=">=1.13.3"
-joblib=">=1.4.2,<2"
+scipy = ">=1.6"
+formulaic = ">=1.1.0"
+pandas = ">=1.1.0"
+numba = ">=0.58.0"
+seaborn = ">=0.13.2"
+tabulate = ">=0.9.0"
+tqdm = ">=4.0.0"
+numpy = ">=1.25.2"
+narwhals = ">=1.13.3"
+joblib = ">=1.4.2"
 
 [pypi-dependencies]
 maketables = ">=0.1.0"
@@ -45,10 +45,9 @@ duckdb = ">=1.1.3"
 polars = ">=1.0.0"
 pyarrow = ">=14.0"
 rpy2 = "==3.5.11"
-pixi-pycharm = ">=0.0.8,<0.0.9"
 r-base = "*"
 r-fixest = ">=0.13.2,<0.14"
-r-stringmagic=">=1.2.0"
+r-stringmagic = ">=1.2.0"
 r-sandwich = ">=3.0_2,<4"
 r-broom = ">=1.0.5,<2"
 r-car = ">=3.1_2,<4"
@@ -58,13 +57,12 @@ pytest-sugar = ">=1.1.1,<2"
 
 [feature.dev.pypi-dependencies]
 wildboottest = ">=0.3.2, <0.4"
-jax = ">=0.4.38, <0.8"
-jaxlib = ">=0.4.38, <0.8"
+jax = ">=0.4.38"
 
 [feature.dev.tasks]
 "tests" = "pytest -rs -n 9 --cov-report=term tests"
 "tests-against-r-core" = 'pytest -rs tests -n 9 -m "against_r_core" --cov=pyfixest --cov-report=xml'
-"tests-against-r-extended"  = 'pytest -rs tests -n 9 -m "against_r_extended" --cov=pyfixest --cov-report=xml'
+"tests-against-r-extended" = 'pytest -rs tests -n 9 -m "against_r_extended" --cov=pyfixest --cov-report=xml'
 "tests-regular" = 'pytest tests -n 9 -m "not (extended or against_r_core or against_r_extended or plots or hac)" --cov=pyfixest --cov-report=xml'
 "tests-extended" = 'pytest tests -n 9 -m "extended" --cov=pyfixest --cov-report=xml'
 "tests-fixest" = "pytest -rs tests/test_vs_fixest.py -n 9 --cov=pyfixest --cov-report=xml"
@@ -80,19 +78,19 @@ jupytext = "==1.16.1"
 watermark = "==2.4.3"
 pylatex = ">=1.4.2,<2"
 pyarrow = ">=14.0"
-quartodoc =">=0.11.0,<0.12"
+quartodoc = ">=0.11.0,<0.12"
 rpy2 = "==3.5.11"
 r-base = "*"
 r-fixest = ">=0.13.2,<0.14"
-r-stringmagic=">=1.2.0"
+r-stringmagic = ">=1.2.0"
 r-broom = ">=1.0.5,<2"
 quarto = ">=1.8.25,<2"
 
 [feature.docs.pypi-dependencies]
 wildboottest = ">=0.3.2"
 causaldata = ">=0.1.4,<0.2"
-marginaleffects= ">=0.0.10"
-lets-plot =">=4.0.0"
+marginaleffects = ">=0.0.10"
+lets-plot = ">=4.0.0"
 
 [feature.docs.tasks]
 docs-build = "quartodoc build --verbose --config docs/_quarto.yml"
@@ -102,8 +100,7 @@ docs-preview = { cmd = "quarto preview docs", env = { PYTHONWARNINGS = "ignore",
 [feature.jax.dependencies]
 
 [feature.jax.pypi-dependencies]
-jax = ">=0.4.15"
-jaxlib= ">=0.4.15"
+jax = ">=0.4.29"
 
 [feature.plots.pypi-dependencies]
 lets-plot = ">=4.0.0"

pixi.toml

@@ -139,7 +139,7 @@ def _get_fixed_effects_prediction_component(
                 fdict.update({x: 0 for x in omitted_cat})
             fixef_dicts[f"C({f})"] = fdict
 
-        _fixef_mat = _apply_fixef_numpy(df_fe.values, fixef_dicts)
+        _fixef_mat = _apply_fixef_numpy(df_fe.values.astype(str), fixef_dicts)
         fe_hat += np.sum(_fixef_mat, axis=1)
 
     return fe_hat

pyfixest/estimation/post_estimation/prediction.py

@@ -10,7 +10,7 @@ readme = "README.md"
 requires-python = ">=3.9"
 license = { text = "MIT" }
 dependencies = [
-    "scipy>=1.6,<1.16",
+    "scipy>=1.6",
     "formulaic>=1.1.0",
     "pandas>=1.1.0",
     "numba>=0.58.0",
@@ -20,13 +20,12 @@ dependencies = [
     "maketables>=0.1.0",
     "numpy>=1.25.2",
     "narwhals>=1.13.3",
-    "joblib>=1.4.2,<2",
+    "joblib>=1.4.2",
 ]
 
 [project.optional-dependencies]
 jax = [
     "jax>=0.4.15",
-    "jaxlib>=0.4.15",
 ]
 plots = [
     "lets-plot>=4.0.0",
@@ -43,25 +42,25 @@ python-source = "."
 
 [tool.pytest.ini_options]
 addopts = [
-  "-v",
-  "--strict-markers",
-  "--strict-config",
-  "--cov=pyfixest",
-  "--cov-report=term-missing",
-  "--color=yes",
+    "-v",
+    "--strict-markers",
+    "--strict-config",
+    "--cov=pyfixest",
+    "--cov-report=term-missing",
+    "--color=yes",
 ]
 
 markers = [
-  "against_r_core: mark test to be part of the test suite that depends on conda available R modules",
-  "against_r_extended: mark test to be part of the test suite that depends on other R modules",
-  "extended: mark test to be part of the extended test suite",
-  "plots: marks all tests for plotting functionality, tests only triggered when using tag in github issue",
-  "hac: marks all tests for HAC SEs"
+    "against_r_core: mark test to be part of the test suite that depends on conda available R modules",
+    "against_r_extended: mark test to be part of the test suite that depends on other R modules",
+    "extended: mark test to be part of the extended test suite",
+    "plots: marks all tests for plotting functionality, tests only triggered when using tag in github issue",
+    "hac: marks all tests for HAC SEs"
 ]
 
 filterwarnings = [
-  "ignore::FutureWarning:rpy2",
-  "ignore::DeprecationWarning:rpy2",
+    "ignore::FutureWarning:rpy2",
+    "ignore::DeprecationWarning:rpy2",
 ]
 
 [tool.ruff]
@@ -74,37 +73,37 @@ target-version = "py39"
 [tool.ruff.lint]
 # docs: https://docs.astral.sh/ruff/rules/
 select = [
-  "B",   # bugbear
-  "F",   # Pyflakes
-  "E",   # pycodestyle errors
-  "W",   # pycodestyle warnings
-  "I",   # isort
-  "D",   # flake8-docstrings
-  "UP",  # pyupgrade
-  "SIM", # flake8-simplify
-  "RUF", # ruff,
-  "TRY", # tryceratops
+    "B", # bugbear
+    "F", # Pyflakes
+    "E", # pycodestyle errors
+    "W", # pycodestyle warnings
+    "I", # isort
+    "D", # flake8-docstrings
+    "UP", # pyupgrade
+    "SIM", # flake8-simplify
+    "RUF", # ruff,
+    "TRY", # tryceratops
 ]
 
 ignore = [
-  # do not enable if formatting
-  # docs: https://docs.astral.sh/ruff/formatter/#conflicting-lint-rules
-  "B028",   # No explicit `stacklevel` keyword argument found
-  "B904",   # Better exception handling
-  "W191",   # tab indentation
-  "E111",   # indentation
-  "E114",   # indentation
-  "E117",   # over indented
-  "D206",   # indent with spaces
-  "D300",   # triple single quotes
-  "E501",   # line length regulated by formatter
-  "D105",   # missing docstring in magic method
-  "D100",   # missing docstring in public module
-  "D104",   # missing docstring in public package
-  "SIM110", # Use all instead of `for` loop
-  "TRY003", # Avoid specifying long messages outside the exception class
-  "D205",   # 1 blank line required between summary line and description
-  "W505",   # Doc line too long
+    # do not enable if formatting
+    # docs: https://docs.astral.sh/ruff/formatter/#conflicting-lint-rules
+    "B028", # No explicit `stacklevel` keyword argument found
+    "B904", # Better exception handling
+    "W191", # tab indentation
+    "E111", # indentation
+    "E114", # indentation
+    "E117", # over indented
+    "D206", # indent with spaces
+    "D300", # triple single quotes
+    "E501", # line length regulated by formatter
+    "D105", # missing docstring in magic method
+    "D100", # missing docstring in public module
+    "D104", # missing docstring in public package
+    "SIM110", # Use all instead of `for` loop
+    "TRY003", # Avoid specifying long messages outside the exception class
+    "D205", # 1 blank line required between summary line and description
+    "W505", # Doc line too long
 ]
 
 [tool.ruff.lint.per-file-ignores]
@@ -129,16 +128,16 @@ docstring-code-line-length = 88
 
 [[tool.mypy.overrides]]
 module = [
-  "pandas.*",
-  "numpy.*",
-  "scipy.*",
-  "numba.*",
-  "lets_plot.*",
-  "formulaic.*",
-  "wildboottest.*",
-  "tabulate.*",
-  "joblib.*",
-  "narwhals.*",
-  "tqdm.*",
+    "pandas.*",
+    "numpy.*",
+    "scipy.*",
+    "numba.*",
+    "lets_plot.*",
+    "formulaic.*",
+    "wildboottest.*",
+    "tabulate.*",
+    "joblib.*",
+    "narwhals.*",
+    "tqdm.*",
 ]
 ignore_missing_imports = true

pyproject.toml

@@ -76,7 +76,9 @@ def test_event_study(data):
         r_df = pd.DataFrame(r_df)
 
     np.testing.assert_allclose(fit_did2s.coef(), r_df[1], atol=1e-05, rtol=1e-05)
-    np.testing.assert_allclose(fit_did2s.se(), float(r_df[2]), atol=1e-05, rtol=1e-05)
+    np.testing.assert_allclose(
+        fit_did2s.se(), r_df[2].squeeze(), atol=1e-05, rtol=1e-05
+    )
 
 
 @pytest.mark.against_r_core
@@ -119,7 +121,7 @@ def test_did2s(data, weights):
     )
 
     np.testing.assert_allclose(
-        fit_did2s_py1.se(), float(fit_did2s_r1.iloc[:, 3]), atol=1e-05, rtol=1e-05
+        fit_did2s_py1.se(), fit_did2s_r1.iloc[:, 3].squeeze(), atol=1e-05, rtol=1e-05
     )
 
     # Model 2