magic-state-cultivation/src/latte/dem_util_test.py at main · Strilanc/magic-state-cultivation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
from typing import Any

import stim

from latte.dem_util import dem_with_replaced_targets, \
    dem_with_compressed_detectors, bernoulli_combo, Symptom


def test_dem_with_replaced_targets():
    dem = stim.DetectorErrorModel("""
        error(0.25) D2 D3 L0
        logical_observable L0
        detector(1, 2) D2
    """)
    dem2 = dem_with_replaced_targets(
        dem,
        {stim.target_logical_observable_id(0): stim.target_relative_detector_id(5)},
    )
    assert dem2 == stim.DetectorErrorModel("""
        error(0.25) D2 D3 D5
        detector D5
        detector(1, 2) D2
    """)


def approx_equal(v1: Any, v2: Any, *, atol: float):
    if isinstance(v1, dict) and isinstance(v2, dict):
        if v1.keys() != v2.keys():
            return False
        for k, s1 in v1.items():
            s2 = v2[k]
            if not approx_equal(s1, s2, atol=atol):
                return False
        return True
    elif (isinstance(v1, tuple) and isinstance(v2, tuple)) or (isinstance(v1, list) and isinstance(v2, list)):
        if len(v1) != len(v2):
            return False
        for s1, s2 in zip(v1, v2):
            if not approx_equal(s1, s2, atol=atol):
                return False
        return True
    elif isinstance(v1, float) and isinstance(v2, float):
        return abs(v2 - v1) <= atol
    else:
        return v1 == v2


def test_bernoulli_combo():
    # assert bernoulli_combo(
    #     errors={
    #         Symptom(obs_mask=1 << 0, dets=frozenset([0])): 0.01,
    #         Symptom(obs_mask=1 << 1, dets=frozenset([0, 1])): 0.01,
    #         Symptom(obs_mask=1 << 2, dets=frozenset([1, 2])): 0.01,
    #         Symptom(obs_mask=1 << 3, dets=frozenset([2, 3])): 0.01,
    #         Symptom(obs_mask=1 << 4, dets=frozenset([3])): 0.01,
    #     },
    #     compressed_dets=frozenset([]),
    #     max_errors=20,
    #     error_size_cutoff=3,
    #     detection_event_cutoff=2,
    # ) == {}
    #
    # assert bernoulli_combo(
    #     errors={
    #         Symptom(obs_mask=1 << 0, dets=frozenset([0])): 0.01,
    #         Symptom(obs_mask=1 << 1, dets=frozenset([0, 1])): 0.01,
    #         Symptom(obs_mask=1 << 2, dets=frozenset([1, 2])): 0.01,
    #         Symptom(obs_mask=1 << 3, dets=frozenset([2, 3])): 0.01,
    #         Symptom(obs_mask=1 << 4, dets=frozenset([3])): 0.01,
    #     },
    #     compressed_dets=frozenset([0]),
    #     max_errors=1,
    #     error_size_cutoff=3,
    #     detection_event_cutoff=2,
    # ) == {}

    v2 = bernoulli_combo(
        errors={
            Symptom(obs_mask=1 << 0, dets=frozenset([0])): 0.01,
            Symptom(obs_mask=1 << 1, dets=frozenset([0, 1])): 0.01,
            Symptom(obs_mask=1 << 2, dets=frozenset([1, 2])): 0.01,
            Symptom(obs_mask=1 << 3, dets=frozenset([2, 3])): 0.01,
            Symptom(obs_mask=1 << 4, dets=frozenset([3])): 0.01,
        },
        compressed_dets=frozenset([0]),
        max_errors=2,
        error_size_cutoff=3,
        detection_event_cutoff=2,
    )
    assert approx_equal(v2, {1: (3, 1e-4)}, atol=1e-5)

    v3 = bernoulli_combo(
        errors={
            Symptom(obs_mask=1 << 0, dets=frozenset([0])): 0.01,
            Symptom(obs_mask=1 << 1, dets=frozenset([0, 1])): 0.01,
            Symptom(obs_mask=1 << 2, dets=frozenset([1, 2])): 0.01,
            Symptom(obs_mask=1 << 3, dets=frozenset([2, 3])): 0.01,
            Symptom(obs_mask=1 << 4, dets=frozenset([3])): 0.01,
        },
        compressed_dets=frozenset([0, 1]),
        max_errors=3,
        error_size_cutoff=3,
        detection_event_cutoff=2,
    )
    assert approx_equal(v3, {2: (7, 1e-6)}, atol=5e-7)


def test_dem_with_compressed_detectors():
    dem = stim.DetectorErrorModel("""
        error(0.01) D0 L0
        error(0.01) D0 D1
        error(0.01) D1 D2
        error(0.01) D2 D3
        error(0.01) D3 D4
        error(0.01) D4
        detector(0) D0
        detector(1) D1
        detector(2) D2
        detector(3) D3
        detector(4) D4
    """)

    dem2 = dem_with_compressed_detectors(
        dem,
        lambda coords: len(coords) > 0 and coords[0] < 1,
        error_size_cutoff=3,
        detection_event_cutoff=4,
        max_compressed_errors=2,
    )
    assert dem2.approx_equals(stim.DetectorErrorModel("""
        error(0.01) D1 D2
        error(0.01) D2 D3
        error(0.01) D3 D4
        error(0.01) D4
        detector(0) D0
        detector(1) D1
        detector(2) D2
        detector(3) D3
        detector(4) D4
        error(0.0001) D1 L0
        error(0.5) D0
    """), atol=0.00001)