Metadata-Version: 2.4
Name: tmg_hmc
Version: 0.0.2
Summary: This package implements exact HMC sampling for truncated multivariate gaussians with quadratic constraints
Project-URL: Homepage, https://github.com/erik-a-bensen/tmg_hmc
Project-URL: Repository, https://github.com/erik-a-bensen/tmg_hmc
Project-URL: Issues, https://github.com/erik-a-bensen/tmg_hmc/issues
Author-email: "Erik A. Bensen" <ebensen@andrew.cmu.edu>
License: MIT License
        
        Copyright (c) 2025 Erik A. Bensen
        
        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.
License-File: LICENSE
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.10
Requires-Dist: numpy
Requires-Dist: scipy
Requires-Dist: torch
Provides-Extra: dev
Requires-Dist: black; extra == 'dev'
Requires-Dist: flake8; extra == 'dev'
Requires-Dist: isort; extra == 'dev'
Requires-Dist: mypy; extra == 'dev'
Requires-Dist: pre-commit; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Provides-Extra: test
Requires-Dist: pytest-cov; extra == 'test'
Requires-Dist: pytest>=7.0; extra == 'test'
Description-Content-Type: text/markdown

# tmg_hmc

[![Python Versions](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)

> **Exact Hamiltonian Monte Carlo sampling for truncated multivariate Gaussians with quadratic constraints**

This package implements the exact HMC algorithm from [Pakman and Paninski (2014)](https://doi.org/10.1080/10618600.2013.788448) for sampling from truncated multivariate Gaussian distributions. 

## How It Works

The algorithm uses Hamiltonian Monte Carlo with
1. **Analytic Hamiltonian Dynamics**: Particles follow deterministic Hamiltonian trajectories that are analytically computable
2. **Exact Bounces**: When a trajectory hits a constraint boundary, the algorithm computes the exact bounce time by solving the quartic equation for the hit time analytically
4. **Perfect Acceptance Probability**: Unlike standard HMC, there's no integration error to solve the Hamiltonian dynamics. This means the acceptance probability is always 1.

See [Pakman & Paninski (2014)](https://doi.org/10.1080/10618600.2013.788448) for mathematical details.

## Features

- ✅ **Flexible constraints** - Supports linear and quadratic inequality constraints
- ✅ **Efficient** - Uses optimized compiled C++ hit time calculation for efficient sampling
- ✅ **GPU acceleration** - Optional PyTorch backend for large-scale problems
<!-- - ✅ **Well-tested** - Comprehensive test suite ensuring correctness -->

## Installation

### From Source
```bash
git clone https://github.com/erik-a-bensen/tmg_hmc.git
cd tmg_hmc 
pip install .
```

<!-- ### Development Installation
```bash
git clone https://github.com/erik-a-bensen/tmg_hmc.git
cd tmg_hmc/src/tmg_hmc/compiled # Need to precompile for editable install
make 
cd ../../..
pip install -e ".[dev]"  # Includes testing dependencies -->
```

**Requirements:**
- Python 3.10+
- numpy
- scipy
- torch

**Build Requirements:**
- C++ compiler (g++, clang, or MSVC)
- make

## Quick Start

### Linearly Constrained Gaussian
Sample a 2D standard normal with the y-component restricted to be positive:
```python
import numpy as np
from tmg_hmc import TMGSampler 

# Define the mean and covariance of the untruncated distribution
mu = np.zeros((2, 1))
Sigma = np.identity(2)
sampler = TMGSampler(mu, Sigma)

# Define the constraint y >= 0
# This corresponds to the constraint: f^T x + c >= 0
# where f = [0, 1] and c = 0
f = np.array([[0], [1]])
sampler.add_constraint(f=f, c=0)

# Sample 100 samples with 100 burn-in iterations
x0 = np.array([[1], [1]])  # Initial point (must satisfy constraints)
samples = sampler.sample(x0, n_samples=100, burn_in=100)

print(f"Sample mean: {samples.mean(axis=1)}")
print(f"Sample covariance:\n{np.cov(samples)}")
```

### Quadratically Constrained Gaussian
Sample from a Gaussian constrained to a circular region:
```python
import numpy as np
from tmg_hmc import TMGSampler

# 2D standard normal
mu = np.zeros((2, 1))
Sigma = np.identity(2)
sampler = TMGSampler(mu, Sigma)

# Add constraint: x^2 + y^2 <= 4 (inside circle of radius 2)
# Quadratic constraint: x^T A x + f^T x + c <= 0
# For x^2 + y^2 - 4 <= 0, we have A = I, f = 0, c = -4
A = np.identity(2)
c = -4
sampler.add_constraint(A=A, c=c)

# Sample
x0 = np.array([[0.5], [0.5]])
samples = sampler.sample(x0, n_samples=1000, burn_in=100)
```

### Multiple Constraints
Combine multiple constraints (e.g., box constraints):
```python
import numpy as np
from tmg_hmc import TMGSampler

mu = np.zeros((2, 1))
Sigma = np.identity(2)
sampler = TMGSampler(mu, Sigma)

# Box constraint: -1 <= x, y <= 1
# x >= -1  =>  [1,0]^T x + 1 >= 0
sampler.add_constraint(f=np.array([[1], [0]]), c=1)
# x <= 1   =>  [-1,0]^T x + 1 >= 0
sampler.add_constraint(f=np.array([[-1], [0]]), c=1)
# y >= -1  =>  [0,1]^T x + 1 >= 0
sampler.add_constraint(f=np.array([[0], [1]]), c=1)
# y <= 1   =>  [0,-1]^T x + 1 >= 0
sampler.add_constraint(f=np.array([[0], [-1]]), c=1)

x0 = np.array([[0], [0]])
samples = sampler.sample(x0, n_samples=1000, burn_in=100)
```

## Examples
TO BE IMPLEMENTED
<!-- See the `examples/` directory for:
- Linear constraint examples
- Quadratic constraint examples  
- High-dimensional sampling
- Comparison with approximate methods -->

## Testing

TO BE IMPLEMENTED
<!-- Run the test suite:
```bash
pytest tests/
``` -->

## Documentation

- [Full API Reference](API_DOCS.md) - Complete documentation of all functions and classes
- [Hit-time Calculations](resources/HMC_exact_soln.pdf) - Mathematica solutions for the hit times of each type of constraint.

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.

<!-- ## Related Projects
- [tmg](https://github.com/brunzema/truncated-mvn-sampler) - R package for approximate truncated Gaussian sampling -->

## License

This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

## Citation

If you use this package in your research, please cite:

**Software:**
> Bensen, E. A. (2025). tmg_hmc: A Python package for Exact HMC Sampling for Truncated Multivariate Gaussians with Linear and Quadratic Constraints. *TBD*. [To be updated upon acceptance]

**Methodology:**
> Pakman, A., & Paninski, L. (2014). Exact Hamiltonian Monte Carlo for Truncated Multivariate Gaussians. *Journal of Computational and Graphical Statistics*, 23(2), 518-542. https://doi.org/10.1080/10618600.2013.788448

<details>
<summary>BibTeX</summary>

```bibtex
@article{Bensen2025tmghmc,
  title={tmg\_hmc: A Python package for Exact HMC Sampling for Truncated Multivariate Gaussians with Linear and Quadratic Constraints},
  author={Bensen, Erik A.},
  journal={TBD},
  year={2025},
  note={[To be updated upon acceptance]}
}

@article{PakmanPaninski2014,
  title={Exact Hamiltonian Monte Carlo for Truncated Multivariate Gaussians},
  author={Pakman, Ari and Paninski, Liam},
  journal={Journal of Computational and Graphical Statistics},
  volume={23},
  number={2},
  pages={518--542},
  year={2014},
  publisher={Taylor \& Francis},
  doi={10.1080/10618600.2013.788448}
}
```
</details>

## Acknowledgments

This implementation is based on the exact HMC algorithm developed by Ari Pakman and Liam Paninski.