anchorboosting 0.2.0__tar.gz → 0.3.0__tar.gz

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/.github/workflows/ci.yml

@@ -8,6 +8,28 @@ name: CI
 
 jobs:
 
+  asv-bench:
+    name: ASV Benchmarks
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - name: Set up micromamba
+        uses: mamba-org/setup-micromamba@d05808540d968a55ca33c798e0661fb98f533c73
+        with:
+          micromamba-version: 1.5.10-0
+          environment-file: benchmarks/environment.yml
+    
+      - name: Install repository
+        run: python -m pip install --no-build-isolation --no-deps --disable-pip-version-check -e .
+      - name: Run benchmarks
+        shell: bash -l {0}
+        run: |
+          asv machine --machine github --yes
+          # asv run --machine github --python=same
+          asv continuous origin/main HEAD -e --factor 1.05 --machine github --python=same
+
   test:
     name: Unit tests - ${{ matrix.PYTHON_VERSION }}
     runs-on: ubuntu-latest

anchorboosting-0.3.0/CHANGELOG.rst

@@ -0,0 +1,24 @@
+Changelog
+=========
+
+0.3.0 (2025-09-15)
+------------------
+
+- Fixed handling of feature names for numpy input in `AnchorBooster`.
+- Added support for pandas DataFrame input in `AnchorBooster`.
+
+0.2.1 (2025-07-01)
+------------------
+
+- `AnchorBooster.refit` no longer raises an error if the model stopped boosting early during initial training.
+
+0.2.0 (2025-06-12)
+------------------
+
+- Sped up boosting.
+
+0.1.0 (2025-05-29)
+------------------
+
+- Initial release of the `anchorboosting` package.
+- Implemented `AnchorBooster` class for anchor regression and anchor probit classification.

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anchorboosting
-Version: 0.2.0
+Version: 0.3.0
 Summary: Efficient tree-based nonlinear anchor regression and classification
 Author-email: Malte Londschien <malte@londschien.ch>
 Project-URL: homepage, https://github.com/mlondschien/anchorboosting/
@@ -15,8 +15,15 @@ Dynamic: license-file
 
 # Anchor Boosting
 
-Efficient tree-based nonlinear anchor regression and classification for Python [1]
+Efficient tree-based nonlinear anchor regression and classification for Python.
 
-## References
+If you use this, please consider citing
 
-[1] Londschien, M., M. Burger, G. Rätsch, P. Bühlmann (2025). Working paper.
+```
+article{londschien2025,
+  title={Domain Generalization and Adaptation in Intensive Care with Anchor Regression},
+  author={Londschien, Malte, Burger, Manuel, R{\"a}tsch, Gunnar and B{\"u}hlmann, Peter},
+  journal={arXiv preprint arXiv:2507.21783},
+  year={2025}
+}
+```

anchorboosting-0.3.0/README.md

@@ -0,0 +1,14 @@
+# Anchor Boosting
+
+Efficient tree-based nonlinear anchor regression and classification for Python.
+
+If you use this, please consider citing
+
+```
+article{londschien2025,
+  title={Domain Generalization and Adaptation in Intensive Care with Anchor Regression},
+  author={Londschien, Malte, Burger, Manuel, R{\"a}tsch, Gunnar and B{\"u}hlmann, Peter},
+  journal={arXiv preprint arXiv:2507.21783},
+  year={2025}
+}
+```

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/anchorboosting/models.py

@@ -161,9 +161,6 @@ class AnchorBooster:
         if self.objective != "regression" and not np.isin(y, [0, 1]).all():
             raise ValueError("For binary classification, y values must be in {0, 1}.")
 
-        if Z is None:
-            Z = np.zeros(shape=(len(y), 0))
-
         y = y.flatten()
 
         if self.objective == "regression":
@@ -177,7 +174,9 @@ class AnchorBooster:
             )
 
         if hasattr(X, "columns"):
-            feature_name = X.columns
+            feature_name = list(X.columns)
+        else:
+            feature_name = None
 
         if _POLARS_INSTALLED and isinstance(X, pl.DataFrame):
             X = X.to_arrow()
@@ -198,7 +197,7 @@ class AnchorBooster:
         f = np.ones(len(y), dtype=np.float64) * self.init_score_
         booster_preds = f.copy()
 
-        Q, _ = np.linalg.qr(Z, mode="reduced")  # P_Z f = Q @ (Q^T @ f)
+        proj = Proj(Z)
 
         if self.objective == "binary":
             y_tilde = np.where(y == 1, 1, -1).astype(np.float64)
@@ -231,7 +230,7 @@ class AnchorBooster:
                 dr = -f * r + r**2  # d^2/df^2 loss(f, y)
                 ddr = (f**2 - 1) * r - 3 * f * r**2 + 2 * r**3  # d^3/df^3 loss
 
-            r_proj = Q @ (Q.T @ r)
+            r_proj = proj(r)
             grad = r + (self.gamma - 1) * r_proj * dr
 
             # We wish to fit one additional tree. Intuitively, one would use
@@ -348,19 +347,8 @@ class AnchorBooster:
             counts += self.params.get("lambda_l2", 0)
             counts[counts == 0] = 1  # counts==0 should never happen.
 
-            # Mdr^T P_Z Mdr = (Mdr^T Q) @ (Mdr^T Q)^T
-            # One could also compute this using bincount, but it appears this
-            # version using a sparse matrix is faster.
-            Mdr = scipy.sparse.csr_matrix(
-                (
-                    dr,
-                    (np.arange(len(leaves)), leaves),
-                ),
-                shape=(len(leaves), num_leaves),
-                dtype=np.float64,
-            )
-            B = Mdr.T.dot(Q)
-            H = (self.gamma - 1) * B @ B.T
+            # (dr * M)^T P_Z (dr * M)
+            H = (self.gamma - 1) * proj.sandwich(leaves, num_leaves, weights=dr)
             H += np.diag(counts)
 
             # Compute the 2nd order update
@@ -449,7 +437,8 @@ class AnchorBooster:
 
         f = np.full(len(y), self.init_score_, dtype="float64")
 
-        for idx in range(self.num_boost_round):
+        # The model might have stopped boosting prematurely.
+        for idx in range(leaves.shape[1]):
             n_obs = np.bincount(leaves[:, idx], minlength=num_leaves[idx])
 
             if self.objective == "regression":
@@ -505,3 +494,101 @@ def _suppress_stderr():
         os.dup2(old_stderr_fd, stderr_fd)  # put the old stderr back
         os.close(old_stderr_fd)
         os.close(devnull_fd)
+
+
+class Proj:
+    """
+    Cache the projection onto the subspace spanned by Z.
+
+    Parameters
+    ----------
+    Z: np.ndarray of dimension (n, d_Z) or (n,), optional, default=None
+        The `Z` matrix or 1d array of integers.
+    """
+
+    def __init__(self, Z):
+        if Z is None:
+            self._type = "none"
+
+        elif len(Z.shape) == 1 and np.issubdtype(Z.dtype, np.integer):
+            self._type = "categorical"
+            self._Z = Z
+            self._n_categories = np.max(Z) + 1
+            counts = np.bincount(Z, minlength=self._n_categories)
+            counts[counts == 0] = 1.0  # Avoid division by zero
+            self._one_over_counts = 1 / counts
+
+        elif len(Z.shape) == 2 and np.issubdtype(Z.dtype, np.floating):
+            self._type = "linear"
+            self._Q, _ = np.linalg.qr(Z, mode="reduced")
+        else:
+            raise ValueError(
+                "Z should be either a 1d array of integers or a 2d array of floats. "
+                f"Got shape {Z.shape} and dtype {Z.dtype}."
+            )
+
+    def __call__(self, f):
+        """
+        Project the input array `f` onto the subspace spanned by `Z`.
+
+        Parameters
+        ----------
+        f: np.ndarray of shape (n,)
+            The input array to project.
+
+        Returns
+        -------
+        np.ndarray of shape (n,)
+            The projected array.
+        """
+        if self._type == "none":
+            return np.zeros_like(f)
+
+        elif self._type == "categorical":
+            sums = np.bincount(self._Z, weights=f, minlength=self._n_categories)
+            means = sums * self._one_over_counts
+            return means[self._Z]
+
+        elif self._type == "linear":
+            return self._Q @ (self._Q.T @ f)
+
+    def sandwich(self, leaves, num_leaves, weights):
+        """
+        For M = weights * one_hot(leaves), return proj(Z, M).T @ proj(Z, M).
+
+        Parameters
+        ----------
+        leaves: np.ndarray of shape (n,)
+            The leaf indices for each sample in `f`. Integers in [0, num_leaves).
+        num_leaves: int
+            The number of leaves in the decision tree.
+        weights: np.ndarray of shape (n,)
+            The input array to project.
+
+        Returns
+        -------
+        np.ndarray of shape (d, d)
+            The sandwich product.
+        """
+        if self._type == "none":
+            return np.zeros((num_leaves, num_leaves), dtype=weights.dtype)
+
+        elif self._type == "categorical":
+            S = np.zeros((self._n_categories, num_leaves), dtype=weights.dtype)
+            np.add.at(S, (self._Z, leaves), weights)
+            return (S * self._one_over_counts[:, np.newaxis]).T @ S
+
+        elif self._type == "linear":
+            # M^T P_Z M = (M^T Q) @ (M^T Q)^T
+            # One could also compute this using bincount, but it appears this
+            # version using a sparse matrix is faster.
+            M = scipy.sparse.csr_matrix(
+                (
+                    weights,
+                    (np.arange(len(leaves)), leaves),
+                ),
+                shape=(len(leaves), num_leaves),
+                dtype=weights.dtype,
+            )
+            B = M.T @ self._Q
+            return B @ B.T

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/anchorboosting/simulate.py

@@ -2,7 +2,7 @@ import numpy as np
 import polars as pl
 
 
-def simulate(f, n=100, shift=0, seed=0):
+def simulate(f, n=100, shift=0, seed=0, return_dtype="polars"):
     rng = np.random.RandomState(seed)
 
     p = 3
@@ -23,6 +23,15 @@ def simulate(f, n=100, shift=0, seed=0):
         schema={"x1": pl.Float64, "x2": pl.Float64, "x3": pl.Int64},
     )
 
+    if return_dtype == "numpy":
+        df = df.to_numpy()
+        a = a.astype(np.float32)
+        y = y.astype(np.float32)
+    elif return_dtype == "pandas":
+        df = df.to_pandas()
+        a = a.astype(np.float32)
+        y = y.astype(np.float32)
+
     return df, y, a
 
 

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/anchorboosting.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anchorboosting
-Version: 0.2.0
+Version: 0.3.0
 Summary: Efficient tree-based nonlinear anchor regression and classification
 Author-email: Malte Londschien <malte@londschien.ch>
 Project-URL: homepage, https://github.com/mlondschien/anchorboosting/
@@ -15,8 +15,15 @@ Dynamic: license-file
 
 # Anchor Boosting
 
-Efficient tree-based nonlinear anchor regression and classification for Python [1]
+Efficient tree-based nonlinear anchor regression and classification for Python.
 
-## References
+If you use this, please consider citing
 
-[1] Londschien, M., M. Burger, G. Rätsch, P. Bühlmann (2025). Working paper.
+```
+article{londschien2025,
+  title={Domain Generalization and Adaptation in Intensive Care with Anchor Regression},
+  author={Londschien, Malte, Burger, Manuel, R{\"a}tsch, Gunnar and B{\"u}hlmann, Peter},
+  journal={arXiv preprint arXiv:2507.21783},
+  year={2025}
+}
+```

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/anchorboosting.egg-info/SOURCES.txt

@@ -3,6 +3,7 @@
 CHANGELOG.rst
 LICENSE
 README.md
+asv.conf.json
 environment.yml
 pyproject.toml
 readthedocs.yaml
@@ -11,12 +12,16 @@ readthedocs.yaml
 anchorboosting/__init__.py
 anchorboosting/models.py
 anchorboosting/simulate.py
-anchorboosting/utils.py
 anchorboosting.egg-info/PKG-INFO
 anchorboosting.egg-info/SOURCES.txt
 anchorboosting.egg-info/dependency_links.txt
 anchorboosting.egg-info/requires.txt
 anchorboosting.egg-info/top_level.txt
+benchmarks/__init__.py
+benchmarks/benchmark.py
+benchmarks/environment.yml
+benchmarks/load.py
+benchmarks/data/.gitignore
 docs/Makefile
 docs/bib.bib
 docs/changelog.rst
@@ -24,5 +29,5 @@ docs/conf.py
 docs/index.rst
 tests/test_anchor_boosting.py
 tests/test_derivatives.py
-tests/test_refit.py
-tests/test_utils.py
+tests/test_proj.py
+tests/test_refit.py

anchorboosting-0.3.0/asv.conf.json

@@ -0,0 +1,18 @@
+{
+    "version": 1,
+    "branches": ["main"],
+    "project": "anchorboosting",
+    "project_url": "https://github.com/mlondschien/anchorboosting",
+    "repo": "..",
+    "environment_type": "conda",
+    "show_commit_url": "https://github.com/mlondschien/anchorboosting/commit/",
+    "pythons": ["3.12"],
+    "conda_channels": ["conda-forge"],
+    "build_command": "",
+    "install_command": "python -mpip install --no-deps .",
+    "conda_environment_file": "benchmarks/environment.yml",
+    "benchmark_dir": "benchmarks",
+    "env_dir": ".asv/env",
+    "results_dir": ".asv/results",
+    "html_dir": ".asv/html"
+}

anchorboosting-0.3.0/benchmarks/benchmark.py

@@ -0,0 +1,77 @@
+import timeit
+from abc import ABC, abstractmethod
+from typing import Any
+
+from lightgbm import LGBMModel
+
+from anchorboosting import AnchorBooster
+
+from .load import load_nyc_taxi
+
+
+class Benchmark(ABC):
+    # Measure wall time instead of CPU usage
+    timer = timeit.default_timer
+
+    name = ""
+
+    params: Any
+    param_names: Any
+
+    def setup(self, *params):
+        self._setup_data(params)
+        self._setup_model(params)
+
+    def _setup_data(self, params):
+        idx = self.param_names.index("n")
+        n = params[idx]
+        self.X, self.Z, self.y, self.y_binary = load_nyc_taxi(n)
+
+    @abstractmethod  # Required for asv to skip.
+    def _setup_model(self, params):
+        pass
+
+
+class AnchorBoosterBenchmark(Benchmark):
+    def __init__(self, n_jobs=2):
+        self.n_jobs = n_jobs  # GitHub CI has two nodes.
+
+    name = "AnchorBooster"
+    param_names = ["objective", "n", "num_boost_round", "max_depth"]
+    params = (["regression", "binary"], [100_000], [100], [3])
+
+    def _setup_model(self, params):
+        objective, _, num_boost_round, max_depth = params
+        self.model = AnchorBooster(
+            gamma=2,
+            num_boost_round=num_boost_round,
+            max_depth=max_depth,
+            n_jobs=self.n_jobs,
+            objective=objective,
+        )
+
+    def time_fit(self, *args):
+        y = self.y_binary if self.model.objective == "binary" else self.y
+        self.model.fit(self.X, y, self.Z)
+
+
+class LGBMBenchmark(Benchmark):
+    def __init__(self, n_jobs=2):
+        self.n_jobs = n_jobs  # GitHub CI has two nodes.
+
+    name = "LGBM Booster"
+    param_names = ["objective", "n", "num_boost_round", "max_depth"]
+    params = (["regression", "binary"], [100_000], [100], [3])
+
+    def _setup_model(self, params):
+        objective, _, num_boost_round, max_depth = params
+        self.model = LGBMModel(
+            num_boost_round=num_boost_round,
+            max_depth=max_depth,
+            objective=objective,
+            n_jobs=self.n_jobs,
+        )
+
+    def time_fit(self, *args):
+        y = self.y_binary if self.model.objective == "binary" else self.y
+        self.model.fit(self.X, y)

anchorboosting-0.3.0/benchmarks/data/.gitignore

@@ -0,0 +1 @@
+*.parquet

anchorboosting-0.3.0/benchmarks/environment.yml

@@ -0,0 +1,14 @@
+name: benchmarks
+channels:
+  - conda-forge
+  - nodefaults
+dependencies:
+  - python =3.12
+  - numpy
+  - lightgbm
+  - asv
+  - scikit-learn <=1.5.2  # https://github.com/microsoft/LightGBM/issues/6849
+  - polars
+  - pyarrow
+  - cffi  # https://github.com/microsoft/LightGBM/pull/6785
+  - libmambapy

anchorboosting-0.3.0/benchmarks/load.py

@@ -0,0 +1,110 @@
+# From https://github.com/xhochy/nyc-taxi-fare-prediction-deployment-example
+from pathlib import Path
+
+import polars as pl
+from sklearn.compose import ColumnTransformer
+from sklearn.preprocessing import OrdinalEncoder
+
+DATA_PATH = Path(__file__).resolve().parent / "data"
+NYC_TAXI_DATASET_PATH = DATA_PATH / "nyc_taxi_data.parquet"
+NYC_TAXI_DATASET_URL = (
+    "https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2019-01.parquet"
+)
+NYC_TAXI_BURROWS_PATH = DATA_PATH / "nyc_taxi_burrows.parquet"
+NYC_TAXI_BURROWS_URL = "https://d37ci6vzurychx.cloudfront.net/misc/taxi_zone_lookup.csv"
+
+
+def load_nyc_taxi_raw():
+    if NYC_TAXI_DATASET_PATH.exists():
+        return pl.read_parquet(NYC_TAXI_DATASET_PATH)
+    else:
+        df = pl.read_parquet(NYC_TAXI_DATASET_URL)
+        df.write_parquet(NYC_TAXI_DATASET_PATH)
+        return df
+
+
+def load_nyc_taxi_burrows():
+    if NYC_TAXI_BURROWS_PATH.exists():
+        return pl.read_parquet(NYC_TAXI_BURROWS_PATH)
+    else:
+        df = pl.read_csv(NYC_TAXI_BURROWS_URL)
+        df.write_parquet(NYC_TAXI_BURROWS_PATH)
+        return df
+
+
+def load_nyc_taxi(n):
+    df = load_nyc_taxi_raw().head(n)
+
+    burrows = load_nyc_taxi_burrows()
+
+    df = (
+        df.filter(pl.col("total_amount").is_not_null() & pl.col("total_amount").gt(0))
+        .join(
+            burrows.with_columns(pl.col("Borough").alias("pickup_borough")),
+            left_on="PULocationID",
+            right_on="LocationID",
+            how="left",
+        )
+        .join(
+            burrows.with_columns(pl.col("Borough").alias("dropoff_borough")),
+            left_on="DOLocationID",
+            right_on="LocationID",
+            how="left",
+        )
+        .with_columns(
+            pl.col("total_amount").log().alias("total_amount_log"),
+            pl.col("tpep_pickup_datetime").dt.weekday().alias("pickup_weekday"),
+            pl.col("tpep_pickup_datetime").dt.hour().alias("pickup_hour"),
+            pl.col("tpep_pickup_datetime").dt.minute().alias("pickup_minute"),
+            pl.col("pickup_borough").fill_null("Unknown").alias("pickup_borough"),
+        )
+    )
+
+    continuous_columns = [
+        "passenger_count",
+        "trip_distance",
+        "extra",
+        "mta_tax",
+        "tip_amount",
+        "tolls_amount",
+        "improvement_surcharge",
+        "congestion_surcharge",
+        "pickup_hour",
+        "pickup_minute",
+    ]
+    df = df.with_columns(pl.col(c).cast(pl.Float64) for c in continuous_columns)
+
+    ordinal_columns = [
+        "VendorID",
+        "RatecodeID",
+        "store_and_fwd_flag",
+        "payment_type",
+        "service_zone",
+        "Borough",
+        "Zone",
+        "pickup_weekday",
+        "pickup_borough",
+        "dropoff_borough",
+    ]
+
+    preprocessor = ColumnTransformer(
+        transformers=[
+            ("continuous", "passthrough", continuous_columns),
+            (
+                "categorical",
+                OrdinalEncoder(
+                    handle_unknown="use_encoded_value",
+                    unknown_value=-1,
+                ),
+                ordinal_columns,
+            ),
+        ]
+    ).set_output(transform="polars")
+
+    y = df.select(pl.col("total_amount_log")).to_numpy().ravel()
+    y_binary = (y > y.mean()).astype(int)
+
+    X = preprocessor.fit_transform(df)
+    Z = X.select("categorical__Zone").to_numpy().astype("int32").ravel()
+
+    return X.drop("categorical__Zone"), Z, y, y_binary

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/docs/index.rst

@@ -30,7 +30,7 @@ API Reference
    :maxdepth: 1
    :caption: Other
 
-   GitHub <https://github.com/mlondschien/ivmodels>
+   GitHub <https://github.com/mlondschien/anchorboosting>
    changelog
 
 

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/environment.yml

@@ -9,9 +9,9 @@ dependencies:
   - lightgbm >=4.0.0
   - setuptools  # build
   - pip  # build
-  - scikit-learn  # tests
   - pytest-cov  # tests
   - pytest  # tests
+  - pandas  # tests
   - polars  # tests
   - pyarrow  # tests
   - sphinx  # docs

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/pyproject.toml

@@ -5,7 +5,7 @@ requires = ['setuptools', 'setuptools-scm', 'wheel']
 name = "anchorboosting"
 description = "Efficient tree-based nonlinear anchor regression and classification"
 readme = "README.md"
-version = "0.2.0"
+version = "0.3.0"
 requires-python = ">=3.9"
 authors = [
     { name = "Malte Londschien", email = "malte@londschien.ch" }

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/tests/test_anchor_boosting.py

@@ -3,19 +3,24 @@ import numpy as np
 import pytest
 import scipy
 
-from anchorboosting import AnchorBooster
+from anchorboosting.models import AnchorBooster, Proj
 from anchorboosting.simulate import f1, simulate
 
 
 @pytest.mark.parametrize("gamma", [1.0, 2.0, 100])
 @pytest.mark.parametrize("objective", ["regression", "binary"])
-def test_anchor_boosting_second_order(gamma, objective):
+@pytest.mark.parametrize("categorical_z", [True, False])
+@pytest.mark.parametrize("input_dtype", ["polars", "numpy"])
+def test_anchor_boosting_second_order(gamma, objective, categorical_z, input_dtype):
     learning_rate = 0.1
     num_leaves = 5
     n = 200
     num_boost_round = 10
 
-    x, y, a = simulate(f1, n=n, shift=0, seed=0)
+    x, y, a = simulate(f1, n=n, shift=0, seed=0, return_dtype=input_dtype)
+    if categorical_z:
+        a = np.digitize(a[:, 0], bins=[-1, 0, 1, 2, 3])
+        assert np.issubdtype(a.dtype, np.integer)
 
     if objective == "binary":
         y = (y > 0).astype(int)
@@ -32,15 +37,17 @@ def test_anchor_boosting_second_order(gamma, objective):
     f = model.predict(x, num_iteration=num_boost_round - 1, raw_score=True)
 
     leaves = model.booster_.predict(
-        x.to_arrow(),
+        x.to_arrow() if hasattr(x, "to_arrow") else x,
         pred_leaf=True,
         start_iteration=num_boost_round - 1,
         num_iteration=1,
     ).flatten()
 
+    proj = Proj(a)
+
     def regression_loss(leaf_values):
         residuals = y - f - leaf_values[leaves]
-        Pa_residuals = a @ np.linalg.solve(a.T @ a, a.T @ residuals)
+        Pa_residuals = proj(residuals)
         return np.sum(np.square(residuals)) + (gamma - 1) * residuals.T @ Pa_residuals
 
     def probit_loss(leaf_values):
@@ -49,7 +56,7 @@ def test_anchor_boosting_second_order(gamma, objective):
         dp = scipy.stats.norm.pdf(scores)
         losses = -np.log(np.where(y == 1, p, 1 - p))
         dl = np.where(y == 1, -dp / p, dp / (1 - p))
-        Pa_dl = a @ np.linalg.solve(a.T @ a, a.T @ dl)
+        Pa_dl = proj(dl)
         return np.sum(losses) + (gamma - 1) / 2 * dl.T @ Pa_dl
 
     if objective == "regression":
@@ -138,8 +145,19 @@ def test_anchor_boosting_decreases_loss(gamma, objective):
         {"lambda_l2": 0.1},
     ],
 )
-def test_compare_anchor_boosting_to_lgbm(parameters):
-    X, y, a = simulate(f1, shift=0, seed=0)
+@pytest.mark.parametrize("input_dtype", ["polars", "numpy", "pandas"])
+def test_compare_anchor_boosting_to_lgbm(parameters, input_dtype):
+    X, y, a = simulate(f1, shift=0, seed=0, return_dtype=input_dtype)
+
+    if input_dtype == "polars":
+        X_arrow = X.to_arrow()
+        categorical_feature = ["x3"]
+    elif input_dtype == "pandas":
+        X_arrow = X
+        categorical_feature = ["x3"]
+    elif input_dtype == "numpy":
+        X_arrow = X
+        categorical_feature = [2]
 
     lgbm_model = lgb.train(
         params={
@@ -147,7 +165,7 @@ def test_compare_anchor_boosting_to_lgbm(parameters):
             "objective": "regression",
             **parameters,
         },
-        train_set=lgb.Dataset(X.to_arrow(), y, categorical_feature=[2]),
+        train_set=lgb.Dataset(X_arrow, y, categorical_feature=[2]),
         num_boost_round=50,
     )
 
@@ -157,7 +175,7 @@ def test_compare_anchor_boosting_to_lgbm(parameters):
         objective="regression",
         learning_rate=0.1,
         **parameters,
-    ).fit(X, y, Z=a, categorical_feature=["x3"])
+    ).fit(X, y, Z=a, categorical_feature=categorical_feature)
 
     anchor_booster_noncat = AnchorBooster(
         gamma=1,
@@ -167,7 +185,7 @@ def test_compare_anchor_boosting_to_lgbm(parameters):
         **parameters,
     ).fit(X, y, Z=a)
 
-    lgbm_pred = lgbm_model.predict(X.to_arrow())
+    lgbm_pred = lgbm_model.predict(X_arrow)
     anchor_booster_pred = anchor_booster.predict(X)
     anchor_booster_noncat_pred = anchor_booster_noncat.predict(X)
 
@@ -175,19 +193,65 @@ def test_compare_anchor_boosting_to_lgbm(parameters):
     assert not np.allclose(lgbm_pred, anchor_booster_noncat_pred)
 
 
+@pytest.mark.parametrize("gamma", [1, 10])
+def test_compare_input_types(gamma):
+    X_polars, y_polars, a_polars = simulate(f1, shift=0, seed=0, return_dtype="polars")
+    X_numpy, y_numpy, a_numpy = simulate(f1, shift=0, seed=0, return_dtype="numpy")
+    X_pandas, y_pandas, a_pandas = simulate(f1, shift=0, seed=0, return_dtype="pandas")
+
+    model_polars = AnchorBooster(
+        gamma=gamma,
+        num_boost_round=10,
+        objective="regression",
+    ).fit(X_polars, y_polars, Z=a_polars)
+
+    model_numpy = AnchorBooster(
+        gamma=gamma,
+        num_boost_round=10,
+        objective="regression",
+    ).fit(X_numpy, y_numpy, Z=a_numpy)
+
+    model_pandas = AnchorBooster(
+        gamma=gamma,
+        num_boost_round=10,
+        objective="regression",
+    ).fit(X_pandas, y_pandas, Z=a_pandas)
+
+    pred_pandas = model_pandas.predict(X_pandas)
+    pred_polars = model_polars.predict(X_polars)
+    pred_numpy = model_numpy.predict(X_numpy)
+
+    np.testing.assert_allclose(pred_polars, pred_numpy, rtol=1e-5)
+    np.testing.assert_allclose(pred_polars, pred_pandas, rtol=1e-5)
+
+
 @pytest.mark.parametrize("objective", ["regression", "binary"])
-def test_anchor_booster_init_score(objective):
+def test_categorical_and_onehot_is_the_same(objective):
     X, y, a = simulate(f1, shift=0, seed=0)
 
+    a = np.digitize(a[:, 0], bins=[-1, 0, 1, 2, 3])
+
     if objective == "binary":
         y = (y > 0).astype(int)
 
-    anchor_booster = AnchorBooster(
+    anchor_booster1 = AnchorBooster(
         gamma=1,
-        num_boost_round=0,
+        num_boost_round=10,
         objective=objective,
         learning_rate=0.1,
     ).fit(X, y, Z=a)
 
-    predictions = anchor_booster.predict(X)
-    assert np.allclose(predictions, np.ones_like(y) * y.mean())
+    predictions1 = anchor_booster1.predict(X)
+
+    M = scipy.sparse.csr_matrix(
+        (np.ones(len(a), dtype=np.float64), (np.arange(len(a)), a)),
+        shape=(len(a), len(np.unique(a))),
+    )
+    anchor_booster2 = AnchorBooster(
+        gamma=1,
+        num_boost_round=10,
+        objective=objective,
+        learning_rate=0.1,
+    ).fit(X, y, Z=M.toarray())
+    predictions2 = anchor_booster2.predict(X)
+    np.testing.assert_allclose(predictions1, predictions2, rtol=1e-5)

anchorboosting-0.3.0/tests/test_proj.py

@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+import scipy.sparse
+
+from anchorboosting.models import Proj
+
+cases = [
+    (
+        np.array([1, 1, 3]),
+        np.array([1.0, 2.0, 3.0]),
+        np.array([1.5, 1.5, 3.0]),
+    ),
+    (
+        np.array([[0.0], [0.0], [1.0]]),
+        np.array([1.0, 2.0, 3.0]),
+        np.array([0.0, 0.0, 3.0]),
+    ),
+    (
+        np.array([[0.0, 1.0], [0.0, 1.0], [1.0, 0.0]]),
+        np.array([1.0, 2.0, 3.0]),
+        np.array([1.5, 1.5, 3]),
+    ),
+    (
+        np.array([2, 2, 1, 7]),
+        np.array([1.0, 2.0, 3.0, 1.0]),
+        np.array([1.5, 1.5, 3.0, 1.0]),
+    ),
+    (
+        np.array([1, 1, 0]),
+        np.array([1.0, 0.0, 3.0]),
+        np.array([0.5, 0.5, 3]),
+    ),
+    (np.array([0]), np.array([1.0]), np.array([1.0])),
+    (np.array([[0.0]]), np.array([0.0]), np.array([0.0])),
+]
+
+
+@pytest.mark.parametrize("Z, f, result", cases)
+def test_cached_proj_result(Z, f, result):
+    np.testing.assert_almost_equal(Proj(Z)(f), result)
+
+
+@pytest.mark.parametrize("Z, f, _", cases)
+def test_cached_proj_dot_product(Z, f, _):
+    np.testing.assert_almost_equal(
+        np.dot(Proj(Z)(f).T, f),
+        np.dot(Proj(Z)(f).T, Proj(Z)(f)),
+    )
+
+
+def test_proj_sandwich():
+    n = 100
+    num_categories = 10
+    num_leaves = 5
+    rng = np.random.RandomState(0)
+
+    Z = rng.choice(np.arange(num_categories), size=n, replace=True)
+    leaves = rng.choice(np.arange(num_leaves), size=n, replace=True)
+    f = rng.normal(size=n)
+
+    M = scipy.sparse.csr_matrix((f, (np.arange(n), leaves)), (n, num_leaves)).toarray()
+    proj = Proj(Z)
+
+    sandwich1 = proj.sandwich(leaves, num_leaves, f)
+    for ldx in range(num_leaves):
+        M[:, ldx] = proj(M[:, ldx])
+    np.testing.assert_almost_equal(sandwich1, M.T @ M)

{anchorboosting-0.2.0 → anchorboosting-0.3.0}/tests/test_refit.py

@@ -9,9 +9,10 @@ from anchorboosting.simulate import f1, f2, simulate
 @pytest.mark.parametrize("seed", [0, 1])
 @pytest.mark.parametrize("fun", [f1, f2])
 @pytest.mark.parametrize("objective", ["regression", "binary"])
-def test_refit(objective, seed, fun):
+@pytest.mark.parametrize("input_dtype", ["numpy", "pandas", "polars"])
+def test_refit(objective, seed, fun, input_dtype):
 
-    X, y, a = simulate(fun, shift=0, seed=seed)
+    X, y, a = simulate(fun, shift=0, seed=seed, return_dtype=input_dtype)
 
     if objective == "binary":
         y = (y > 0).astype(int)

anchorboosting-0.2.0/CHANGELOG.rst

@@ -1,13 +0,0 @@
-Changelog
-=========
-
-0.2.0 (2025-06-12)
-------------------
-
-- Sped up boosting.
-
-0.1.0 (2025-05-29)
-------------------
-
-- Initial release of the `anchorboosting` package.
-- Implemented `AnchorBooster` class for anchor regression and anchor probit classification.

anchorboosting-0.2.0/README.md

@@ -1,7 +0,0 @@
-# Anchor Boosting
-
-Efficient tree-based nonlinear anchor regression and classification for Python [1]
-
-## References
-
-[1] Londschien, M., M. Burger, G. Rätsch, P. Bühlmann (2025). Working paper.

anchorboosting-0.2.0/anchorboosting/utils.py

@@ -1,84 +0,0 @@
-# copy-pasted from https://github.com/mlondschien/ivmodels/blob/09f8bc5a78b3793ae5\
-# 1f0bb7f0331b70a0971ebc/ivmodels/utils.py
-import numpy as np
-import scipy
-
-
-def proj(Z, f, n_categories=None, copy=False):
-    """
-    Project f onto the subspace spanned by Z.
-
-    If `Z` is a matrix of floats, uses a copy of `ivmodel`'s `proj:
-    https://github.com/mlondschien/ivmodels/blob/main/ivmodels/utils.py#L12
-    If `Z` is a 1d array of unsigned integers, assumes that `Z` is categorical with
-    values 0, 1, ..., n_categories-1. In this case, the `n_categories` argument must be
-    be supplied. The projection is done by averaging the values of f within each
-    category. If `Z` is None, returns np.zeros_like(f).
-
-    Parameters
-    ----------
-    Z: np.ndarray of dimension (n, d_Z) or (n,)
-        The `Z` matrix or 1d array of integers. If None, returns np.zeros_like(f).
-    *args: np.ndarrays of dimension (n, d_f) or (n,)
-        Vectors or matrices to project.
-    n_categories: int, optional, default=None
-        If not None, then Z is assumed to be categorical with categories 0, 1, ...,
-        n_categories-1. The projection is done by averaging the values of f within each
-        category.
-    copy: bool, default=False
-        If False, overwrites *args.
-
-    Returns
-    -------
-    np.ndarray of dimension (n, d_f) or (n,)
-        Projection of args onto the subspace spanned by Z. Same number of
-        outputs as args. Same dimension as args
-    """
-    if Z is None:
-        return np.zeros_like(f)
-
-    if len(f.shape) > 2:
-        raise ValueError(f"*args should have shapes (n, d_f) or (n,). Got {f.shape}.")
-    if f.shape[0] != Z.shape[0]:
-        raise ValueError(f"Shape mismatch: Z.shape={Z.shape}, f.shape={f.shape}.")
-
-    if n_categories is not None:
-        f = f.copy() if copy else f
-
-        if len(Z.shape) != 1 or not np.issubdtype(Z.dtype, np.integer):
-            raise ValueError(
-                "If n_categories is not None, then Z should be a single column of "
-                f"integers. Got shape {Z.shape} and dtype {Z.dtype}."
-            )
-
-        if len(f.shape) == 1:
-            means = np.zeros((n_categories,))
-            counts = np.zeros(n_categories)
-
-            # np.add.at(a, indices, b) is equivalent to a[indices] += b
-            np.add.at(means, Z, f)
-            np.add.at(counts, Z, 1)
-
-            means[counts > 0] = means[counts > 0] / counts[counts > 0]
-
-            f[:] = means[Z]
-        else:
-            means = np.zeros((n_categories, f.shape[1]))
-            counts = np.zeros(n_categories)
-
-            # np.add.at(a, indices, b) is equivalent to a[indices] += b
-            np.add.at(means, Z, f)
-            np.add.at(counts, Z, 1)
-
-            means[counts > 0, :] = means[counts > 0, :] / counts[counts > 0, None]
-
-            f[:] = means[Z, :]
-
-        return f
-
-    # The gelsy driver raises in this case - we handle it separately
-    if len(f.shape) == 2 and f.shape[1] == 0:
-        return np.zeros_like(f)
-
-    # return np.dot(Z, scipy.linalg.pinv(Z.T @ Z) @ Z.T @ args[0])
-    return np.dot(Z, scipy.linalg.lstsq(Z, f, cond=None, lapack_driver="gelsy")[0])

anchorboosting-0.2.0/tests/test_utils.py

@@ -1,74 +0,0 @@
-import numpy as np
-import pytest
-from scipy.linalg import norm
-from sklearn.preprocessing import OneHotEncoder
-
-from anchorboosting.utils import proj
-
-cases = [
-    (
-        np.array([[0], [0], [1]]),
-        np.array([[1.0], [2.0], [3.0]]),
-        np.array([[1.5], [1.5], [3]]),
-    ),
-    (
-        np.array([[2], [2], [1], [3]]),
-        np.array([[1.0], [2.0], [3.0], [1.0]]),
-        np.array([[1.5], [1.5], [3], [1]]),
-    ),
-    (
-        np.array([[1], [1], [0]]),
-        np.array([[1.0, 2.0], [0.0, 2.0], [3.0, 4.0]]),
-        np.array([[0.5, 2.0], [0.5, 2.0], [3.0, 4.0]]),
-    ),
-    (
-        np.array([[1], [1], [0]]),
-        np.array([1.0, 0.0, 3.0]),
-        np.array([0.5, 0.5, 3]),
-    ),
-    (
-        np.array([[1], [1], [0]]),
-        np.array([[1.0, 2.0], [0.0, 2.0], [3.0, 4.0]]),
-        np.array([[0.5, 2.0], [0.5, 2.0], [3.0, 4.0]]),
-    ),
-    (
-        np.array([[2], [2], [1]]),
-        np.array([[1.0, 2.0], [0.0, 2.0], [3.0, 4.0]]),
-        np.array([[0.5, 2.0], [0.5, 2.0], [3.0, 4.0]]),
-    ),
-    (np.array([[0]]), np.array([[1]]), np.array([[1]])),
-]
-
-
-@pytest.mark.parametrize("anchor, residuals, result", cases)
-@pytest.mark.parametrize("n_categories", [4, None])
-def test_proj_result(n_categories, anchor, residuals, result):
-    if n_categories is None:
-        if anchor.shape[1] != 1:
-            pytest.skip("one_hot only makes sense for 1D anchors")
-        encoder = OneHotEncoder(sparse_output=False, drop=None)
-        anchor = encoder.fit_transform(anchor).astype(float)
-    else:
-        anchor = anchor.flatten()
-
-    np.testing.assert_almost_equal(
-        proj(anchor, residuals.copy(), n_categories=n_categories), result
-    )
-
-
-@pytest.mark.parametrize("anchor, residuals, _", cases)
-def test_proj_dot_product(anchor, residuals, _):
-    np.testing.assert_almost_equal(
-        np.dot(proj(anchor, residuals).T, residuals),
-        np.dot(proj(anchor, residuals).T, proj(anchor, residuals)),
-    )
-
-
-@pytest.mark.parametrize("anchor, residuals, result", cases)
-@pytest.mark.parametrize("gamma", [0.1, 1, 2, 10, 100, 1000])
-def test_proj_orthogonal(anchor, residuals, result, gamma):
-    np.testing.assert_almost_equal(
-        norm(residuals - proj(anchor, residuals)) ** 2
-        + gamma * norm(proj(anchor, residuals)) ** 2,
-        norm(residuals) ** 2 + (gamma - 1) * norm(proj(anchor, residuals)) ** 2,
-    )