dmx-learn

Distributed Mixture Learning - A package for distributed heterogeneous density estimation. With only a few lines of code you can specify and fit complex models on variable-length heterogeneous data.

---

📚 Documentation

View the full documentation on Read the Docs:

👉 https://dmx-learn.readthedocs.io/en/latest/

---

🚀 Installation

Using Poetry (Recommended)

dmx-learn uses Poetry for dependency management. Install Poetry first if you haven't already:

curl -sSL https://install.python-poetry.org | python3 -

Then install dmx-learn:

# Clone or navigate to the package directory
cd /path/to/dmx-learn

# Install core dependencies only
poetry install

# Or install with optional features using extras
poetry install -E torch            # Include PyTorch support
poetry install -E optional         # Include MPI and UMAP support
poetry install -E all              # Include all optional features

# Or using dependency groups (for dev/docs)
poetry install --with dev          # Include dev tools (pytest)
poetry install --with docs         # Include documentation tools

Using pip

You can also install with pip using extras:

# Basic installation
pip install /path/to/dmx-learn

# With PyTorch support
pip install /path/to/dmx-learn[torch]

# With MPI and UMAP support
pip install /path/to/dmx-learn[optional]

# With all optional features
pip install /path/to/dmx-learn[all]

---

📦 Optional Dependencies

dmx-learn supports several optional dependency extras and groups:

Extras (for pip and Poetry)

Extra	Install Command	Includes	Use Case
torch	pip install .[torch] or poetry install -E torch	PyTorch	GPU-accelerated distributions
optional	pip install .[optional] or poetry install -E optional	mpi4py, umap-learn	Distributed computing
all	pip install .[all] or poetry install -E all	All above	Everything

Groups (Poetry only - for development)

Group	Install Command	Includes	Use Case
dev	poetry install --with dev	pytest, pytest-dependency	Testing
docs	poetry install --with docs	sphinx, sphinx-rtd-theme	Documentation building

Install multiple extras and groups:

# Poetry: extras with -E, groups with --with
poetry install -E all --with dev,docs

# pip: multiple extras with comma
pip install .[torch,optional]

Check what's installed:

poetry show

---

📖 Usage Examples

Stats Examples (scipy-based)

Examples using stats distributions (always available) are located in ./examples/stats_examples/:

# Using Poetry
poetry run python ./examples/stats_examples/mixture_example.py

# Or set PYTHONPATH
export PYTHONPATH=/path/to/dmx-learn:$PYTHONPATH
python ./examples/stats_examples/mixture_example.py

Bayesian Stats Examples

Examples using bstats are located in ./examples/bstats_examples/:

poetry run python ./examples/bstats_examples/dpm_auto_example1.py

Utility Examples

Examples for automatic model selection, detailed estimation, embeddings, and random graphs are located in ./examples/utils_examples/:

poetry run python ./examples/utils_examples/auto_example.py

Torch Examples (PyTorch-based)

Examples using torch_stats (requires PyTorch) are located in ./examples/torch_examples/:

poetry run python ./examples/torch_examples/stats_examples/mixture_example.py

---

🌐 Running with Spark

Examples that run with Apache Spark are located in ./examples/spark_examples/.

First, build a wheel:

cd /path/to/dmx-learn
pip install setuptools wheel
python setup.py bdist_wheel

Then run with Spark:

/path/to/spark/bin/spark-submit \
  --master local[*] \
  --py-files /path/to/dmx-learn/dist/dmx-learn-1.0.0-py3-none-any.whl \
  ./examples/spark_examples/mixture_example.py

---

🚄 Running with MPI4PY

Examples using mpi4py are located in ./examples/mpi4py_examples/.

Install MPI support:

# Using Poetry
poetry install -E optional

# Using pip
pip install /path/to/dmx-learn[optional]

Run with MPI:

# Run with 4 processes
mpiexec -n 4 poetry run python ./examples/mpi4py_examples/estimation_example.py

---

🔥 PyTorch Support (Optional)

dmx-learn includes torch_stats, a PyTorch-based implementation of statistical distributions for GPU-accelerated computation.

Installing with PyTorch

# Using Poetry (recommended)
poetry install -E torch

# Using pip
pip install /path/to/dmx-learn[torch]

Using torch_stats

from dmx.torch_stats import GaussianDistribution, MixtureDistribution

# Create distributions (runs on GPU if available)
dist = GaussianDistribution(mu=0.0, sigma2=1.0)

# Automatic device detection
from dmx.torch_utils import detect_device
device = detect_device()  # Returns 'cuda', 'mps', or 'cpu'

Running Torch Examples

Examples using torch_stats are located in ./examples/torch_examples/:

poetry run python ./examples/torch_examples/stats_examples/mixture_example.py

Testing with PyTorch

Run torch_stats tests on different devices:

# Default (CPU)
poetry run pytest tests/torch_stats

# On Apple Silicon MPS
export NUMBA_DISABLE_JIT=1
TEST_TORCH_DEVICE=mps poetry run pytest tests/torch_stats

# On CUDA
export NUMBA_DISABLE_JIT=1
TEST_TORCH_DEVICE=cuda poetry run pytest tests/torch_stats

See tests/torch_stats/README.md for detailed testing information.

---

🧪 Testing

Run All Tests

# Using Poetry
poetry run pytest tests/

# Using the test runner script
./test_runner.sh all

Run Specific Test Suites

# Stats tests (always available)
poetry run pytest tests/stats

# Torch tests (requires PyTorch)
export NUMBA_DISABLE_JIT=1
poetry run pytest tests/torch_stats

# Quick sanity check
./test_runner.sh quick

Test Runner Script

The test_runner.sh helper script makes testing easier:

./test_runner.sh stats        # Run stats tests
./test_runner.sh torch        # Run torch tests (CPU)
./test_runner.sh torch-mps    # Run torch tests (MPS)
./test_runner.sh all          # Run all tests
./test_runner.sh quick        # Quick sanity check

---

🔧 Development Setup

For development with all features:

# Clone the repository
git clone https://github.com/llnl/dmx-learn.git
cd dmx-learn

# Install Poetry
curl -sSL https://install.python-poetry.org | python3 -

# Install all dependencies including dev tools
poetry install -E all --with dev,docs

# Run tests
poetry run pytest tests/

# Build documentation
cd docs
poetry run make html

---

📋 Requirements

• Python: >= 3.10
• Core Dependencies: numpy, scipy, numba, pandas, pyspark, mpmath
• Optional: torch (for GPU acceleration), mpi4py (for distributed computing)

---

🆘 Troubleshooting

PyTorch Not Found

If you see errors about PyTorch not being installed:

# Install torch extra with Poetry
poetry install -E torch

# Or with pip
pip install /path/to/dmx-learn[torch]

Tests Failing

Make sure you have the test dependencies:

poetry install --with dev

For torch tests, set the environment variable:

export NUMBA_DISABLE_JIT=1

Import Errors

Make sure the package is installed:

poetry install

Or add to PYTHONPATH:

export PYTHONPATH=/path/to/dmx-learn:$PYTHONPATH

License

dmx-learn is distributed under the terms of the BSD 3-Clause license. See LICENSE and NOTICE for details.

SPDX-License-Identifier: BSD-3-Clause

LLNL-CODE-844837