pmf-dark 0.1.0__tar.gz

pmf_dark-0.1.0/PKG-INFO

@@ -0,0 +1,233 @@
+Metadata-Version: 2.4
+Name: pmf-dark
+Version: 0.1.0
+Summary: Modelling species dark diversity using bayesian Probablistic Matrix Factorisation
+License: Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+Author: davidyshen
+Author-email: contact@davidyshen.com
+Requires-Python: >=3.13,<3.15
+Classifier: License :: Other/Proprietary License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: matplotlib (>=3.10.9,<4.0.0)
+Requires-Dist: numpy (>=2.4.6,<3.0.0)
+Requires-Dist: pandas (>=3.0.3,<4.0.0)
+Requires-Dist: pyro-ppl (>=1.9.0,<2.0.0)
+Description-Content-Type: text/markdown
+
+# PMF-dark: Using matrix factorisation for dark diversity estimation
+
+## Overview
+
+This repository implements a **PMF-dark** using Bayesian Probabilistic Matrix Factorisation to estimate **dark diversity** - the set of species absent from a site despite having suitable environmental conditions. The method uses **counterfactual predictions** to reconstruct the potential species pool by separating environmental effects from unmeasured drivers of absence (e.g., land-use degradation, dispersal limitation, biotic interactions).
+
+## The Problem: What is Dark Diversity?
+
+Traditional biodiversity assessments only count observed species (alpha diversity). However, many species are absent from sites where they *could* thrive based on environmental conditions. This "**dark diversity**" represents:
+
+- Species lost due to historical or ongoing land-use degradation
+- Species unable to reach suitable sites due to dispersal limitation
+- Species suppressed by biotic interactions
+
+Quantifying dark diversity is crucial for:
+- Conservation planning and restoration potential assessment
+- Understanding true biodiversity patterns
+- Identifying areas with highest restoration value
+
+## Methodology
+
+### Core Model
+
+The framework decomposes species occurrence probabilities into **three additive components**:
+
+$$\text{logit}(p_{ij}) = \underbrace{\alpha_j}_{\text{Intercept}} + \underbrace{f_j(\mathbf{x}_i)}_{\text{Environmental Effects}} + \underbrace{\mathbf{w}_i^\top \mathbf{z}_j}_{\text{Latent Factors}}$$
+
+Where:
+- **$\alpha_j$**: Species-specific baseline prevalence
+- **$f_j(\mathbf{x}_i)$**: Environmental response function to measured abiotic variables (temperature, pH, elevation, etc.), which can be modelled as linear, Gaussian niche, or non-linear (e.g. Bayesian neural network)
+- **$\mathbf{w}_i^\top \mathbf{z}_j$**: Latent factors capturing unmeasured drivers of absence
+
+### Key Innovation: Counterfactual Predictions
+
+1. **Full Predictions**: Include all components (environment + latent factors)
+   - Represents observed diversity with all drivers active
+   
+2. **Environment-Only Predictions**: Exclude latent factors
+   - Represents potential diversity (setting $\mathbf{w}_i^\top \mathbf{z}_j = 0$)
+   
+3. **Dark Diversity Proxy**: Difference between full and environment-only predictions
+   - Quantifies species lost to unmeasured stressors
+
+### Inference: Stochastic Variational Inference (SVI)
+
+The model is fit using **Pyro-based SVI**, which:
+- Handles high-dimensional ecological matrices efficiently
+- Treats inference as an optimisation problem (ELBO maximisation)
+- Scales to thousands of sites and species
+- Requires minimal computational resources
+
+## Repository Structure
+
+```
+PMF_dark/
+├── README.md                                    # This file
+├── mat_fact_dark_div.ipynb                     # Main analysis notebook
+├── data/
+│   ├── survey.csv                              # Species presence/absence matrix (sites × species)
+│   ├── env.csv                                 # Environmental predictors (sites × covariates)
+│   └── truth.csv                               # Ground truth data (if available)
+└── output/
+    ├── mat_fact_predicted_probabilities_full.csv           # Full model predictions
+    ├── mat_fact_predicted_probabilities_env_only.csv       # Environment-only predictions
+    └── mat_fact_dark_diversity_proxy.csv                   # Dark diversity estimates
+```
+
+## Installation
+
+### Requirements
+
+- Python 3.13 or 3.14
+- PyTorch (with CUDA support if using GPU)
+- Pyro (pyro-ppl)
+- Pandas
+- NumPy
+- scikit-learn
+- scipy
+
+### Setup
+
+To ensure PyTorch is installed with the correct CUDA version for your system, it is recommended to install PyTorch manually **first** before installing the package or its other dependencies.
+
+#### 1. Setup Virtual Environment
+```bash
+# Clone the repository
+git clone https://github.com/davidyshen/PMF_dark.git
+cd PMF_dark
+
+# Create virtual environment (optional but recommended)
+python -m venv .venv
+source .venv/bin/activate  # On Windows: .venv\Scripts\activate
+```
+
+#### 2. Install PyTorch with CUDA
+Visit the [PyTorch Getting Started guide](https://pytorch.org/get-started/locally/) to select the correct command for your CUDA version and OS. For example, to install PyTorch with CUDA 12.4 support on Windows/Linux:
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cu124
+```
+
+If not using CUDA, simply install the CPU version:
+```bash
+pip install torch
+```
+
+
+#### 3. Install remaining dependencies
+If installing via pip:
+```bash
+pip install pyro-ppl pandas numpy scikit-learn scipy jupyter
+```
+
+If using Poetry:
+```bash
+# This will install the package and its remaining dependencies into your environment
+poetry install
+```
+
+## Usage
+
+### Running the Full Analysis
+
+1. Prepare your data in `data/` directory:
+   - `survey.csv`: Species presence/absence (rows = sites, columns = species, values = 0/1)
+   - `env.csv`: Environmental predictors (rows = sites, columns = variables)
+
+2. Open and run the Jupyter notebook:
+   ```bash
+   jupyter notebook mat_fact_dark_div.ipynb
+   ```
+
+3. The notebook will:
+   - Load and standardise data
+   - Fit the matrix factorisation model (2,500 iterations)
+   - Generate predictions and save CSV outputs
+
+### Customisation
+
+Key parameters in the notebook:
+
+```python
+# Model parameters
+num_factors = 5                # Number of latent factors (adjust based on data complexity)
+num_iterations = 2500          # Model iterations
+
+# Learning rate
+Adam({"lr": 0.01})            # Adjust if convergence is slow
+```
+
+## Output Files
+
+- **mat_fact_predicted_probabilities_full.csv**: Predicted species occurrence probabilities including all effects
+- **mat_fact_predicted_probabilities_env_only.csv**: Predicted probabilities using only environmental effects
+- **mat_fact_dark_diversity_proxy.csv**: Dark diversity estimates (full - env_only)
+
+## Data Format
+
+### survey.csv
+```
+site_id,species_1,species_2,...,species_n,ID,x,y
+site_1,0,1,0,...,1,id_1,100.5,200.3
+site_2,1,0,1,...,0,id_2,101.2,201.5
+...
+```
+- Rows: Sites/locations
+- Columns: Species (0/1 presence/absence) + ID + spatial coordinates
+- **Note**: ID and spatial coordinates are automatically extracted/dropped
+
+### env.csv
+```
+site_id,temp,pH,elevation,...,ID,landuse
+site_1,15.2,7.1,500,...,id_1,degraded
+site_2,14.8,6.9,520,...,id_2,pristine
+...
+```
+- Rows: Sites matching survey.csv
+- Columns: Environmental predictors + ID + land-use
+- **Note**: ID and land-use columns are dropped; only abiotic predictors are used
+
+## Interpretation of Results
+
+### Dark Diversity Proxy Values
+- **High values (close to 1)**: Species should be present based on environment but are absent—candidate for restoration
+- **Low values (close to 0)**: Species absence explained by environmental conditions
+- **Negative values**: Model predicts species should be absent (rare, indicates environmental unsuitability)
+
+### Key Metrics
+- **AUC (Area Under ROC Curve)**: Overall model discrimination (0.5 = random, 1.0 = perfect)
+- **Brier Score**: Prediction calibration error (lower is better)
+- **F1 Score**: Balance between precision and recall
+
+## Advantages of This Approach
+
+✓ **No subjective benchmarking**: Automated separation of environmental vs. unmeasured effects  
+✓ **Mathematically principled**: Latent factors naturally absorb degradation signals  
+✓ **Scalable**: SVI handles thousands of species and sites  
+✓ **Species-specific**: Each species can have unique environmental responses  
+✓ **Reproducible**: Fully probabilistic framework with clear assumptions
+
+## Limitations
+
+- Assumes species responses are log-linear (logit link)
+- Requires sufficient environmental variation to estimate effects reliably
+- May overestimate dark diversity if detection is imperfect
+- Computational cost increases with number of species and sites
+- Requires careful tuning of number of latent factors
+
+## References & Theoretical Background
+
+### Key Concepts
+- **Joint Species Distribution Models (JSDMs)**: Latent variable models for multivariate species data
+- **Matrix Factorisation**: Low-rank decomposition of high-dimensional species matrices
+- **Stochastic Variational Inference**: Scalable Bayesian inference for probabilistic models
+- **Counterfactual Predictions**: Causal inference approach to estimate potential outcomes
+

pmf_dark-0.1.0/README.md

@@ -0,0 +1,214 @@
+# PMF-dark: Using matrix factorisation for dark diversity estimation
+
+## Overview
+
+This repository implements a **PMF-dark** using Bayesian Probabilistic Matrix Factorisation to estimate **dark diversity** - the set of species absent from a site despite having suitable environmental conditions. The method uses **counterfactual predictions** to reconstruct the potential species pool by separating environmental effects from unmeasured drivers of absence (e.g., land-use degradation, dispersal limitation, biotic interactions).
+
+## The Problem: What is Dark Diversity?
+
+Traditional biodiversity assessments only count observed species (alpha diversity). However, many species are absent from sites where they *could* thrive based on environmental conditions. This "**dark diversity**" represents:
+
+- Species lost due to historical or ongoing land-use degradation
+- Species unable to reach suitable sites due to dispersal limitation
+- Species suppressed by biotic interactions
+
+Quantifying dark diversity is crucial for:
+- Conservation planning and restoration potential assessment
+- Understanding true biodiversity patterns
+- Identifying areas with highest restoration value
+
+## Methodology
+
+### Core Model
+
+The framework decomposes species occurrence probabilities into **three additive components**:
+
+$$\text{logit}(p_{ij}) = \underbrace{\alpha_j}_{\text{Intercept}} + \underbrace{f_j(\mathbf{x}_i)}_{\text{Environmental Effects}} + \underbrace{\mathbf{w}_i^\top \mathbf{z}_j}_{\text{Latent Factors}}$$
+
+Where:
+- **$\alpha_j$**: Species-specific baseline prevalence
+- **$f_j(\mathbf{x}_i)$**: Environmental response function to measured abiotic variables (temperature, pH, elevation, etc.), which can be modelled as linear, Gaussian niche, or non-linear (e.g. Bayesian neural network)
+- **$\mathbf{w}_i^\top \mathbf{z}_j$**: Latent factors capturing unmeasured drivers of absence
+
+### Key Innovation: Counterfactual Predictions
+
+1. **Full Predictions**: Include all components (environment + latent factors)
+   - Represents observed diversity with all drivers active
+   
+2. **Environment-Only Predictions**: Exclude latent factors
+   - Represents potential diversity (setting $\mathbf{w}_i^\top \mathbf{z}_j = 0$)
+   
+3. **Dark Diversity Proxy**: Difference between full and environment-only predictions
+   - Quantifies species lost to unmeasured stressors
+
+### Inference: Stochastic Variational Inference (SVI)
+
+The model is fit using **Pyro-based SVI**, which:
+- Handles high-dimensional ecological matrices efficiently
+- Treats inference as an optimisation problem (ELBO maximisation)
+- Scales to thousands of sites and species
+- Requires minimal computational resources
+
+## Repository Structure
+
+```
+PMF_dark/
+├── README.md                                    # This file
+├── mat_fact_dark_div.ipynb                     # Main analysis notebook
+├── data/
+│   ├── survey.csv                              # Species presence/absence matrix (sites × species)
+│   ├── env.csv                                 # Environmental predictors (sites × covariates)
+│   └── truth.csv                               # Ground truth data (if available)
+└── output/
+    ├── mat_fact_predicted_probabilities_full.csv           # Full model predictions
+    ├── mat_fact_predicted_probabilities_env_only.csv       # Environment-only predictions
+    └── mat_fact_dark_diversity_proxy.csv                   # Dark diversity estimates
+```
+
+## Installation
+
+### Requirements
+
+- Python 3.13 or 3.14
+- PyTorch (with CUDA support if using GPU)
+- Pyro (pyro-ppl)
+- Pandas
+- NumPy
+- scikit-learn
+- scipy
+
+### Setup
+
+To ensure PyTorch is installed with the correct CUDA version for your system, it is recommended to install PyTorch manually **first** before installing the package or its other dependencies.
+
+#### 1. Setup Virtual Environment
+```bash
+# Clone the repository
+git clone https://github.com/davidyshen/PMF_dark.git
+cd PMF_dark
+
+# Create virtual environment (optional but recommended)
+python -m venv .venv
+source .venv/bin/activate  # On Windows: .venv\Scripts\activate
+```
+
+#### 2. Install PyTorch with CUDA
+Visit the [PyTorch Getting Started guide](https://pytorch.org/get-started/locally/) to select the correct command for your CUDA version and OS. For example, to install PyTorch with CUDA 12.4 support on Windows/Linux:
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cu124
+```
+
+If not using CUDA, simply install the CPU version:
+```bash
+pip install torch
+```
+
+
+#### 3. Install remaining dependencies
+If installing via pip:
+```bash
+pip install pyro-ppl pandas numpy scikit-learn scipy jupyter
+```
+
+If using Poetry:
+```bash
+# This will install the package and its remaining dependencies into your environment
+poetry install
+```
+
+## Usage
+
+### Running the Full Analysis
+
+1. Prepare your data in `data/` directory:
+   - `survey.csv`: Species presence/absence (rows = sites, columns = species, values = 0/1)
+   - `env.csv`: Environmental predictors (rows = sites, columns = variables)
+
+2. Open and run the Jupyter notebook:
+   ```bash
+   jupyter notebook mat_fact_dark_div.ipynb
+   ```
+
+3. The notebook will:
+   - Load and standardise data
+   - Fit the matrix factorisation model (2,500 iterations)
+   - Generate predictions and save CSV outputs
+
+### Customisation
+
+Key parameters in the notebook:
+
+```python
+# Model parameters
+num_factors = 5                # Number of latent factors (adjust based on data complexity)
+num_iterations = 2500          # Model iterations
+
+# Learning rate
+Adam({"lr": 0.01})            # Adjust if convergence is slow
+```
+
+## Output Files
+
+- **mat_fact_predicted_probabilities_full.csv**: Predicted species occurrence probabilities including all effects
+- **mat_fact_predicted_probabilities_env_only.csv**: Predicted probabilities using only environmental effects
+- **mat_fact_dark_diversity_proxy.csv**: Dark diversity estimates (full - env_only)
+
+## Data Format
+
+### survey.csv
+```
+site_id,species_1,species_2,...,species_n,ID,x,y
+site_1,0,1,0,...,1,id_1,100.5,200.3
+site_2,1,0,1,...,0,id_2,101.2,201.5
+...
+```
+- Rows: Sites/locations
+- Columns: Species (0/1 presence/absence) + ID + spatial coordinates
+- **Note**: ID and spatial coordinates are automatically extracted/dropped
+
+### env.csv
+```
+site_id,temp,pH,elevation,...,ID,landuse
+site_1,15.2,7.1,500,...,id_1,degraded
+site_2,14.8,6.9,520,...,id_2,pristine
+...
+```
+- Rows: Sites matching survey.csv
+- Columns: Environmental predictors + ID + land-use
+- **Note**: ID and land-use columns are dropped; only abiotic predictors are used
+
+## Interpretation of Results
+
+### Dark Diversity Proxy Values
+- **High values (close to 1)**: Species should be present based on environment but are absent—candidate for restoration
+- **Low values (close to 0)**: Species absence explained by environmental conditions
+- **Negative values**: Model predicts species should be absent (rare, indicates environmental unsuitability)
+
+### Key Metrics
+- **AUC (Area Under ROC Curve)**: Overall model discrimination (0.5 = random, 1.0 = perfect)
+- **Brier Score**: Prediction calibration error (lower is better)
+- **F1 Score**: Balance between precision and recall
+
+## Advantages of This Approach
+
+✓ **No subjective benchmarking**: Automated separation of environmental vs. unmeasured effects  
+✓ **Mathematically principled**: Latent factors naturally absorb degradation signals  
+✓ **Scalable**: SVI handles thousands of species and sites  
+✓ **Species-specific**: Each species can have unique environmental responses  
+✓ **Reproducible**: Fully probabilistic framework with clear assumptions
+
+## Limitations
+
+- Assumes species responses are log-linear (logit link)
+- Requires sufficient environmental variation to estimate effects reliably
+- May overestimate dark diversity if detection is imperfect
+- Computational cost increases with number of species and sites
+- Requires careful tuning of number of latent factors
+
+## References & Theoretical Background
+
+### Key Concepts
+- **Joint Species Distribution Models (JSDMs)**: Latent variable models for multivariate species data
+- **Matrix Factorisation**: Low-rank decomposition of high-dimensional species matrices
+- **Stochastic Variational Inference**: Scalable Bayesian inference for probabilistic models
+- **Counterfactual Predictions**: Causal inference approach to estimate potential outcomes

pmf_dark-0.1.0/pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "pmf-dark"
+version = "0.1.0"
+description = "Modelling species dark diversity using bayesian Probablistic Matrix Factorisation"
+authors = [
+    {name = "davidyshen",email = "contact@davidyshen.com"}
+]
+license = {text = "Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE."}
+readme = "README.md"
+requires-python = ">=3.13,<3.15"
+dependencies = [
+    "numpy (>=2.4.6,<3.0.0)",
+    "pandas (>=3.0.3,<4.0.0)",
+    "matplotlib (>=3.10.9,<4.0.0)",
+    "pyro-ppl (>=1.9.0,<2.0.0)"
+]
+
+[tool.poetry]
+packages = [{ include = "pmf_dark", from = "src" }]
+
+[tool.poetry.group.dev.dependencies]
+ipykernel = "^7.2.0"
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"

pmf_dark-0.1.0/src/pmf_dark/__init__.py

@@ -0,0 +1,10 @@
+import torch
+from .darkdiv import compute_dark_diversity
+
+# Print CUDA availability on import
+if torch.cuda.is_available():
+    print(f"pmf-dark: CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
+else:
+    print("pmf-dark: CUDA is not available. Using CPU.")
+
+__all__ = ["compute_dark_diversity"]

pmf_dark-0.1.0/src/pmf_dark/darkdiv.py

@@ -0,0 +1,264 @@
+import torch
+import pyro
+import matplotlib.pyplot as plt
+
+import numpy as np
+import pandas as pd
+
+
+def infer_y_type(y):
+
+    # torch tensor
+    if isinstance(y, torch.Tensor):
+
+        values = y.detach().cpu().numpy()
+
+    # pandas dataframe
+    elif hasattr(y, "to_numpy"):
+
+        values = y.to_numpy()
+
+    # numpy or other
+    else:
+
+        values = np.asarray(y)
+
+    # Check missing values
+    if np.isnan(values).any():
+        raise ValueError("y contains missing values.")
+
+    # Binary data
+    if np.isin(values, [0, 1]).all():
+        return "presence_absence"
+
+    # Count data
+    elif (values >= 0).all():
+        return "count"
+
+    else:
+        raise ValueError("y must contain either binary or count data.")
+
+
+def prepare_data(x, y, cuda=False):
+
+    # Keep names
+    x_columns = x.columns
+    y_columns = y.columns
+    site_index = y.index
+
+    # Standardize x
+    x = (x - x.mean()) / x.std()
+
+    # Convert to tensors
+    x_tensor = torch.tensor(
+        x.to_numpy(),
+        dtype=torch.float32,
+    )
+
+    y_tensor = torch.tensor(
+        y.to_numpy(),
+        dtype=torch.float32,
+    )
+
+    # Device
+    device = torch.device("cuda" if cuda and torch.cuda.is_available() else "cpu")
+
+    x_tensor = x_tensor.to(device)
+    y_tensor = y_tensor.to(device)
+
+    return {
+        "x": x_tensor,
+        "y": y_tensor,
+        "x_columns": x_columns,
+        "y_columns": y_columns,
+        "site_index": site_index,
+    }
+
+
+def compute_predictions(
+    samples, x, model_type="gaussian", include_latent=True, y_type="presence_absence"
+):
+
+    if model_type == "linear":
+        alpha = samples["alpha"].squeeze(1)
+        beta = samples["beta"].squeeze(1)
+        eta = alpha[:, None, :] + torch.einsum("ij,sjk->sik", x, beta)
+
+    elif model_type == "gaussian":
+        alpha = samples["alpha"].squeeze(1)
+        mu = samples["mu"].squeeze(1)
+        gamma = samples["gamma"].squeeze(1)
+
+        # x:     [sites, env]
+        # mu:    [samples, env, species]
+        # gamma: [samples, env, species]
+        diff = x[None, :, :, None] - mu[:, None, :, :]
+        env_effect = -torch.sum(gamma[:, None, :, :] * diff**2, dim=2)
+
+        eta = alpha[:, None, :] + env_effect
+
+    elif model_type == "bnn":
+        w1 = samples["w1"].squeeze(1)
+        b1 = samples["b1"].squeeze(1)
+        w2 = samples["w2"].squeeze(1)
+        b2 = samples["b2"].squeeze(1)
+
+        hidden = torch.tanh(torch.einsum("ij,sjh->sih", x, w1) + b1[:, None, :])
+
+        eta = torch.einsum("sih,shk->sik", hidden, w2) + b2[:, None, :]
+
+    else:
+        raise ValueError(f"Unknown model_type: {model_type}")
+
+    if include_latent:
+        W = samples["W"].squeeze(1)
+        Z = samples["Z"].squeeze(1)
+        eta = eta + torch.einsum("sik,sjk->sij", W, Z)
+
+    if y_type == "presence_absence":
+        return torch.sigmoid(eta)
+
+    elif y_type == "count":
+        return torch.exp(eta)
+
+    else:
+        raise ValueError("y_type must be 'presence_absence' or 'count'")
+
+
+def compute_dark_diversity(
+    y,
+    x,
+    model_type="linear",
+    num_factors=1,
+    method="svi",
+    cuda=False,
+    include_latent=True,
+    return_means=True,
+    batch_size=None,
+    pred_batch_size=None,
+    **kwargs,
+):
+
+    if cuda and not torch.cuda.is_available():
+        import warnings
+
+        warnings.warn(
+            "CUDA was requested (cuda=True), but PyTorch cannot detect a CUDA-enabled GPU. "
+            "Please check your NVIDIA drivers and reinstall PyTorch with CUDA support. "
+            "Falling back to CPU."
+        )
+
+    # For bnn only svi is supported
+    if method == "mcmc" and model_type == "bnn":
+        raise ValueError(
+            "MCMC is currently not supported for the Bayesian neural network model. "
+            "Please use method='svi' instead."
+        )
+
+    # Check if y is presence/absence or count data
+    y_type = infer_y_type(y)
+    print(y_type)
+
+    data = prepare_data(x, y, cuda=cuda)
+
+    x = data["x"]
+    y = data["y"]
+
+    # Load the model
+    if model_type == "linear":
+        from .models import linear_model
+
+        model = linear_model
+    elif model_type == "gaussian":
+        from .models import gaussian
+
+        model = gaussian
+    elif model_type == "bnn":
+        from .models import bnn_model
+
+        model = bnn_model
+
+    # Inference
+    if method == "svi":
+        from .inference import fit_svi
+
+        fit = fit_svi(
+            model,
+            y,
+            x,
+            num_factors,
+            y_type=y_type,
+            cuda=cuda,
+            batch_size=batch_size,
+            **kwargs,
+        )
+    elif method == "mcmc":
+        from .inference import fit_mcmc
+
+        fit = fit_mcmc(
+            model, y, x, num_factors, y_type=y_type, batch_size=batch_size, **kwargs
+        )
+
+    # Compute probabilities
+    if pred_batch_size is not None:
+        n_sites = x.shape[0]
+        pred_chunks = []
+        for i in range(0, n_sites, pred_batch_size):
+            x_chunk = x[i : i + pred_batch_size]
+            samples_chunk = fit["samples"].copy()
+            if "W" in fit["samples"]:
+                w_tensor = fit["samples"]["W"]
+                if w_tensor.dim() == 4:
+                    samples_chunk["W"] = w_tensor[:, :, i : i + pred_batch_size, :]
+                else:
+                    samples_chunk["W"] = w_tensor[:, i : i + pred_batch_size, :]
+
+            pred_chunk = compute_predictions(
+                samples_chunk,
+                x_chunk,
+                model_type=model_type,
+                include_latent=include_latent,
+                y_type=y_type,
+            )
+
+            if return_means:
+                pred_chunk_processed = pred_chunk.mean(dim=0).detach().cpu().numpy()
+            else:
+                pred_chunk_processed = pred_chunk.detach().cpu().numpy()
+
+            pred_chunks.append(pred_chunk_processed)
+
+        if return_means:
+            pred_np = np.concatenate(pred_chunks, axis=0)
+            pred = pd.DataFrame(
+                pred_np,
+                index=data["site_index"],
+                columns=data["y_columns"],
+            )
+        else:
+            # If returning full samples, shape of each chunk is (num_samples, batch_size, n_species)
+            # We concatenate along the site dimension (axis 1)
+            pred = np.concatenate(pred_chunks, axis=1)
+
+    else:
+        pred = compute_predictions(
+            fit["samples"],
+            x,
+            model_type=model_type,
+            include_latent=include_latent,
+            y_type=y_type,
+        )
+        if return_means:
+            pred = pred.mean(dim=0)
+            pred = pred.detach().cpu().numpy()
+
+            pred = pd.DataFrame(
+                pred,
+                index=data["site_index"],
+                columns=data["y_columns"],
+            )
+
+        else:
+            pred = pred.detach().cpu().numpy()
+
+    return pred

pmf_dark-0.1.0/src/pmf_dark/inference.py

@@ -0,0 +1,107 @@
+import numpy as np
+import pyro
+from pyro.infer import SVI, Trace_ELBO, MCMC, NUTS
+from pyro.infer.autoguide import AutoNormal
+from pyro.optim import Adam
+from torch import device
+import torch
+from pyro.infer import Predictive
+
+
+def fit_svi(
+    model,
+    y,
+    x,
+    num_factors,
+    y_type,
+    lr=0.01,
+    num_iterations=2500,
+    cuda=False,
+    batch_size=None,
+    **kwargs,
+):
+
+    device = torch.device("cuda" if cuda and torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # Define create_plates helper for autoguide subsampling
+    def create_plates(*args, **kwargs):
+        # args[1] is the species presence/absence matrix Y
+        y_data = args[1]
+        n_sites = y_data.shape[0]
+        batch_size_val = kwargs.get("batch_size", None)
+        return pyro.plate("sites", n_sites, subsample_size=batch_size_val)
+
+    # Setup Inference
+    pyro.clear_param_store()
+    guide = AutoNormal(model, create_plates=create_plates)
+    optimizer = Adam({"lr": lr})
+    svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
+
+    # Training Loop
+    losses = []
+    for i in range(num_iterations):
+        loss = svi.step(x, y, num_factors, y_type, batch_size=batch_size, **kwargs)
+        losses.append(loss)
+        if i % 500 == 0:
+            print(f"Iteration {i} - Loss: {loss:.2f}")
+
+    # Simple convergence check
+    first = np.mean(losses[-200:-100])
+    last = np.mean(losses[-100:])
+    relative_change = abs(last - first) / abs(first)
+    if relative_change < 0.01:
+        print("SVI converged successfully.")
+    else:
+        print("SVI may not have converged.")
+
+    # return samples
+    predictive = Predictive(
+        model,
+        guide=guide,
+        num_samples=1000,
+    )
+
+    samples = predictive(x, y, num_factors, y_type, batch_size=None, **kwargs)
+
+    return {
+        "method": "svi",
+        "guide": guide,
+        "losses": losses,
+        "samples": samples,
+    }
+
+
+def fit_mcmc(
+    model,
+    y,
+    x,
+    num_factors,
+    y_type,
+    num_samples=1000,
+    warmup_steps=500,
+    num_chains=1,
+    batch_size=None,
+):
+    if batch_size is not None:
+        raise ValueError(
+            "MCMC does not support mini-batching. Please use SVI (method='svi') or set batch_size=None."
+        )
+    pyro.clear_param_store()
+
+    kernel = NUTS(model)
+
+    mcmc = MCMC(
+        kernel,
+        num_samples=num_samples,
+        warmup_steps=warmup_steps,
+        num_chains=num_chains,
+    )
+
+    mcmc.run(x, y, num_factors, y_type)
+
+    return {
+        "method": "mcmc",
+        "mcmc": mcmc,
+        "samples": mcmc.get_samples(),
+    }

pmf_dark-0.1.0/src/pmf_dark/models.py

@@ -0,0 +1,229 @@
+# 2. Define the Model
+import pyro
+import pyro.distributions as dist
+import torch
+
+
+def observation_dist(eta, y_type):
+    if y_type == "presence_absence":
+        return dist.Bernoulli(logits=eta)
+
+    elif y_type == "count":
+        return dist.Poisson(rate=torch.exp(torch.clamp(eta, -20, 20)))
+    # elif y_type == "count":
+    #
+    #    mu = torch.exp(torch.clamp(eta, -20, 20))
+    #
+    #    dispersion = pyro.sample(
+    #        "dispersion",
+    #        dist.LogNormal(
+    #            torch.zeros(eta.shape[-1], device=eta.device),
+    #            torch.ones(eta.shape[-1], device=eta.device),
+    #        ).to_event(1),
+    #    )
+    #
+    #    return dist.NegativeBinomial(
+    #        total_count=dispersion,
+    #        logits=torch.log(mu) - torch.log(dispersion),
+    #    )
+
+    else:
+        raise ValueError(
+            "y_type must be 'presence_absence', 'count', or 'zero_inflated_count'"
+        )
+
+
+def linear_model(X, Y, num_factors, y_type="presence_absence", batch_size=None):
+    device = X.device
+
+    n_sites, n_species = Y.shape
+    n_env = X.shape[1]
+
+    alpha = pyro.sample(
+        "alpha",
+        dist.Normal(
+            torch.zeros(n_species, device=device),
+            torch.ones(n_species, device=device),
+        ).to_event(1),
+    )
+
+    beta = pyro.sample(
+        "beta",
+        dist.Normal(
+            torch.zeros(n_env, n_species, device=device),
+            torch.ones(n_env, n_species, device=device),
+        ).to_event(2),
+    )
+
+    Z = pyro.sample(
+        "Z",
+        dist.Normal(
+            torch.zeros(n_species, num_factors, device=device),
+            torch.ones(n_species, num_factors, device=device),
+        ).to_event(2),
+    )
+
+    with pyro.plate("sites", n_sites, subsample_size=batch_size) as ind:
+        X_batch = X[ind]
+        Y_batch = Y[ind]
+
+        W_batch = pyro.sample(
+            "W",
+            dist.Normal(
+                torch.zeros(num_factors, device=device),
+                torch.ones(num_factors, device=device),
+            ).to_event(1),
+        )
+
+        eta_batch = (
+            alpha
+            + torch.matmul(X_batch, beta)
+            + torch.matmul(W_batch, Z.transpose(-1, -2))
+        )
+
+        pyro.sample(
+            "obs",
+            observation_dist(eta_batch, y_type).to_event(1),
+            obs=Y_batch,
+        )
+
+
+def gaussian(
+    X, Y, num_factors, y_type="presence_absence", batch_size=None
+):
+    device = X.device
+
+    n_sites, n_species = Y.shape
+    n_env = X.shape[1]
+
+    alpha = pyro.sample(
+        "alpha",
+        dist.Normal(
+            torch.zeros(n_species, device=device),
+            torch.ones(n_species, device=device),
+        ).to_event(1),
+    )
+
+    # species-specific environmental optimum
+    mu = pyro.sample(
+        "mu",
+        dist.Normal(
+            torch.zeros(n_env, n_species, device=device),
+            torch.ones(n_env, n_species, device=device),
+        ).to_event(2),
+    )
+
+    # positive niche strength / inverse width
+    gamma = pyro.sample(
+        "gamma",
+        dist.LogNormal(
+            torch.zeros(n_env, n_species, device=device),
+            torch.ones(n_env, n_species, device=device),
+        ).to_event(2),
+    )
+
+    Z = pyro.sample(
+        "Z",
+        dist.Normal(
+            torch.zeros(n_species, num_factors, device=device),
+            torch.ones(n_species, num_factors, device=device),
+        ).to_event(2),
+    )
+
+    with pyro.plate("sites", n_sites, subsample_size=batch_size) as ind:
+        X_batch = X[ind]
+        Y_batch = Y[ind]
+
+        W_batch = pyro.sample(
+            "W",
+            dist.Normal(
+                torch.zeros(num_factors, device=device),
+                torch.ones(num_factors, device=device),
+            ).to_event(1),
+        )
+
+        # quadratic environmental response
+        diff = X_batch[..., None] - mu[..., None, :, :]
+        env_effect = -torch.sum(gamma[..., None, :, :] * diff**2, dim=-2)
+
+        eta_batch = alpha + env_effect + torch.matmul(W_batch, Z.transpose(-1, -2))
+
+        pyro.sample(
+            "obs",
+            observation_dist(eta_batch, y_type).to_event(1),
+            obs=Y_batch,
+        )
+
+
+def bnn_model(
+    X, Y, num_factors=1, y_type="presence_absence", hidden_size=10, batch_size=None
+):
+    device = X.device
+
+    n_sites, n_species = Y.shape
+    n_env = X.shape[1]
+
+    w1 = pyro.sample(
+        "w1",
+        dist.Normal(
+            torch.zeros(n_env, hidden_size, device=device),
+            torch.ones(n_env, hidden_size, device=device),
+        ).to_event(2),
+    )
+
+    b1 = pyro.sample(
+        "b1",
+        dist.Normal(
+            torch.zeros(hidden_size, device=device),
+            torch.ones(hidden_size, device=device),
+        ).to_event(1),
+    )
+
+    w2 = pyro.sample(
+        "w2",
+        dist.Normal(
+            torch.zeros(hidden_size, n_species, device=device),
+            torch.ones(hidden_size, n_species, device=device),
+        ).to_event(2),
+    )
+
+    b2 = pyro.sample(
+        "b2",
+        dist.Normal(
+            torch.zeros(n_species, device=device),
+            torch.ones(n_species, device=device),
+        ).to_event(1),
+    )
+
+    Z = None
+    if num_factors > 0:
+        Z = pyro.sample(
+            "Z",
+            dist.Normal(
+                torch.zeros(n_species, num_factors, device=device),
+                torch.ones(n_species, num_factors, device=device),
+            ).to_event(2),
+        )
+
+    with pyro.plate("sites", n_sites, subsample_size=batch_size) as ind:
+        X_batch = X[ind]
+        Y_batch = Y[ind]
+
+        hidden = torch.tanh(X_batch @ w1 + b1)
+        eta_batch = hidden @ w2 + b2
+
+        if num_factors > 0:
+            W_batch = pyro.sample(
+                "W",
+                dist.Normal(
+                    torch.zeros(num_factors, device=device),
+                    torch.ones(num_factors, device=device),
+                ).to_event(1),
+            )
+            eta_batch = eta_batch + W_batch @ Z.transpose(-1, -2)
+
+        pyro.sample(
+            "obs",
+            observation_dist(eta_batch, y_type).to_event(1),
+            obs=Y_batch,
+        )