pypolymix logo

Docs	Code style	Coverage

Pypolymix

pypolymix is a light-weight PyTorch companion for building stochastic surrogate models.

Documentation

Project documentation lives at https://sandialabs.github.io/pypolymix.

Installation

From PyPI:

python -m pip install pypolymix

Optional extras:

python -m pip install "pypolymix[examples]"
python -m pip install "pypolymix[docs]"
python -m pip install "pypolymix[dev]"

For local development:

python -m pip install -e ".[dev,docs,examples]"

Quickstart

Fit a 1D polynomial chaos model with a mix of stochastic and deterministic coefficients:

import torch

from pypolymix.parameter_groups import DeterministicGroup, IIDGaussianGroup
from pypolymix.surrogate_models import PolynomialChaosExpansion
from pypolymix import StochasticModel

torch.manual_seed(0)
x = torch.linspace(-1, 1, 200).unsqueeze(-1)
y = torch.sin(3 * x) + 0.1 * torch.randn_like(x)

surrogate_model = PolynomialChaosExpansion(num_inputs=1, num_outputs=1, degree=5)
num_params = surrogate_model.num_params()

parameter_groups = [
    IIDGaussianGroup("stochastic", 2),
    DeterministicGroup("deterministic", num_params - 2),
]

model = StochasticModel(surrogate_model=surrogate_model, parameter_groups=parameter_groups)

optimizer = torch.optim.Adam(model.parameters(), lr=5e-3)
for _ in range(500):
    optimizer.zero_grad()
    preds = model(x, num_samples=16).mean(dim=0)
    data_loss = torch.mean((preds - y) ** 2)
    loss = data_loss + 1e-3 * model.distribution_loss()
    loss.backward()
    optimizer.step()

with torch.no_grad():
    samples = model(x, num_samples=100)
    mean = samples.mean(dim=0)
    std = samples.std(dim=0)

Polynomial chaos example

Development

If you use Poetry:

poetry install --with dev,docs,examples

Install the pre-commit hook:

pre-commit install

Examples

Notebooks live in docs/examples/ and are rendered on the documentation site.

For a full end-to-end workflow rather than a single notebook, see docs/examples/constitutive_model/: a runnable mixture-of-experts surrogate for a constitutive model, covering the data contract, transforms, training with checkpointing, and evaluation by integrating the learned right-hand side forward in time.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pypolymix-0.1.9.tar.gz (16.2 kB view details)

Uploaded Sep 9, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

pypolymix-0.1.9-py3-none-any.whl (21.7 kB view details)

Uploaded Sep 9, 2026 Python 3

File details

Details for the file pypolymix-0.1.9.tar.gz.

File metadata

Download URL: pypolymix-0.1.9.tar.gz
Upload date: Sep 9, 2026
Size: 16.2 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for pypolymix-0.1.9.tar.gz
Algorithm	Hash digest
SHA256	`3522543f671c88d9c73c0804d5baaf990a53f484656fc2bb063c41a31ce4629c`
MD5	`dd724b5fcc7e51149a8e400d6937b64e`
BLAKE2b-256	`8fd68902140ed6ddf802c16f8a3f4c70983f21d0964816fe989876961cae743c`

See more details on using hashes here.

Provenance

The following attestation bundles were made for pypolymix-0.1.9.tar.gz:

Publisher: publish-on-tag.yml on sandialabs/pypolymix

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: pypolymix-0.1.9.tar.gz
- Subject digest: 3522543f671c88d9c73c0804d5baaf990a53f484656fc2bb063c41a31ce4629c
- Sigstore transparency entry: 2765877577
- Sigstore integration time: Sep 9, 2026
Source repository:
- Permalink: sandialabs/pypolymix@74ade28dfc6243dfb6873e20df02ffac80c4e508
- Branch / Tag: refs/tags/v0.1.9
- Owner: https://github.com/sandialabs
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish-on-tag.yml@74ade28dfc6243dfb6873e20df02ffac80c4e508
- Trigger Event: push