PML-261: Add post-selection to MerLin utils.

[anthonyrtw]

Summary

Allow Perceval-style post-selection rules to a batch of probability distributions indexed by a number of keys. May be immediately applied to the output of a SLOSComputeGraph. Additionally a unit test has been added for this feature under tests/utils/test_postselection.py.

Type of change

• New feature

Proposed changes

Add new function in utils/post_selection.py called post_select_probs. It takes a batch of probability vectors and a post-selection rules.

import torch
from merlin.utils import post_select_probs

keys = [(2, 0), (1, 1), (0, 2)]
probs = torch.tensor([
    [0.5, 0.0, 0.5],
    [0, 0.25, 0.75]
])

post_selection_rule = "[0] == 2"

new_keys, new_probs = post_select_probs(post_selection_rule, keys, probs)
print(new_keys)
print(new_probs)

[(2, 0), (1, 1), (0, 2)]
tensor([[1., 0., 0.],
        [0., 0., 0.]])

Additionally, one can decide whether they want the output keys to be written in the post-selected basis or the original basis using the same_keys argument. Default: True

...
post_selection_rule = "[0] == 2 OR [1] == 2"

new_keys, new_probs = post_select_probs(post_selection_rule, keys, probs, same_keys=False)
print(new_keys)
print(new_probs)

[(2, 0), (0, 2)]
tensor([[0.5000, 0.5000],
        [0.2500, 0.7500]])

The code works by masking the each probability value based on whether its key obeys the post-selection rule. The probability vectors are then normalized.

How to test

cd merlin
pytest tests/utils/test_post_selection.py

Performance considerations

PostSelect is handled by exqalibur (written in C++). It stands to reason that if exqalibur is updated to allow evaluation over a batch of keys (list of tuples/lists), the performance would be increased.