diff --git a/.github/release.yml b/.github/release.yml new file mode 100644 index 0000000..ee16508 --- /dev/null +++ b/.github/release.yml @@ -0,0 +1,22 @@ +# Config for GitHub's auto-generated release notes. +changelog: + exclude: + labels: + - no releasenotes + categories: + - title: Major Changes 🛠 + labels: + - major + - title: New Features 🎉 + labels: + - enhancements + - feature request + - title: Bugfixes 🪲 + labels: + - bug + - title: Documentation 📖 + labels: + - docs + - title: Maintenance 🔧 + labels: + - "*" diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml deleted file mode 100644 index b06b942..0000000 --- a/.github/workflows/CI.yml +++ /dev/null @@ -1,221 +0,0 @@ -name: CI - -on: - push: - branches: - - main - - master - tags: - - "*" - pull_request: - workflow_dispatch: - -permissions: - contents: read - -jobs: - linux: - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - target: [x86_64, aarch64] - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 - with: - python-version: "3.10" - - - name: Build wheels - uses: PyO3/maturin-action@v1 - with: - target: ${{ matrix.target }} - args: --release --out dist --find-interpreter - manylinux: 2_28 - - - name: Upload wheels - uses: actions/upload-artifact@v4 - with: - name: linux-${{ matrix.target }}-wheels - path: dist - - # - name: pytest - # if: ${{ startsWith(matrix.target, 'x86_64') }} - # shell: bash - # run: | - # set -e - # pip install --find-links dist --force-reinstall 'pymc-bart-rs[all]' - # pip install pytest - # pytest - - # - name: pytest - # if: ${{ !startsWith(matrix.target, 'x86') }} - # uses: uraimo/run-on-arch-action@v2.8.1 - # with: - # arch: ${{ matrix.target }} - # distro: ubuntu22.04 - # githubToken: ${{ github.token }} - # install: | - # apt-get update - # apt-get install -y --no-install-recommends python3 python3-pip - # pip3 install -U pip pytest - # run: | - # set -e - # pip3 install --find-links dist --force-reinstall 'pymc-bart-rs[all]' - # pytest - - windows: - runs-on: windows-latest - strategy: - fail-fast: false - matrix: - target: [x64] - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 - with: - python-version: "3.10" - architecture: ${{ matrix.target }} - - - name: Build wheels - uses: PyO3/maturin-action@v1 - with: - target: ${{ matrix.target }} - args: --release --out dist --find-interpreter - sccache: "true" - - - name: Upload wheels - uses: actions/upload-artifact@v4 - with: - name: windows-${{ matrix.target }}-wheels - path: dist - - # - name: pytest - # if: ${{ !startsWith(matrix.target, 'aarch64') }} - # shell: bash - # run: | - # set -e - # pip install "pymc-bart-rs[all]" --find-links dist --force-reinstall - # pip install pytest - # pytest - - macos: - runs-on: macos-latest - strategy: - fail-fast: false - matrix: - target: [x86_64, aarch64] - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 - with: - python-version: "3.10" - - - name: Build wheels - uses: PyO3/maturin-action@v1 - with: - target: ${{ matrix.target }} - args: --release --out dist --find-interpreter - sccache: "true" - - - name: Upload wheels - uses: actions/upload-artifact@v4 - with: - name: macos-${{ matrix.target }}-wheels - path: dist - - # - name: pytest - # if: ${{ startsWith(matrix.target, 'aarch64') }} - # shell: bash - # run: | - # set -e - # pip install 'pymc-bart-rs[all]' --find-links dist --force-reinstall - # pip install pytest - # pytest - - sdist: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - name: Build sdist - uses: PyO3/maturin-action@v1 - with: - command: sdist - args: --out dist - - - name: Upload sdist - uses: actions/upload-artifact@v4 - with: - name: sdist - path: dist - - release: - name: Release - runs-on: ubuntu-latest - if: "startsWith(github.ref, 'refs/tags/')" - needs: [linux, windows, macos, sdist] - # Combine all the wheels and sdists into a single directory - steps: - - name: Download artifacts - uses: actions/download-artifact@v4 - with: - name: linux-x86_64-wheels - path: dist - - uses: actions/download-artifact@v4 - with: - name: linux-aarch64-wheels - path: dist - - uses: actions/download-artifact@v4 - with: - name: windows-x64-wheels - path: dist - - uses: actions/download-artifact@v4 - with: - name: macos-x86_64-wheels - path: dist - - uses: actions/download-artifact@v4 - with: - name: macos-aarch64-wheels - path: dist - - uses: actions/download-artifact@v4 - with: - name: sdist - path: dist - - name: Publish to PyPI - uses: PyO3/maturin-action@v1 - env: - MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }} - with: - command: upload - args: --skip-existing dist/* - - test: - runs-on: ubuntu-latest - needs: [linux] - steps: - - name: Checkout - uses: actions/checkout@v4 - - # Download the artifact from your Linux x86_64 job (adjust if you need other arches) - - name: Download Wheels - uses: actions/download-artifact@v4 - with: - name: linux-x86_64-wheels - path: dist - - - name: Install Python & Dependencies - uses: actions/setup-python@v5 - with: - python-version: "3.10" - - # Optionally run your Rust tests here (or in the linux job itself) - - name: Rust tests - run: cargo test --all - - - name: Install wheel and Python test dependencies - run: | - pip install --force-reinstall --find-links dist 'pymc-bart-rs[all]' - pip install pytest - - - name: Run Python tests - run: pytest diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml new file mode 100644 index 0000000..6ee2cef --- /dev/null +++ b/.github/workflows/build-and-test.yml @@ -0,0 +1,77 @@ +name: Build and Test + +on: + push: + branches: + - main + pull_request: + workflow_dispatch: + +permissions: + contents: read + +jobs: + linux: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: | + 3.10 + 3.11 + 3.12 + - name: Build wheels + uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 + with: + target: x86_64 + args: --release --out dist --find-interpreter + manylinux: 2_28 + - name: Upload wheels + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: linux-x86_64-wheels + path: dist + + package: + name: Build & inspect package + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - uses: hynek/build-and-inspect-python-package@d44ca7d91762de7a7d5436ddae667c6da6d1c3df # v2.18.0 + + rust-tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - name: Rust tests + run: cargo test --all + + python-tests: + runs-on: ubuntu-latest + needs: [linux] + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12"] + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: linux-x86_64-wheels + path: dist + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: ${{ matrix.python-version }} + - name: Install package and test dependencies + run: pip install --find-links dist "bartrs[dev]" + - name: Run Python tests + run: pytest diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml new file mode 100644 index 0000000..5735d75 --- /dev/null +++ b/.github/workflows/pypi.yml @@ -0,0 +1,129 @@ +name: PyPI + +# Wheel version comes from Cargo.toml; bump it and tag the release to match. +on: + release: + types: [published] + workflow_dispatch: + +permissions: + contents: read + +jobs: + linux: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + target: [x86_64, aarch64] + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: | + 3.10 + 3.11 + 3.12 + - name: Build wheels + uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 + with: + target: ${{ matrix.target }} + args: --release --out dist --find-interpreter + manylinux: 2_28 + - name: Upload wheels + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: wheels-linux-${{ matrix.target }} + path: dist + + windows: + runs-on: windows-latest + strategy: + fail-fast: false + matrix: + target: [x64] + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: | + 3.10 + 3.11 + 3.12 + architecture: ${{ matrix.target }} + - name: Build wheels + uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 + with: + target: ${{ matrix.target }} + args: --release --out dist --find-interpreter + sccache: "true" + - name: Upload wheels + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: wheels-windows-${{ matrix.target }} + path: dist + + macos: + runs-on: macos-latest + strategy: + fail-fast: false + matrix: + target: [x86_64, aarch64] + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: | + 3.10 + 3.11 + 3.12 + - name: Build wheels + uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 + with: + target: ${{ matrix.target }} + args: --release --out dist --find-interpreter + sccache: "true" + - name: Upload wheels + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: wheels-macos-${{ matrix.target }} + path: dist + + sdist: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + persist-credentials: false + - name: Build sdist + uses: PyO3/maturin-action@e83996d129638aa358a18fbd1dfb82f0b0fb5d3b # v1.51.0 + with: + command: sdist + args: --out dist + - name: Upload sdist + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: wheels-sdist + path: dist + + publish: + name: Publish to PyPI + needs: [linux, windows, macos, sdist] + runs-on: ubuntu-latest + if: github.repository_owner == 'pymc-devs' && github.event_name == 'release' && github.event.action == 'published' + environment: pypi + permissions: + id-token: write # trusted publishing (OIDC) + PEP 740 attestations + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + pattern: wheels-* + merge-multiple: true + path: dist + - uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml deleted file mode 100644 index 5983ce2..0000000 --- a/.github/workflows/test.yml +++ /dev/null @@ -1,39 +0,0 @@ -name: Run tests - -on: - push: - branches: - - main - pull_request: - -jobs: - test: - runs-on: ubuntu-latest - needs: [linux] - steps: - - name: Checkout - uses: actions/checkout@v4 - - # Download the artifact from your Linux x86_64 job (adjust if you need other arches) - - name: Download Wheels - uses: actions/download-artifact@v4 - with: - name: linux-x86_64-wheels - path: dist - - - name: Install Python & Dependencies - uses: actions/setup-python@v5 - with: - python-version: "3.10" - - # Optionally run your Rust tests here (or in the linux job itself) - - name: Rust tests - run: cargo test --all - - - name: Install wheel and Python test dependencies - run: | - pip install --force-reinstall --find-links dist 'pymc-bart-rs[all]' - pip install pytest - - - name: Run Python tests - run: pytest diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..69ba254 --- /dev/null +++ b/.gitignore @@ -0,0 +1,27 @@ +# Rust +/target/ + +# Python build artifacts +/dist/ +build/ +*.egg-info/ +__pycache__/ +*.py[cod] + +# Compiled extension modules +*.so +*.pyd +*.dll + +# Virtual environments +.venv/ +venv/ + +# Test / coverage +.pytest_cache/ +.coverage +htmlcov/ + +# Editors +.idea/ +.vscode/ diff --git a/pyproject.toml b/pyproject.toml index d45612f..0533f87 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "bartrs" -version = "0.2.0" +dynamic = ["version"] authors = [ {name = "Otto Vintola", email="hello@ottovintola.com" }, ] @@ -17,7 +17,7 @@ classifiers = [ ] license = "Apache-2.0" license-files = ["LICEN[CS]E*"] -readme = "README.md" +readme = {file = "README.md", content-type = "text/markdown"} dependencies = [ "numba >= 0.60.0", diff --git a/python/tests/test_compile_pymc.py b/python/tests/test_compile_pymc.py deleted file mode 100644 index 2974b8c..0000000 --- a/python/tests/test_compile_pymc.py +++ /dev/null @@ -1,33 +0,0 @@ -import pymc as pm -import numpy as np - -from pymc_bartrs.compile_pymc import CompiledPyMCModel - -n_obs = 1000 -X = np.random.randn(n_obs, 3) -y_observed = np.random.randn(n_obs) - -with pm.Model() as model: - # These are SAMPLED (not shared) - beta = pm.Normal("beta", 0, 1, shape=3) - sigma = pm.HalfNormal("sigma", 1) - - mu = pm.Deterministic("mu", pm.math.dot(X, beta)) # X is shared - y = pm.Normal("y", mu, sigma, observed=y_observed) # y_observed is shared - -compiled_model = CompiledPyMCModel(model, [beta, sigma]) - -# Check what's in the PyTensor function storage -print(f"Number of inputs: {len(compiled_model.logp_fn_ptr.input_storage)}") -# print(f"Input 0 (parameters): shape = {compiled_model.logp_fn_ptr.input_storage[0].storage[0].shape}") - -# Check the shared arrays (these are the observed data) -for i, storage in enumerate(compiled_model.logp_fn_ptr.input_storage[1:], 1): - data = storage.storage[0] - print(f"Input {i} (shared): shape = {data.shape}, dtype = {data.dtype}") - print(f" First few values: {data.flat[:5]}") - -# Pre-allocated shared arrays -print(f"Number of logp_args: {len(compiled_model.logp_args)}") -for i, arr in enumerate(compiled_model.logp_args): - print(f"logp_args[{i}]: shape = {arr.shape}, dtype = {arr.dtype}") diff --git a/python/tests/test_tree.py b/python/tests/test_tree.py deleted file mode 100644 index 0e23151..0000000 --- a/python/tests/test_tree.py +++ /dev/null @@ -1,26 +0,0 @@ -import numpy as np - - - - -depth = 10 -value_tree = np.zeros(2 ** depth) -print(value_tree.shape) - - -rules = {0: "ContinuousSplit", 1: "Hey", 2 :"ContinuousSplit"} -required = set(range(3)) -passed = set(list(rules.keys())) - -if not required.issubset(passed): - missing = sorted(list(required - passed)) - raise ValueError(f"Missing: {missing}") - -supported_split_rules = {"ContinuousSplit", "OneHotSplit", "SubsetSplit"} -passed_rules = set(rules.values()) -invalid_rules = passed_rules - supported_split_rules - -if invalid_rules: - raise ValueError( - f"rule(s) must be one of {supported_split_rules}. Received invalid rule(s): {invalid_rules}" - ) diff --git a/src/tree.rs b/src/tree.rs index 20e56ed..8842d20 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -19,7 +19,7 @@ use crate::response::{LeafKind, LeafPayload, LeafProposal}; /// Internal nodes store split variable and threshold. Leaf nodes store /// predicted values. The `leaf_indices` vector maps each training sample /// to its assigned leaf node. -#[pyclass(module = "pymc_bart.pymc_bart", get_all)] +#[pyclass(module = "bartrs.bartrs", get_all)] #[derive(Clone, Debug)] pub struct TreeArrays { /// Split variable per node (u32::MAX = leaf sentinel) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/tests/test_bart_models.py b/tests/test_pgbart.py similarity index 77% rename from python/tests/test_bart_models.py rename to tests/test_pgbart.py index e5999c8..5ecd51c 100644 --- a/python/tests/test_bart_models.py +++ b/tests/test_pgbart.py @@ -1,15 +1,17 @@ import numpy as np import pandas as pd import pymc as pm -import pymc_bartrs as pmb +import pymc_bart as pmb +from bartrs import PGBART -NUM_TUNE = 300 -NUM_DRAWS = 600 -NUM_CHAINS = 4 + +NUM_TUNE = 50 +NUM_DRAWS = 50 +NUM_CHAINS = 2 BATCH_SIZE = (0.1, 0.1) -NUM_TREES = 50 -NUM_PARTICLES = 10 +NUM_TREES = 10 +NUM_PARTICLES = 5 RANDOM_SEED = 42 @@ -20,14 +22,14 @@ def test_bikes(): with pm.Model() as model: alpha = pm.Exponential("alpha", 1.0) - mu = pmb.BART("mu", X, np.log(Y), m=NUM_TREES) + mu = pmb.BART("mu", X, np.log(Y), m=NUM_TREES, response="gaussian") y = pm.NegativeBinomial("y", mu=pm.math.exp(mu), alpha=alpha, observed=Y) idata = pm.sample( tune=NUM_TUNE, draws=NUM_DRAWS, chains=NUM_CHAINS, - step=[pmb.PGBART([mu], batch=BATCH_SIZE, num_particles=NUM_PARTICLES)], + step=[PGBART([mu], batch=BATCH_SIZE, num_particles=NUM_PARTICLES)], random_seed=RANDOM_SEED, ) @@ -43,7 +45,7 @@ def test_coal(): y_data = hist with pm.Model() as model: - mu = pmb.BART("mu", X=x_data, Y=np.log(y_data), m=NUM_TREES) + mu = pmb.BART("mu", X=x_data, Y=np.log(y_data), m=NUM_TREES, response="gaussian") exp_mu = pm.Deterministic("exp_mu", pm.math.exp(mu)) y_pred = pm.Poisson("y_pred", mu=exp_mu, observed=y_data) @@ -51,6 +53,6 @@ def test_coal(): tune=NUM_TUNE, draws=NUM_DRAWS, chains=NUM_CHAINS, - step=[pmb.PGBART([mu], batch=BATCH_SIZE, num_particles=NUM_PARTICLES)], + step=[PGBART([mu], batch=BATCH_SIZE, num_particles=NUM_PARTICLES)], random_seed=RANDOM_SEED, )