Allow setting backend at runtime

conftest.py

@@ -0,0 +1,43 @@
+import pytest
+
+from ttsim.interface_dag_elements.backend import dnp as ttsim_dnp
+from ttsim.interface_dag_elements.backend import xnp as ttsim_xnp
+
+
+# content of conftest.py
+def pytest_addoption(parser):
+    parser.addoption(
+        "--backend",
+        action="store",
+        default="numpy",
+        help="The backend to test against (e.g., --backend=numpy --backend=jax)",
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "skipif_jax: skip test if backend is jax")
+
+
+@pytest.fixture
+def backend(request):
+    backend = request.config.getoption("--backend")
+    return backend
+
+
+@pytest.fixture
+def xnp(request):
+    backend = request.config.getoption("--backend")
+    return ttsim_xnp(backend)
+
+
+@pytest.fixture
+def dnp(request):
+    backend = request.config.getoption("--backend")
+    return ttsim_dnp(backend)
+
+
+@pytest.fixture(autouse=True)
+def skipif_jax(request, backend):
+    """Automatically skip tests marked with skipif_jax when backend is jax."""
+    if request.node.get_closest_marker("skipif_jax") and backend == "jax":
+        pytest.skip("Cannot run this test with Jax")

pyproject.toml

@@ -150,6 +150,9 @@ python = "3.12.*"
 jax = ">=0.4.20"
 jaxlib = ">=0.4.20"
 
+[tool.pixi.feature.jax.pypi-dependencies]
+jax-datetime = { git = "https://github.com/google/jax-datetime.git" }
+
 [tool.pixi.feature.jax.target.win-64.pypi-dependencies]
 jax = { version = ">=0.4.20", extras = ["cpu"] }
 jaxlib = ">=0.4.20"
@@ -165,6 +168,9 @@ types-pytz = "*"
 [tool.pixi.feature.test.tasks]
 tests = "pytest"
 
+[tool.pixi.feature.jax.tasks]
+tests-jax = "pytest --backend=jax"
+
 [tool.pixi.feature.mypy.tasks]
 mypy = "mypy --ignore-missing-imports"
 

src/_gettsim/arbeitslosengeld_2/einkommen.py

@@ -14,6 +14,8 @@
 )
 
 if TYPE_CHECKING:
+    from types import ModuleType
+
     from ttsim.interface_dag_elements.typing import RawParam
 
 
@@ -147,12 +149,14 @@
     einkommensteuer__einkünfte__aus_nichtselbstständiger_arbeit__bruttolohn_m: float,
     nettoquote: float,
     parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg: PiecewisePolynomialParamValue,
+    xnp: ModuleType,
 ) -> float:
     """Share of income which remains to the individual."""
     return piecewise_polynomial(
         x=einkommensteuer__einkünfte__aus_nichtselbstständiger_arbeit__bruttolohn_m,
         parameters=parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg,
         rates_multiplier=nettoquote,
+        xnp=xnp,
     )
 
 
@@ -164,6 +168,7 @@
     einkommensteuer__anzahl_kinderfreibeträge: int,
     parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg: PiecewisePolynomialParamValue,
     parameter_anrechnungsfreies_einkommen_mit_kindern_in_bg: PiecewisePolynomialParamValue,
+    xnp: ModuleType,
 ) -> float:
     """Calculate share of income, which remains to the individual since 10/2005.
 
@@ -184,32 +189,37 @@
         out = piecewise_polynomial(
             x=eink_erwerbstätigkeit,
             parameters=parameter_anrechnungsfreies_einkommen_mit_kindern_in_bg,
+            xnp=xnp,
         )
     else:
         out = piecewise_polynomial(
             x=eink_erwerbstätigkeit,
             parameters=parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg,
+            xnp=xnp,
         )
     return out
 
 
 @param_function(start_date="2005-01-01")
 def parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg(
     raw_parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg: RawParam,
+    xnp: ModuleType,
 ) -> PiecewisePolynomialParamValue:
     """Parameter for calculation of income not subject to transfer withdrawal when
     children are not in the Bedarfsgemeinschaft."""
     return get_piecewise_parameters(
         leaf_name="parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg",
         func_type="piecewise_linear",
         parameter_dict=raw_parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg,
+        xnp=xnp,
     )
 
 
 @param_function(start_date="2005-10-01")
 def parameter_anrechnungsfreies_einkommen_mit_kindern_in_bg(
     raw_parameter_anrechnungsfreies_einkommen_mit_kindern_in_bg: RawParam,
     raw_parameter_anrechnungsfreies_einkommen_ohne_kinder_in_bg: RawParam,
+    xnp: ModuleType,
 ) -> PiecewisePolynomialParamValue:
     """Parameter for calculation of income not subject to transfer withdrawal when
     children are in the Bedarfsgemeinschaft."""
@@ -221,4 +231,5 @@
         leaf_name="parameter_anrechnungsfreies_einkommen_mit_kindern_in_bg",
         func_type="piecewise_linear",
         parameter_dict=updated_parameters,
+        xnp=xnp,
     )

src/_gettsim/arbeitslosengeld_2/kindergeldübertrag.py

@@ -7,7 +7,9 @@
 from ttsim.tt_dag_elements import AggType, agg_by_p_id_function, join, policy_function
 
 if TYPE_CHECKING:
-    from ttsim.config import numpy_or_jax as np
+    from types import ModuleType
+
+    import numpy
 
 
 @agg_by_p_id_function(start_date="2005-01-01", agg_type=AggType.SUM)
@@ -57,9 +59,10 @@
 @policy_function(start_date="2005-01-01", vectorization_strategy="not_required")
 def kindergeld_zur_bedarfsdeckung_m(
     kindergeld_pro_kind_m: float,
-    kindergeld__p_id_empfänger: np.ndarray,  # int
-    p_id: np.ndarray,  # int
-) -> np.ndarray:  # float
+    kindergeld__p_id_empfänger: numpy.ndarray,  # int
+    p_id: numpy.ndarray,  # int
+    xnp: ModuleType,
+) -> numpy.ndarray:  # float
     """Kindergeld that is used to cover the SGB II Regelbedarf of the child.
 
     Even though the Kindergeld is paid to the parent (see function
@@ -71,10 +74,11 @@
 
     """
     return join(
-        kindergeld__p_id_empfänger,
-        p_id,
-        kindergeld_pro_kind_m,
+        foreign_key=kindergeld__p_id_empfänger,
+        primary_key=p_id,
+        target=kindergeld_pro_kind_m,
         value_if_foreign_key_is_missing=0.0,
+        xnp=xnp,
     )
 
 
@@ -119,10 +123,11 @@
 
 @policy_function(start_date="2005-01-01", vectorization_strategy="not_required")
 def in_anderer_bg_als_kindergeldempfänger(
-    p_id: np.ndarray,  # int
-    kindergeld__p_id_empfänger: np.ndarray,  # int
-    bg_id: np.ndarray,  # int
-) -> np.ndarray:  # bool
+    p_id: numpy.ndarray,  # int
+    kindergeld__p_id_empfänger: numpy.ndarray,  # int
+    bg_id: numpy.ndarray,  # int
+    xnp: ModuleType,  # Will become necessary for Jax.  # noqa: ARG001
+) -> numpy.ndarray:  # bool
     """True if the person is in a different Bedarfsgemeinschaft than the
     Kindergeldempfänger of that person.
     """

src/_gettsim/arbeitslosengeld_2/regelbedarf.py

@@ -14,6 +14,8 @@
 )
 
 if TYPE_CHECKING:
+    from types import ModuleType
+
     from _gettsim.grundsicherung.bedarfe import Regelbedarfsstufen
     from ttsim.tt_dag_elements import RawParam
 
@@ -450,11 +452,12 @@
 def berechtigte_wohnfläche_eigentum(
     parameter_berechtigte_wohnfläche_eigentum: RawParam,
     wohngeld__max_anzahl_personen: dict[str, int],
+    xnp: ModuleType,
 ) -> ConsecutiveInt1dLookupTableParamValue:
     """Berechtigte Wohnfläche für Eigenheim."""
     tmp = parameter_berechtigte_wohnfläche_eigentum.copy()
     je_weitere_person = tmp.pop("je_weitere_person")
     max_anzahl_direkt = tmp.pop("max_anzahl_direkt")
     for i in range(wohngeld__max_anzahl_personen["indizierung"] - max_anzahl_direkt):
         tmp[i] = tmp[max_anzahl_direkt] + i * je_weitere_person
-    return get_consecutive_int_1d_lookup_table_param_value(raw=tmp)
+    return get_consecutive_int_1d_lookup_table_param_value(raw=tmp, xnp=xnp)

src/_gettsim/einkommensteuer/abzüge/alter.py

@@ -11,6 +11,8 @@
 )
 
 if TYPE_CHECKING:
+    from types import ModuleType
+
     from ttsim.tt_dag_elements import ConsecutiveInt1dLookupTableParamValue
 
 
@@ -94,6 +96,7 @@
 @param_function(start_date="2005-01-01")
 def altersentlastungsquote_gestaffelt(
     raw_altersentlastungsquote_gestaffelt: dict[str | int, int | float],
+    xnp: ModuleType,
 ) -> ConsecutiveInt1dLookupTableParamValue:
     """Convert the raw parameters for the age-based tax deduction allowance to a dict."""
     spec = raw_altersentlastungsquote_gestaffelt.copy()
@@ -104,12 +107,14 @@
         raw=spec_int_float,
         left_tail_key=first_birthyear_to_consider,
         right_tail_key=last_birthyear_to_consider,
+        xnp=xnp,
     )
 
 
 @param_function(start_date="2005-01-01")
 def maximaler_altersentlastungsbetrag_gestaffelt(
     raw_maximaler_altersentlastungsbetrag_gestaffelt: dict[str | int, int | float],
+    xnp: ModuleType,
 ) -> ConsecutiveInt1dLookupTableParamValue:
     """Convert the raw parameters for the age-based tax deduction allowance to a dict."""
     spec = raw_maximaler_altersentlastungsbetrag_gestaffelt.copy()
@@ -120,13 +125,15 @@
         raw=spec_int_float,
         left_tail_key=first_birthyear_to_consider,
         right_tail_key=last_birthyear_to_consider,
+        xnp=xnp,
     )
 
 
 def get_consecutive_int_1d_lookup_table_with_filled_up_tails(
     raw: dict[int, float],
     left_tail_key: int,
     right_tail_key: int,
+    xnp: ModuleType,
 ) -> ConsecutiveInt1dLookupTableParamValue:
     """Create a consecutive integer lookup table with filled tails.
 
@@ -148,5 +155,6 @@
         range(max_key_in_spec + 1, right_tail_key + 1), raw[max_key_in_spec]
     )
     return get_consecutive_int_1d_lookup_table_param_value(
-        {**consecutive_dict_start, **raw, **consecutive_dict_end}
+        raw={**consecutive_dict_start, **raw, **consecutive_dict_end},
+        xnp=xnp,
     )

src/_gettsim/einkommensteuer/abzüge/behinderung.py

@@ -7,15 +7,20 @@
 from ttsim.tt_dag_elements import piecewise_polynomial, policy_function
 
 if TYPE_CHECKING:
+    from types import ModuleType
+
     from ttsim.tt_dag_elements import PiecewisePolynomialParam
 
 
 @policy_function()
 def pauschbetrag_behinderung_y(
-    behinderungsgrad: int, parameter_behindertenpauschbetrag: PiecewisePolynomialParam
+    behinderungsgrad: int,
+    parameter_behindertenpauschbetrag: PiecewisePolynomialParam,
+    xnp: ModuleType,
 ) -> float:
     """Assign tax deduction allowance for handicaped to different handicap degrees."""
     return piecewise_polynomial(
         x=behinderungsgrad,
         parameters=parameter_behindertenpauschbetrag,
+        xnp=xnp,
     )

src/_gettsim/einkommensteuer/abzüge/vorsorge.py

@@ -6,6 +6,8 @@
 from ttsim.tt_dag_elements.column_objects_param_function import param_function
 
 if TYPE_CHECKING:
+    from types import ModuleType
+
     from ttsim.tt_dag_elements.param_objects import PiecewisePolynomialParamValue
 
 
@@ -201,6 +203,7 @@
 def rate_abzugsfähige_altersvorsorgeaufwendungen(
     evaluationsjahr: int,
     parameter_einführungsfaktor_altersvorsorgeaufwendungen: PiecewisePolynomialParamValue,
+    xnp: ModuleType,
 ) -> dict[str, Any]:
     """Calculate introductory factor for pension expense deductions which depends on the
     current year as follows:
@@ -215,6 +218,7 @@
     return piecewise_polynomial(
         x=evaluationsjahr,
         parameters=parameter_einführungsfaktor_altersvorsorgeaufwendungen,
+        xnp=xnp,
     )