Merge branch 'main' into 2024-12/deps

dev_tools/qualtran_dev_tools/notebook_specs.py

@@ -66,6 +66,7 @@
 import qualtran.bloqs.chemistry.hubbard_model.qubitization
 import qualtran.bloqs.chemistry.pbc.first_quantization.prepare_t
 import qualtran.bloqs.chemistry.pbc.first_quantization.prepare_uv
+import qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta
 import qualtran.bloqs.chemistry.pbc.first_quantization.projectile.select_and_prepare
 import qualtran.bloqs.chemistry.pbc.first_quantization.select_t
 import qualtran.bloqs.chemistry.pbc.first_quantization.select_uv
@@ -119,6 +120,8 @@
 import qualtran.bloqs.rotations.phasing_via_cost_function
 import qualtran.bloqs.rotations.programmable_rotation_gate_array
 import qualtran.bloqs.rotations.quantum_variable_rotation
+import qualtran.bloqs.rotations.rz_via_phase_gradient
+import qualtran.bloqs.rotations.zpow_via_phase_gradient
 import qualtran.bloqs.state_preparation.black_box_prepare
 import qualtran.bloqs.state_preparation.prepare_base
 import qualtran.bloqs.state_preparation.prepare_uniform_superposition
@@ -309,6 +312,7 @@
             qualtran.bloqs.chemistry.pbc.first_quantization.prepare_uv._PREPARE_UV,
             qualtran.bloqs.chemistry.pbc.first_quantization.select_t._SELECT_T,
             qualtran.bloqs.chemistry.pbc.first_quantization.select_uv._SELECT_UV,
+            qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta._PREPARE_ZETA,
         ],
         directory=f'{SOURCE_DIR}/bloqs/chemistry/pbc/first_quantization',
     ),
@@ -634,6 +638,18 @@
             qualtran.bloqs.rotations.hamming_weight_phasing._HAMMING_WEIGHT_PHASING_VIA_PHASE_GRADIENT_DOC,
         ],
     ),
+    NotebookSpecV2(
+        title='ZPow Rotation via Phase Gradient',
+        module=qualtran.bloqs.rotations.zpow_via_phase_gradient,
+        bloq_specs=[
+            qualtran.bloqs.rotations.zpow_via_phase_gradient._ZPOW_CONST_VIA_PHASE_GRADIENT_DOC
+        ],
+    ),
+    NotebookSpecV2(
+        title='Rz Rotation via Phase Gradient',
+        module=qualtran.bloqs.rotations.rz_via_phase_gradient,
+        bloq_specs=[qualtran.bloqs.rotations.rz_via_phase_gradient._RZ_VIA_PHASE_GRADIENT_DOC],
+    ),
     NotebookSpecV2(
         title='Programmable Rotation Gate Array',
         module=qualtran.bloqs.rotations.programmable_rotation_gate_array,

docs/bloqs/index.rst

@@ -108,6 +108,8 @@ Bloqs Library
     rotations/phasing_via_cost_function.ipynb
     rotations/phase_gradient.ipynb
     rotations/hamming_weight_phasing.ipynb
+    rotations/zpow_via_phase_gradient.ipynb
+    rotations/rz_via_phase_gradient.ipynb
     rotations/programmable_rotation_gate_array.ipynb
     qft/two_bit_ffft.ipynb
     qft/approximate_qft.ipynb

qualtran/_infra/composite_bloq.py

@@ -1234,12 +1234,13 @@ def split(self, soq: Soquet) -> NDArray[Soquet]:  # type: ignore[type-var]
 
         return self.add(Split(dtype=soq.reg.dtype), reg=soq)
 
-    def join(self, soqs: NDArray[Soquet], dtype: Optional[QDType] = None) -> Soquet:  # type: ignore[type-var]
+    def join(self, soqs: SoquetInT, dtype: Optional[QDType] = None) -> Soquet:
         from qualtran.bloqs.bookkeeping import Join
 
         try:
+            soqs = np.asarray(soqs)
             (n,) = soqs.shape
-        except AttributeError:
+        except (AttributeError, ValueError):
             raise ValueError("`join` expects a 1-d array of input soquets to join.") from None
 
         if not all(soq.reg.bitsize == 1 for soq in soqs):

qualtran/_infra/composite_bloq_test.py

@@ -389,6 +389,12 @@ def test_util_convenience_methods():
     assert len(cbloq.connections) == 1 + 10 + 1
 
 
+def test_join_list():
+    bb = BloqBuilder()
+    qs = [bb.allocate() for _ in range(10)]
+    _ = bb.join(qs)
+
+
 def test_util_convenience_methods_errors():
     bb = BloqBuilder()
 

qualtran/bloqs/basic_gates/rotation.py

@@ -108,7 +108,7 @@ def _z_pow() -> ZPowGate:
 
 @frozen
 class CZPowGate(CirqGateAsBloqBase):
-    exponent: float = 1.0
+    exponent: SymbolicFloat = 1.0
     global_shift: float = 0.0
     eps: SymbolicFloat = 1e-11
 

qualtran/bloqs/chemistry/pbc/first_quantization/first_quantization.ipynb

@@ -822,6 +822,119 @@
     "show_call_graph(select_uv_g)\n",
     "show_counts_sigma(select_uv_sigma)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "69c98b9f",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.bloq_doc.md"
+   },
+   "source": [
+    "## `PrepareZetaState`\n",
+    "PREPARE the superpostion over $l$ weighted by $\\zeta_l$.\n",
+    "\n",
+    "See https://github.com/quantumlib/Qualtran/issues/473.\n",
+    "#### Parameters\n",
+    " - `num_bits_p`: The number of bits to represent each dimension of the momentum register.\n",
+    " - `eta`: The number of electrons.\n",
+    " - `m_param`: $\\mathcal{M}$ in the reference.\n",
+    " - `lambda_zeta`: sum of nuclear charges. \n",
+    "\n",
+    "#### Registers\n",
+    " - `l`: the register indexing the atomic number. \n",
+    "\n",
+    "#### References\n",
+    " - [Fault-Tolerant Quantum Simulations of Chemistry in First Quantization](https://arxiv.org/abs/2105.12767). page 23-24, last 3 paragraphs.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "074bc1b8",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.bloq_doc.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta import PrepareZetaState"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fd3071f7",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.example_instances.md"
+   },
+   "source": [
+    "### Example Instances"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31ee49a4",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.prepare_zeta"
+   },
+   "outputs": [],
+   "source": [
+    "num_atoms = 10\n",
+    "lambda_zeta = 10\n",
+    "num_bits_nuc_pos = 8\n",
+    "\n",
+    "prepare_zeta = PrepareZetaState(\n",
+    "    num_atoms=num_atoms, lambda_zeta=lambda_zeta, num_bits_nuc_pos=num_bits_nuc_pos\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4fb19c71",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.graphical_signature.md"
+   },
+   "source": [
+    "#### Graphical Signature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6498eee0",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.graphical_signature.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.drawing import show_bloqs\n",
+    "show_bloqs([prepare_zeta],\n",
+    "           ['`prepare_zeta`'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1e644753",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.call_graph.md"
+   },
+   "source": [
+    "### Call Graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fbef8e6e",
+   "metadata": {
+    "cq.autogen": "PrepareZetaState.call_graph.py"
+   },
+   "outputs": [],
+   "source": [
+    "from qualtran.resource_counting.generalizers import ignore_split_join\n",
+    "prepare_zeta_g, prepare_zeta_sigma = prepare_zeta.call_graph(max_depth=1, generalizer=ignore_split_join)\n",
+    "show_call_graph(prepare_zeta_g)\n",
+    "show_counts_sigma(prepare_zeta_sigma)"
+   ]
   }
  ],
  "metadata": {

qualtran/bloqs/chemistry/pbc/first_quantization/prepare_t_test.py

@@ -12,14 +12,14 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from qualtran.bloqs.basic_gates import Toffoli
 from qualtran.bloqs.chemistry.pbc.first_quantization.prepare import (
     UniformSuperpostionIJFirstQuantization,
 )
 from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_t import (
     _prepare_t,
     PrepareTFirstQuantization,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_prepare_t(bloq_autotester):
@@ -33,15 +33,12 @@ def test_prepare_kinetic_t_counts():
     n_eta = (eta - 1).bit_length()
     expected_cost = (14 * n_eta + 8 * b_r - 36) + 2 * (2 * num_bits_p + 9)
     uni = UniformSuperpostionIJFirstQuantization(eta, num_bits_rot_aa=b_r)
-    _, counts = uni.call_graph()
-    qual_cost = counts[Toffoli()]
+
+    qual_cost = get_cost_value(uni, QECGatesCost()).total_toffoli_only()
     uni = UniformSuperpostionIJFirstQuantization(eta, num_bits_rot_aa=b_r).adjoint()
-    _, counts = uni.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(uni, QECGatesCost()).total_toffoli_only()
     prep = PrepareTFirstQuantization(num_bits_p, eta, num_bits_rot_aa=b_r)
-    _, counts = prep.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_toffoli_only()
     prep = PrepareTFirstQuantization(num_bits_p, eta, num_bits_rot_aa=b_r).adjoint()
-    _, counts = prep.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_toffoli_only()
     assert qual_cost == expected_cost

qualtran/bloqs/chemistry/pbc/first_quantization/prepare_zeta.py

@@ -19,7 +19,7 @@
 import numpy as np
 from attrs import evolve, frozen
 
-from qualtran import Bloq, QAny, Register, Signature
+from qualtran import Bloq, bloq_example, BloqDocSpec, QAny, Register, Signature
 from qualtran.bloqs.basic_gates import Toffoli
 
 if TYPE_CHECKING:
@@ -65,3 +65,22 @@ def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         else:
 
             return {Toffoli(): self.lambda_zeta}
+
+
+@bloq_example
+def _prepare_zeta() -> PrepareZetaState:
+    num_atoms = 10
+    lambda_zeta = 10
+    num_bits_nuc_pos = 8
+
+    prepare_zeta = PrepareZetaState(
+        num_atoms=num_atoms, lambda_zeta=lambda_zeta, num_bits_nuc_pos=num_bits_nuc_pos
+    )
+    return prepare_zeta
+
+
+_PREPARE_ZETA = BloqDocSpec(
+    bloq_cls=PrepareZetaState,
+    import_line='from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta import PrepareZetaState',
+    examples=(_prepare_zeta,),
+)

qualtran/bloqs/chemistry/pbc/first_quantization/prepare_zeta_test.py

@@ -11,8 +11,8 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta import PrepareZetaState
+from qualtran.bloqs.chemistry.pbc.first_quantization.prepare_zeta import _prepare_zeta
 
 
-def test_uniform_superposition_ij():
-    prep = PrepareZetaState(num_atoms=10, lambda_zeta=20, num_bits_nuc_pos=8)
+def test_prepare_zeta(bloq_autotester):
+    bloq_autotester(_prepare_zeta)

qualtran/bloqs/chemistry/pbc/first_quantization/projectile/prepare_t_test.py

@@ -14,13 +14,13 @@
 
 import pytest
 
-from qualtran.bloqs.basic_gates import Toffoli
 from qualtran.bloqs.chemistry.pbc.first_quantization.projectile.prepare_t import (
     _prep_power_two_proj,
     _prep_t_proj,
     PreparePowerTwoStateWithProj,
     PrepareTFirstQuantizationWithProj,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_prep_t_proj(bloq_autotester):
@@ -39,13 +39,11 @@ def test_prepare_kinetic_t_proj_counts():
     expected_cost = 2 * (2 * num_bits_n + 9) + 2 * (num_bits_n - num_bits_p) + 20
     qual_cost = 0
     prep = PrepareTFirstQuantizationWithProj(num_bits_p, num_bits_n, eta, num_bits_rot_aa=b_r)
-    _, counts = prep.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_toffoli_only()
     prep = PrepareTFirstQuantizationWithProj(
         num_bits_p, num_bits_n, eta, num_bits_rot_aa=b_r
     ).adjoint()
-    _, counts = prep.call_graph()
-    qual_cost += counts[Toffoli()]
+    qual_cost += get_cost_value(prep, QECGatesCost()).total_toffoli_only()
     assert qual_cost == expected_cost
 
 

qualtran/bloqs/chemistry/pbc/first_quantization/projectile/select_t_test.py

@@ -11,11 +11,11 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from qualtran.bloqs.basic_gates import Toffoli
 from qualtran.bloqs.chemistry.pbc.first_quantization.projectile.select_t import (
     _sel_t_proj,
     SelectTFirstQuantizationWithProj,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_sel_t_proj(bloq_autotester):
@@ -25,5 +25,5 @@ def test_sel_t_proj(bloq_autotester):
 def test_select_kinetic_t_counts():
     num_bits_n = 6
     sel = SelectTFirstQuantizationWithProj(num_bits_n, 10)
-    _, counts = sel.call_graph()
-    assert counts[Toffoli()] == 5 * (num_bits_n - 1) + 2 + 1
+    toffolis = get_cost_value(sel, QECGatesCost()).total_toffoli_only()
+    assert toffolis == 5 * (num_bits_n - 1) + 2 + 1

qualtran/bloqs/chemistry/pbc/first_quantization/select_t_test.py

@@ -11,11 +11,11 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from qualtran.bloqs.basic_gates import Toffoli
 from qualtran.bloqs.chemistry.pbc.first_quantization.select_t import (
     _select_t,
     SelectTFirstQuantization,
 )
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 
 
 def test_select_t(bloq_autotester):
@@ -25,5 +25,5 @@ def test_select_t(bloq_autotester):
 def test_select_kinetic_t_counts():
     num_bits_p = 6
     sel = SelectTFirstQuantization(num_bits_p, 10)
-    _, counts = sel.call_graph()
-    assert counts[Toffoli()] == 5 * (num_bits_p - 1) + 2
+    toffolis = get_cost_value(sel, QECGatesCost()).total_toffoli_only()
+    assert toffolis == 5 * (num_bits_p - 1) + 2