Metadata-Version: 2.4
Name: torchstats
Version: 0.1.0
Summary: Statistical functions, mostly probability distributions, in PyTorch.
Author-email: David Boetius <david.boetius@uni-konstanz.de>
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
License-File: LICENSE
Requires-Dist: torch>=1.12.1,<3.0
Requires-Dist: numpy>=1.25,<3.0
Requires-Dist: scipy>=1.11,<2.0
Requires-Dist: scikit-learn>=1.3.2,<2.0
Requires-Dist: rust_enum==1.1.5
Requires-Dist: torchstats[test, develop] ; extra == "all"
Requires-Dist: ruff==0.11.11 ; extra == "develop"
Requires-Dist: pytest>=7.4,<7.5 ; extra == "test"
Requires-Dist: nox==2025.5.1 ; extra == "test"
Provides-Extra: all
Provides-Extra: develop
Provides-Extra: test

# torchstats
This package implements various probability distributions in PyTorch, 
among those univariate distributions from scipy, categorical distributions, mixture models, and Bayesian networks.
It implements computing probabilities and densities, and rudimentary functionality for fitting distributions to data.

## Examples
Categorical distributions in a [one-hot](https://en.wikipedia.org/wiki/One-hot) event space.
```python
import torch
from torchstats import CategoricalOneHot

d = CategoricalOneHot([0.1, 0.5, 0.2, 0.2])
events = d.sample(5)
# tensor([[0., 1., 0., 0.],
#         [0., 0., 0., 1.],
#         [0., 1., 0., 0.],
#         [0., 0., 0., 1.],
#         [0., 1., 0., 0.]], dtype=torch.float64)
d.probability((events, events))
# tensor([0.5000, 0.2000, 0.5000, 0.2000, 0.5000], dtype=torch.float64)

# Probabilities and densities are always computed for hyperrectangles in event space.
# These hyperrectangles are represented as tuples of lower bound and upper bounds.
# This allows us to compute, for example, the probability of all categories except the first.
d.probability((torch.tensor([0.0, 0.0, 0.0, 0.0]), torch.tensor([0.0, 1.0, 1.0, 1.0])))
# tensor([0.9000], dtype=torch.float64
```

A mixture model of truncated Gaussian distributions.
```python
import torch
from torchstats import MixtureModel

data = torch.rand(100)
d = MixtureModel.fit_truncnorm_mixture(data, bounds=(0.0, 1.0), n_components=3, n_init=3)
d.weights
# tensor([0.4785, 0.1701, 0.3514], dtype=torch.float64)
d.density(data)
# ...
```

A Bayesian network
```python
import torch
from torchstats import BayesianNetwork, Categorical

factory = BayesianNetwork.Factory()
n1 = factory.new_node("var1")
n1.discrete_event_space(0, 1, 2)
n1.set_conditional_probability({}, Categorical([0.1, 0.2, 0.7]))

n2 = factory.new_node("var2")
n2.discrete_event_space(0, 1, 2)
n2.set_conditional_probability({}, Categorical([0.8, 0.1, 0.1]))

n3 = factory.new_node("var3")
n3.set_parents(n1, n2)
n3.discrete_event_space(0, 1)
n3.set_conditional_probability({n1: 0, n2: 0}, Categorical([1.0, 0.0])) 
n3.set_conditional_probability({n1: 0, n2: (1, 2)}, Categorical([0.1, 0.9]))
n3.set_conditional_probability({n1: (1, 2), n2: (0, 2)}, Categorical([0.0, 1.0]))

n1.hidden = n2.hidden = True  # make n1 and n2 hidden/latent variables

d = factory.create()
type(d)
# <class 'torchstats.bayesian_network.BayesianNetwork'>
d.output_space.input_shape
# torch.Size([1])
d.probability((torch.tensor(0), torch.tensor(0)))
# tensor([0.0820], dtype=torch.float64)
d.probability((torch.tensor(1), torch.tensor(1)))
# tensor([0.0820], dtype=torch.float64)
d.probability((torch.tensor(1), torch.tensor(1)))
# tensor([0.9180], dtype=torch.float64)
```

## Installation
Install `torchstats` from PyPI using
```bash
pip install torchstats
```

To extend `torchstats`, clone this repository and run
```bash
pip install -e .[all]
```
to install development and testing dependencies.
To run the tests, run either `pytest` or `nox`.

