fix tests (#73)

* update dependencies and fix errors

.github/workflows/auto-publish.yml

@@ -6,20 +6,30 @@ jobs:
   release:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v2
-      with:
-        python-version: 3.8
-    - name: Install Tools
-      run: |
-        python -m pip install --upgrade pip
-        pip install setuptools wheel twine
-    - name: Package and Upload
-      env:
-        STACKMANAGER_VERSION: ${{ github.event.release.tag_name }}
-        TWINE_USERNAME: __token__
-        TWINE_PASSWORD: ${{ secrets.PYPI_APIKEY }}
-      run: |
-        python setup.py sdist bdist_wheel
-        twine upload dist/*
+      - uses: actions/checkout@v2
+
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
+      - name: Install Poetry
+        run: |
+          curl -sSL https://install.python-poetry.org | python3 -
+          echo "$HOME/.local/bin" >> $GITHUB_PATH
+
+      - name: Configure Poetry
+        run: |
+          poetry config pypi-token.pypi ${{ secrets.PYPI_APIKEY }}
+
+      - name: Install dependencies
+        run: |
+          poetry install
+
+      - name: Build package
+        run: |
+          poetry build
+
+      - name: Publish package
+        run: |
+          poetry publish

.github/workflows/testing.yml

@@ -3,25 +3,29 @@ name: Code Checks
 on:
   push:
     branches:
-    - main
+      - main
   pull_request:
     branches:
-    - main
+      - main
 
 jobs:
   build:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7]
-
+        python-version: ["3.9", "3.10", "3.11"]
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install Testing Dependencies
-      run: make install-dev
-    - name: Automated checks
-      run: make check
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+        with:
+          version: 1.5.1  # You can specify the Poetry version you want to use
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+      - name: Install dependencies
+        run: poetry install --no-interaction --no-root
+      - name: Automated checks
+        run: poetry run make check

.pre-commit-config.yaml

@@ -9,12 +9,8 @@ repos:
         -   id: check-yaml
         -   id: check-added-large-files
             args: ["--maxkb=2000"]
-    -   repo: https://gitlab.com/pycqa/flake8
-        rev: 8f9b4931b9a28896fb43edccb23016a7540f5b82
+    -   repo: https://github.com/astral-sh/ruff-pre-commit
+        rev: v0.5.0
         hooks:
-        -   id: flake8
-    -   repo: https://github.com/psf/black
-        rev: 20.8b1
-        hooks:
-          - id: black
-            language_version: python3.7
+        -    id: ruff
+        -    id: ruff-format

Makefile

@@ -1,13 +1,13 @@
-black:
-	black xgbse setup.py tests/ --check
+format:
+	ruff format xgbse tests/ --check
 
-flake:
-	flake8 xgbse setup.py tests/
+lint:
+	ruff check xgbse tests/
 
 test:
 	pytest --cov-report term-missing --cov=xgbse tests/
 
-check: black flake test clean
+check: format lint test clean
 
 install:
 	python -m pip install -e .

README.md

@@ -409,7 +409,7 @@ To cite this repository:
   author = {Davi Vieira and Gabriel Gimenez and Guilherme Marmerola and Vitor Estima},
   title = {XGBoost Survival Embeddings: improving statistical properties of XGBoost survival analysis implementation},
   url = {http://github.com/loft-br/xgboost-survival-embeddings},
-  version = {0.2.3},
+  version = {0.3.1},
   year = {2021},
 }
 ```

examples/benchmarks/benchmark.py

@@ -11,7 +11,6 @@
 
 
 def dataframe_to_xy(dataf, event_column, time_column):
-
     e = dataf.loc[:, event_column]
     t = dataf.loc[:, time_column]
     return dataf.drop([event_column, time_column], axis=1), convert_to_structured(t, e)
@@ -65,7 +64,6 @@ def predict(self):
         pass
 
     def test(self):
-
         self.predict()
         try:
             c_index = concordance_index(
@@ -172,7 +170,6 @@ def __init__(
         )
 
     def train(self):
-
         start = time.time()
         params = {"objective": self.objective}
 
@@ -226,14 +223,12 @@ def __init__(
         self.objective = objective
 
     def train(self):
-
         start = time.time()
 
         if self.model.__class__.__name__ not in [
             "XGBSEKaplanTree",
             "XGBSEBootstrapEstimator",
         ]:
-
             self.model.fit(
                 self.X_train,
                 self.y_train,
@@ -244,7 +239,6 @@ def train(self):
             )
 
         else:
-
             self.model.fit(self.X_train, self.y_train, time_bins=self.time_bins)
 
         self.training_time = time.time() - start
@@ -284,7 +278,6 @@ def __init__(
         )
 
     def train(self):
-
         T = self.train_dataset[self.time_column]
         E = self.train_dataset[self.event_column]
 

mkdocs.yml

@@ -19,13 +19,12 @@ nav:
       - converters: modules/converters.md
       - metrics: modules/metrics.md
   - Examples:
+      - Basic usage: examples/basic_usage.md
       - Confidence intervals: examples/confidence_interval.md
       - Extrapolation: examples/extrapolation_example.md
   - Benchmarks: benchmarks/benchmarks.md
 plugins:
-  - mkdocstrings:
-      watch:
-        - xgbse
+  - mkdocstrings
   - search
 copyright:
 theme:

pyproject.toml

@@ -0,0 +1,36 @@
+[tool.poetry]
+name = "xgbse"
+version = "0.3.1"
+description = "Improving XGBoost survival analysis with embeddings and debiased estimators"
+authors = ["Loft Data Science Team <[email protected]>"]
+readme = "README.md"
+packages = [{ include = "xgbse" }]
+repository = "https://github.com/loft-br/xgboost-survival-embeddings"
+
+[tool.poetry.dependencies]
+python = ">=3.9"
+xgboost = "^2.1.0"
+numpy = "^1.26.4"
+scikit-learn = "^1.5.0"
+pandas = "^2.2.0"
+joblib = "^1.4.2"
+lifelines = "^0.29.0"
+
+[tool.poetry.group.docs.dependencies]
+mkdocs = "^1.6.0"
+mkdocs-material = "^9.5.28"
+mkdocstrings = { version = ">=0.18", extras = ["python-legacy"] }
+
+
+[tool.poetry.group.dev.dependencies]
+pre-commit = "^3.7.1"
+pytest = "^8.2.2"
+pytest-cov = "^5.0.0"
+ruff = "^0.5.0"
+
+[tool.poetry.group.benchmark.dependencies]
+pycox = "0.2.1"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

tests/test_feature_extractors.py

@@ -31,18 +31,21 @@ def test_no_objective():
 
 def test_predict_leaves_early_stop():
     xgbse = FeatureExtractor()
+    early_stopping_rounds = 10
     xgbse.fit(
         X_train,
         y_train,
-        num_boost_round=10000,
+        num_boost_round=1000,
         validation_data=(X_valid, y_valid),
-        early_stopping_rounds=10,
+        early_stopping_rounds=early_stopping_rounds,
         verbose_eval=0,
     )
     prediction = xgbse.predict_leaves(X_test)
-    assert prediction.shape == (
-        X_test.shape[0],
-        xgbse.bst.best_iteration + 1,
+    assert prediction.shape[0] == X_test.shape[0]
+    assert (
+        xgbse.bst.best_iteration
+        <= prediction.shape[1]
+        <= xgbse.bst.best_iteration + 1 + early_stopping_rounds
     )
 
 
@@ -64,7 +67,7 @@ def test_predict_hazard_early_stop():
     xgbse.fit(
         X_train,
         y_train,
-        num_boost_round=10000,
+        num_boost_round=1000,
         validation_data=(X_valid, y_valid),
         early_stopping_rounds=10,
         verbose_eval=0,

tests/test_metrics.py

@@ -68,7 +68,7 @@ def is_dist_cal_return_correct_len():
 
 def is_dist_cal_return_correct_type():
     result = dist_calibration_score(y_train, km_survival, returns="all")
-    return type(result) == dict
+    return isinstance(result, dict)
 
 
 # testing

xgbse/__init__.py

@@ -2,11 +2,10 @@
 
 from ._debiased_bce import XGBSEDebiasedBCE
 from ._kaplan_neighbors import XGBSEKaplanNeighbors, XGBSEKaplanTree
-from ._stacked_weibull import XGBSEStackedWeibull
 from ._meta import XGBSEBootstrapEstimator
+from ._stacked_weibull import XGBSEStackedWeibull
 
-
-__version__ = "0.2.3"
+__version__ = "0.3.1"
 
 __all__ = [
     "XGBSEDebiasedBCE",

xgbse/_debiased_bce.py

@@ -292,7 +292,7 @@ def predict(self, X: pd.DataFrame, return_interval_probs: bool = False):
                 XGBoost model.
 
             return_interval_probs (Bool): Boolean indicating if interval probabilities
-             are to be returned. If False the cumulative survival is returned.
+                are to be returned. If False the cumulative survival is returned.
                 Default is False.
 
         Returns: