Extending Samplers

This guide explains how to add a new configuration sampler to qvartools.

Architecture

All samplers extend the Sampler ABC defined in qvartools/samplers/sampler.py:

Sampler (ABC)
├── NFSampler              (classical/)
├── TransformerNFSampler   (classical/)
├── TrotterSampler         (quantum/)
├── CudaQSampler           (quantum/)
└── LUCJSampler            (quantum/)

Step 1: Implement the Sampler

# samplers/classical/my_sampler.py

from __future__ import annotations

import torch
from qvartools.samplers.sampler import Sampler, SamplerResult


class MySampler(Sampler):
    """My custom configuration sampler.

    Parameters
    ----------
    num_sites : int
        Number of lattice / orbital sites.
    temperature : float
        Sampling temperature controlling exploration.

    Examples
    --------
    >>> sampler = MySampler(num_sites=4, temperature=1.0)
    >>> result = sampler.sample(n_samples=1000)
    >>> print(f"Sampled {result.configs.shape[0]} configurations")
    """

    def __init__(
        self,
        num_sites: int,
        temperature: float = 1.0,
    ) -> None:
        self._num_sites = num_sites
        self._temperature = temperature

    def sample(self, n_samples: int) -> SamplerResult:
        """Sample configurations.

        Parameters
        ----------
        n_samples : int
            Number of configurations to sample.

        Returns
        -------
        SamplerResult
            Container with sampled configurations, counts, and metadata.
        """
        # Your sampling algorithm here
        configs = self._generate_configs(n_samples)

        # Count unique configurations
        unique, inverse, counts = torch.unique(
            configs, dim=0, return_inverse=True, return_counts=True
        )

        return SamplerResult(
            configs=unique,
            counts={
                "".join(str(int(b)) for b in cfg): int(c)
                for cfg, c in zip(unique, counts)
            },
            metadata={
                "temperature": self._temperature,
                "total_sampled": n_samples,
                "unique_count": len(unique),
            },
        )

    def _generate_configs(self, n_samples: int) -> torch.Tensor:
        ...

Step 2: Register the Sampler

Export from samplers/classical/__init__.py and samplers/__init__.py.

Step 3: Add Tests

# tests/test_samplers/test_my_sampler.py

from qvartools.samplers.classical.my_sampler import MySampler

def test_my_sampler_shape():
    sampler = MySampler(num_sites=4)
    result = sampler.sample(100)
    assert result.configs.ndim == 2
    assert result.configs.shape[1] == 4

def test_my_sampler_valid_values():
    sampler = MySampler(num_sites=8)
    result = sampler.sample(500)
    assert (result.configs >= 0).all()
    assert (result.configs <= 1).all()

Integration with Pipelines

To use your sampler in a pipeline, you can either:

1. Use it as a standalone component by calling sampler.sample() and passing the configurations to a solver.

2. Integrate with the pipeline by modifying FlowGuidedKrylovPipeline to accept a custom sampler in the training or basis extraction stages.

Design Guidelines

• Return SamplerResult with configs as a 2D integer tensor

• Populate counts for reproducibility and analysis

• Use metadata for sampler-specific diagnostics

• Support a device parameter for GPU sampling

• Validate particle-number constraints when applicable