Uniform superposition preparation moved to std (#2195)

Moved uniform superposition preparation into Std.StatePreparaion
(simplified version based on df chemistry)
Improved checks and updated doc comments
Added tests

---------

Co-authored-by: Dmitry Vasilevsky <[email protected]>

Author: DmitryVasilevsky

library/chemistry/src/JordanWigner/JordanWignerOptimizedBlockEncoding.qs

@@ -8,9 +8,6 @@ export
     MixedStatePreparation,
     MixedStatePreparationRequirements,
     PurifiedMixedState,
-    ObliviousAmplitudeAmplificationFromStatePreparation,
-    ObliviousAmplitudeAmplificationFromPartialReflections,
-    ApplyObliviousAmplitudeAmplification,
     PurifiedMixedStateRequirements,
     BlockEncodingByLCU,
     QuantumWalkByQubitization,
@@ -27,6 +24,7 @@ import Std.Math.*;
 import Std.Convert.IntAsDouble;
 import Std.Arithmetic.ApplyIfGreaterLE;
 import Std.StatePreparation.PreparePureStateD;
+import Std.StatePreparation.PrepareUniformSuperposition;
 import Std.Diagnostics.Fact;
 import Utils.GeneratorIndex;
 import Utils.GeneratorSystem;
@@ -628,207 +626,6 @@ operation PrepareQuantumROMState(
     }
 }
 
-/// # Summary
-/// Creates a uniform superposition over states that encode 0 through `nIndices - 1`.
-///
-/// # Description
-///
-/// This operation can be described by a unitary matrix $U$ that creates
-/// a uniform superposition over all number states
-/// $0$ to $M-1$, given an input state $\ket{0\cdots 0}$. In other words,
-/// $$
-/// \begin{align}
-///     U \ket{0} = \frac{1}{\sqrt{M}} \sum_{j=0}^{M - 1} \ket{j}.
-/// \end{align}
-/// $$.
-///
-/// # Input
-/// ## nIndices
-/// The desired number of states $M$ in the uniform superposition.
-/// ## indexRegister
-/// The qubit register that stores the number states in `LittleEndian` format.
-/// This register must be able to store the number $M-1$, and is assumed to be
-/// initialized in the state $\ket{0\cdots 0}$.
-///
-/// # Remarks
-/// The operation is adjointable, but requires that `indexRegister` is in a uniform
-/// superposition over the first `nIndices` basis states in that case.
-///
-/// # Example
-/// The following example prepares the state $\frac{1}{\sqrt{6}}\sum_{j=0}^{5}\ket{j}$
-/// on $3$ qubits.
-/// ``` Q#
-/// let nIndices = 6;
-/// using(indexRegister = Qubit[3]) {
-///     PrepareUniformSuperposition(nIndices, LittleEndian(indexRegister));
-///     // ...
-/// }
-/// ```
-operation PrepareUniformSuperposition(nIndices : Int, indexRegister : Qubit[]) : Unit is Adj + Ctl {
-    if nIndices == 0 {
-        fail "Cannot prepare uniform superposition over 0 basis states.";
-    } elif nIndices == 1 {
-        // Superposition over one state, so do nothing.
-    } elif nIndices == 2 {
-        H(indexRegister[0]);
-    } else {
-        let nQubits = BitSizeI(nIndices - 1);
-        if nQubits > Length(indexRegister) {
-            fail $"Cannot prepare uniform superposition over {nIndices} states as it is larger than the qubit register.";
-        }
-
-        use flagQubit = Qubit[3];
-        let targetQubits = indexRegister[0..nQubits - 1];
-        let qubits = flagQubit + targetQubits;
-        let stateOracle = PrepareUniformSuperpositionOracle(nIndices, nQubits, 0, _);
-
-        let phases = ([0.0, PI()], [PI(), 0.0]);
-
-        ObliviousAmplitudeAmplificationFromStatePreparation(
-            phases,
-            stateOracle,
-            (qs0, qs1) => I(qs0[0]),
-            0
-        )(qubits, []);
-
-
-        ApplyToEachCA(X, flagQubit);
-    }
-}
-
-function ObliviousAmplitudeAmplificationFromStatePreparation(
-    phases : (Double[], Double[]),
-    startStateOracle : Qubit[] => Unit is Adj + Ctl,
-    signalOracle : (Qubit[], Qubit[]) => Unit is Adj + Ctl,
-    idxFlagQubit : Int
-) : (Qubit[], Qubit[]) => Unit is Adj + Ctl {
-    let startStateReflection = (phase, qs) => {
-        within {
-            ApplyToEachCA(X, qs);
-        } apply {
-            Controlled R1(Rest(qs), (phase, qs[0]));
-        }
-    };
-    let targetStateReflection = (phase, qs) => R1(phase, qs[idxFlagQubit]);
-    let obliviousSignalOracle = (qs0, qs1) => { startStateOracle(qs0); signalOracle(qs0, qs1); };
-    return ObliviousAmplitudeAmplificationFromPartialReflections(
-        phases,
-        startStateReflection,
-        targetStateReflection,
-        obliviousSignalOracle
-    );
-}
-
-function ObliviousAmplitudeAmplificationFromPartialReflections(
-    phases : (Double[], Double[]),
-    startStateReflection : (Double, Qubit[]) => Unit is Adj + Ctl,
-    targetStateReflection : (Double, Qubit[]) => Unit is Adj + Ctl,
-    signalOracle : (Qubit[], Qubit[]) => Unit is Adj + Ctl
-) : (Qubit[], Qubit[]) => Unit is Adj + Ctl {
-    return ApplyObliviousAmplitudeAmplification(
-        phases,
-        startStateReflection,
-        targetStateReflection,
-        signalOracle,
-        _,
-        _
-    );
-}
-
-/// # Summary
-/// Oblivious amplitude amplification by specifying partial reflections.
-///
-/// # Input
-/// ## phases
-/// Phases of partial reflections
-/// ## startStateReflection
-/// Reflection operator about start state of auxiliary register
-/// ## targetStateReflection
-/// Reflection operator about target state of auxiliary register
-/// ## signalOracle
-/// Unitary oracle $O$ of type `ObliviousOracle` that acts jointly on the
-/// auxiliary and system registers.
-/// ## auxiliaryRegister
-/// Auxiliary register
-/// ## systemRegister
-/// System register
-///
-/// # Remarks
-/// Given a particular auxiliary start state $\ket{\text{start}}\_a$, a
-/// particular auxiliary target state $\ket{\text{target}}\_a$, and any
-/// system state $\ket{\psi}\_s$, suppose that
-/// \begin{align}
-/// O\ket{\text{start}}\_a\ket{\psi}\_s= \lambda\ket{\text{target}}\_a U \ket{\psi}\_s + \sqrt{1-|\lambda|^2}\ket{\text{target}^\perp}\_a\cdots
-/// \end{align}
-/// for some unitary $U$.
-/// By a sequence of reflections about the start and target states on the
-/// auxiliary register interleaved by applications of `signalOracle` and its
-/// adjoint, the success probability of applying $U$ may be altered.
-///
-/// In most cases, `auxiliaryRegister` is initialized in the state $\ket{\text{start}}\_a$.
-///
-/// # References
-/// - See [*D.W. Berry, A.M. Childs, R. Cleve, R. Kothari, R.D. Somma*](https://arxiv.org/abs/1312.1414)
-/// for the standard version.
-/// - See [*G.H. Low, I.L. Chuang*](https://arxiv.org/abs/1610.06546)
-/// for a generalization to partial reflections.
-operation ApplyObliviousAmplitudeAmplification(
-    phases : (Double[], Double[]),
-    startStateReflection : (Double, Qubit[]) => Unit is Adj + Ctl,
-    targetStateReflection : (Double, Qubit[]) => Unit is Adj + Ctl,
-    signalOracle : (Qubit[], Qubit[]) => Unit is Adj + Ctl,
-    auxiliaryRegister : Qubit[],
-    systemRegister : Qubit[]
-) : Unit is Adj + Ctl {
-    let (aboutStart, aboutTarget) = phases;
-    Fact(Length(aboutStart) == Length(aboutTarget), "number of phases about start and target state must be equal");
-    let numPhases = Length(aboutStart);
-
-    for idx in 0..numPhases-1 {
-        if aboutStart[idx] != 0.0 {
-            startStateReflection(aboutStart[idx], auxiliaryRegister);
-        }
-
-        if aboutTarget[idx] != 0.0 {
-            // In the last iteration we do not need to apply `Adjoint signalOracle`
-            if idx == numPhases - 1 {
-                signalOracle(auxiliaryRegister, systemRegister);
-                targetStateReflection(aboutTarget[idx], auxiliaryRegister);
-            } else {
-                within {
-                    signalOracle(auxiliaryRegister, systemRegister);
-                } apply {
-                    targetStateReflection(aboutTarget[idx], auxiliaryRegister);
-                }
-            }
-        }
-    }
-
-    // We do need one more `signalOracle` call, if the last phase about the target state was 0.0
-    if numPhases == 0 or aboutTarget[numPhases - 1] == 0.0 {
-        signalOracle(auxiliaryRegister, systemRegister);
-    }
-}
-
-operation PrepareUniformSuperpositionOracle(nIndices : Int, nQubits : Int, idxFlag : Int, qubits : Qubit[]) : Unit is Adj + Ctl {
-    let targetQubits = qubits[3...];
-    let flagQubit = qubits[0];
-    let auxillaryQubits = qubits[1..2];
-    let theta = ArcSin(Sqrt(IntAsDouble(2^nQubits) / IntAsDouble(nIndices)) * Sin(PI() / 6.0));
-
-    ApplyToEachCA(H, targetQubits);
-    use compareQubits = Qubit[nQubits] {
-        within {
-            ApplyXorInPlace(nIndices - 1, compareQubits);
-        } apply {
-            ApplyIfGreaterLE(X, targetQubits, compareQubits, auxillaryQubits[0]);
-            X(auxillaryQubits[0]);
-        }
-    }
-    Ry(2.0 * theta, auxillaryQubits[1]);
-    (Controlled X)(auxillaryQubits, flagQubit);
-}
-
 // Classical processing
 // This discretizes the coefficients such that
 // |coefficient[i] * oneNorm - discretizedCoefficient[i] * discretizedOneNorm| * nCoeffs <= 2^{1-bitsPrecision}.

library/chemistry/src/Tests.qs

@@ -2,7 +2,6 @@ import Std.StatePreparation.ApproximatelyPreparePureStateCP;
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import JordanWigner.JordanWignerOptimizedBlockEncoding.PrepareUniformSuperposition;
 import JordanWigner.JordanWignerClusterOperatorEvolutionSet.JordanWignerClusterOperatorPQRSTermSigns;
 import JordanWigner.OptimizedBEOperator.OptimizedBEXY;
 import JordanWigner.OptimizedBEOperator.SelectZ;
@@ -23,20 +22,6 @@ import Std.Math.Lg;
 import Std.Arrays.Reversed;
 import Std.Math.Sqrt;
 
-@Config(Unrestricted)
-@Test()
-operation PrepareUniformSuperpositionTest() : Unit {
-    let NBits = 4;
-    use qs = Qubit[NBits];
-    for NStates in 1..2^NBits-1 {
-        PrepareUniformSuperposition(NStates, qs);
-        let t = Std.Math.Sqrt(1.0 / IntAsDouble(NStates));
-        Adjoint Std.StatePreparation.PreparePureStateD(Repeated(t, NStates), Reversed(qs));
-        Fact(CheckAllZero(qs), "All Z");
-        ResetAll(qs);
-    }
-}
-
 @Config(Unrestricted)
 @Test()
 operation PrepareSparseMultiConfigurationalState0Test() : Unit {

library/src/tests/resources/src/state_preparation.qs

@@ -106,4 +106,22 @@ namespace Test {
         }
     }
 
+    operation TestPrepareUniformSuperposition(nStates : Int) : Unit {
+        use qs = Qubit[10];
+        PrepareUniformSuperposition(nStates, qs);
+        DumpMachine();
+        ResetAll(qs);
+    }
+
+    operation TestPrepareUniformSuperpositionExhaustive() : Unit {
+        let NBits = 4;
+        for NStates in 1..2^NBits-1 {
+            use qs = Qubit[NBits];
+            PrepareUniformSuperposition(NStates, qs);
+            let t = Std.Math.Sqrt(1.0 / IntAsDouble(NStates));
+            Adjoint Std.StatePreparation.PreparePureStateD(Repeated(t, NStates), Reversed(qs));
+            Fact(CheckAllZero(qs), "All qubits must be in zero state.");
+        }
+    }
+
 }

library/src/tests/state_preparation.rs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 use super::test_expression;
+use super::test_expression_fails;
 use super::test_expression_with_lib;
 use expect_test::expect;
 use qsc::interpret::Value;
@@ -254,3 +255,55 @@ fn check_preparation_doc_sample() {
     "#]]
     .assert_eq(&out);
 }
+
+#[test]
+fn check_uniform_superposition_preparation() {
+    let out = test_expression_with_lib(
+        "Test.TestPrepareUniformSuperposition(5)",
+        STATE_PREPARATION_TEST_LIB,
+        &Value::Tuple(vec![].into()),
+    );
+
+    expect![[r#"
+        STATE:
+        |0000000000⟩: 0.4472+0.0000𝑖
+        |0010000000⟩: 0.4472+0.0000𝑖
+        |0100000000⟩: 0.4472+0.0000𝑖
+        |1000000000⟩: 0.4472+0.0000𝑖
+        |1100000000⟩: 0.4472+0.0000𝑖
+    "#]]
+    .assert_eq(&out);
+}
+
+#[test]
+fn check_uniform_superposition_preparation_exhaustive() {
+    let _ = test_expression_with_lib(
+        "Test.TestPrepareUniformSuperpositionExhaustive()",
+        STATE_PREPARATION_TEST_LIB,
+        &Value::Tuple(vec![].into()),
+    );
+}
+
+#[test]
+fn check_uniform_superposition_short_array() {
+    let out = test_expression_fails(
+        "{
+            use qs=Qubit[2];
+            Std.StatePreparation.PrepareUniformSuperposition(5, qs);
+        }",
+    );
+
+    expect!["program failed: Qubit register is too short to prepare 5 states."].assert_eq(&out);
+}
+
+#[test]
+fn check_uniform_superposition_invalid_state_count() {
+    let out = test_expression_fails(
+        "{
+            use qs=Qubit[2];
+            Std.StatePreparation.PrepareUniformSuperposition(0, qs);
+        }",
+    );
+
+    expect!["program failed: Number of basis states must be positive."].assert_eq(&out);
+}