merlin.core.probability_distribution

ProbabilityDistribution

class merlin.core.probability_distribution.ProbabilityDistribution(tensor, n_modes, n_photons, computation_space=ComputationSpace.FOCK, logical_performance=None, _custom_basis=None)

Bases: object

Probability tensor bundled with Fock metadata and post-filter tracking.

Parameters

tensor:: Dense or sparse probabilities; leading dimensions are treated as batch axes.
n_modes:: Number of modes in the Fock space.
n_photons:: Total photon number represented by the distribution.
computation_space:: Basis enumeration used to order amplitudes (fock, unbunched, dual_rail).
logical_performance:: Optional per-batch scalar tracking kept/total probability after filtering.

Notes

Instances are normalized on construction; arithmetic-style temporary unnormalized states are not supported (unlike StateVector). Only shape, device, dtype, and requires_grad are delegated to the underlying torch.Tensor; tensor-like helpers to, clone, detach, and requires_grad_ mirror tensor semantics while keeping metadata and logical performance aligned. Layout-changing tensor operations should be done on tensor directly, then wrapped again via from_tensor to maintain a consistent basis.

tensor: Tensor

n_modes: int

n_photons: int

computation_space: ComputationSpace = 'fock'

logical_performance: Optional[Tensor] = None

property basis: Combinadics | FilteredBasis | tuple[tuple[int, ...], ...]

property basis_size: int

to(*args, **kwargs)

Return a new ProbabilityDistribution with tensor (and logical_performance) moved/cast via torch.Tensor.to.

Return type:: ProbabilityDistribution

clone()

Return a cloned ProbabilityDistribution with metadata and logical performance copied.

Return type:: ProbabilityDistribution

detach()

Return a detached ProbabilityDistribution sharing data without gradients.

Return type:: ProbabilityDistribution

requires_grad_(requires_grad=True)

Set requires_grad on underlying tensors and return self.

Return type:: ProbabilityDistribution

property is_sparse: bool

property is_normalized: bool

memory_bytes()

Return the tensor’s approximate memory footprint in bytes.

Return type:: int

normalize()

In-place normalization; safe for zero-mass batches.

Return type:: ProbabilityDistribution

Returns

ProbabilityDistribution: The same instance, normalized along the basis dimension.

to_dense()

Return a dense, normalized tensor representation.

Return type:: Tensor

probabilities()

Alias for to_dense() for readability.

Return type:: Tensor

classmethod from_tensor(tensor, *, n_modes, n_photons, computation_space=None, dtype=None, device=None)

Build a distribution from an explicit probability tensor.

Return type:: ProbabilityDistribution

Parameters

tensor:: Dense or sparse probability tensor; last dimension must match the basis size.
n_modes / n_photons:: Metadata for basis construction.
computation_space:: Optional basis scheme; defaults to fock.
dtype / device:: Optional overrides for output tensor placement and precision.

Raises

ValueError: If the last dimension does not match the expected basis size.

classmethod from_state_vector(state_vector, *, dtype=None, device=None, computation_space=None)

Convert a StateVector to a probability distribution.

Return type:: ProbabilityDistribution

Parameters

state_vector:: Source amplitudes; must expose to_dense, n_modes, and n_photons.
dtype / device:: Optional overrides for output tensor placement and precision.
computation_space:: Optional basis scheme; defaults to fock.

classmethod from_perceval(distribution, *, dtype=None, device=None, sparse=None)

Construct from a Perceval BSDistribution.

Validates that all entries share the same photon number and mode count.

Return type:: ProbabilityDistribution

Parameters

distribution:: Input Perceval distribution.
dtype / device:: Optional overrides for output tensor placement and precision.
sparse:: Force dense or sparse output; default auto-selects based on fill ratio.

Raises

ValueError: If the distribution is empty or inconsistent in shape/photon number.

to_perceval(): Convert to Perceval BSDistribution (single) or list for batches.

filter(rule)

Apply post-selection filter and renormalize probabilities.

logical_performance records kept_mass / original_mass per batch.

Return type:: ProbabilityDistribution

Parameters

rule:: Computation space alias (fock, unbunched, dual_rail), a predicate, an explicit iterable of allowed states, or a tuple (space, predicate) to combine a computation-space constraint with an additional predicate.

Returns

ProbabilityDistribution: A new, normalized distribution; may shrink its basis when filtering to unbunched or dual_rail in the dense case.

Raises

ValueError: If dual_rail is selected with incompatible n_modes/n_photons or an unknown computation space is requested.

Notes and Examples

Constructors

from_tensor — wrap a probability tensor with Fock metadata. The last dimension must match the basis size for the chosen computation space. The distribution is normalized on construction:

import torch
from merlin.core.probability_distribution import ProbabilityDistribution

probs = torch.tensor([0.2, 0.3, 0.5])
pd = ProbabilityDistribution.from_tensor(probs, n_modes=2, n_photons=2)

The optional computation_space parameter selects the basis ordering (defaults to FOCK):

from merlin.core.computation_space import ComputationSpace

pd = ProbabilityDistribution.from_tensor(
    probs, n_modes=2, n_photons=2,
    computation_space=ComputationSpace.FOCK,
)

Batched inputs are supported — leading dimensions are treated as batch axes:

batch = torch.rand(16, 3)    # 16 samples, basis_size = 3
pd = ProbabilityDistribution.from_tensor(batch, n_modes=2, n_photons=2)
assert pd.shape == (16, 3)

from_state_vector — compute \(|a_i|^2\) from a StateVector:

from merlin.core.state_vector import StateVector

sv = StateVector.from_basic_state([1, 0, 1, 0], sparse=False)
pd = ProbabilityDistribution.from_state_vector(sv)
assert pd.n_modes == 4 and pd.n_photons == 2

from_perceval — convert from a Perceval BSDistribution:

import perceval as pcvl

dist = pcvl.BSDistribution()
dist[pcvl.BasicState([1, 0])] = 0.8
dist[pcvl.BasicState([0, 1])] = 0.2

pd = ProbabilityDistribution.from_perceval(dist)

Properties and metadata

n_modes and n_photons are set at construction and immutable. shape, device, dtype, and requires_grad are delegated to the underlying tensor:

pd.n_modes        # 2
pd.n_photons      # 2
pd.shape          # torch.Size([3])
pd.device         # device(type='cpu')

basis returns the Fock ordering for the current computation space (or a filtered basis after filter()). basis_size is len(basis):

pd.basis_size            # 3 for (2 modes, 2 photons, FOCK)
list(pd.basis)           # [(2, 0), (1, 1), (0, 2)]

computation_space records which basis scheme was used:

pd.computation_space     # ComputationSpace.FOCK

is_normalized is always True — distributions are normalized on construction and after every filter call.

Accessing probabilities

probabilities() and to_dense() both return a dense, normalized tensor:

dense = pd.probabilities()   # shape: (3,)
dense = pd.to_dense()        # equivalent

Use bracket syntax for a single Fock state’s probability:

import perceval as pcvl

p = pd[[1, 1]]                        # scalar tensor
p = pd[pcvl.BasicState([1, 1])]       # equivalent

For batched distributions, the returned tensor matches the batch shape:

batch_pd = ProbabilityDistribution.from_tensor(
    torch.rand(8, 3), n_modes=2, n_photons=2,
)
p = batch_pd[[1, 1]]   # shape: (8,)

Filtering and post-selection

filter() applies post-selection and returns a new, renormalized distribution. The logical_performance attribute records the fraction of probability mass kept per batch element.

Filter by computation space:

from merlin.core.computation_space import ComputationSpace

pd = ProbabilityDistribution.from_tensor(
    torch.tensor([0.5, 0.25, 0.25]), n_modes=2, n_photons=2,
)
filtered = pd.filter(ComputationSpace.UNBUNCHED)
filtered.basis_size               # 1 — only |1,1⟩ survives
filtered.logical_performance      # tensor(0.25) — 25% of mass kept

String aliases also work: "fock", "unbunched", "dual_rail".

Filter by predicate:

# Keep only states where mode 0 has at least 1 photon
filtered = pd.filter(lambda state: state[0] >= 1)

Filter by explicit allowed states:

filtered = pd.filter([(1, 1), (2, 0)])

Combined space + predicate — pass a tuple (space, predicate):

# Unbunched states where mode 0 is occupied
filtered = pd.filter((ComputationSpace.UNBUNCHED, lambda s: s[0] == 1))

logical_performance is None on unfiltered distributions and is set by filter() to a tensor of kept / total mass per batch element.

PyTorch-like helpers

to, clone, detach, and requires_grad_ mirror the standard torch.Tensor API while preserving metadata and logical_performance:

pd_cuda = pd.to("cuda")               # moves tensor + logical_performance
pd_copy = pd.clone()                   # independent copy
pd_det  = pd.detach()                  # shares data, no gradient graph
pd.requires_grad_(True)               # enable gradients in-place

QuantumLayer integration

With return_object=True and a probability measurement strategy, the layer returns a ProbabilityDistribution instead of a bare tensor:

import merlin as ML
from merlin.core.computation_space import ComputationSpace

layer = ML.QuantumLayer(
    builder=builder,
    input_state=[1, 0, 1, 0],
    n_photons=2,
    measurement_strategy=ML.MeasurementStrategy.probs(ML.ComputationSpace.FOCK),
    return_object=True,
)

pd = layer(x)                                       # ProbabilityDistribution
pd.probabilities()                                   # dense tensor
pd_ub = pd.filter(ComputationSpace.UNBUNCHED)        # post-select
pd_ub.logical_performance                             # fraction of mass kept

Perceval interoperability

import perceval as pcvl
from merlin.core.probability_distribution import ProbabilityDistribution

# Perceval → Merlin
dist = pcvl.BSDistribution()
dist[pcvl.BasicState([1, 0])] = 0.8
dist[pcvl.BasicState([0, 1])] = 0.2
pd = ProbabilityDistribution.from_perceval(dist)

# Merlin → Perceval
pcvl_back = pd.to_perceval()
assert pcvl_back[pcvl.BasicState([1, 0])] == 0.8

For batched distributions, to_perceval() returns a list of pcvl.BSDistribution objects.