Repository: PennyLaneAI/pennylane
Pull request #3956 (2023-03-31 22:08:22)

Add noncubic lattice for all electron resource estimation
**Context:**
Adding non-cubic lattice support for all electron resource estimation.

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/3956?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> :exclamation: No coverage uploaded for pull request base (`resource_estimation_pseudop@f7f7823`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).
> The diff coverage is `n/a`.

```diff
@@                      Coverage Diff                       @@
##             resource_estimation_pseudop    #3956   +/-   ##
==============================================================
  Coverage                               ?   99.68%           
==============================================================
  Files                                  ?      351           
  Lines                                  ?    31037           
  Branches                               ?        0           
==============================================================
  Hits                                   ?    30940           
  Misses                                 ?       97           
  Partials                               ?        0           
```




soranjh: [sc-36664]

Patch:
diff --git a/pennylane/resource/first_quantization.py b/pennylane/resource/first_quantization.py
index 898e7346a..b696b5538 100644
--- a/pennylane/resource/first_quantization.py
+++ b/pennylane/resource/first_quantization.py
@@ -18,6 +18,7 @@ non-Clifford gates for quantum algorithms in first quantization using a plane-wa
 # pylint: disable=no-self-use disable=too-many-arguments disable=too-many-instance-attributes
 import numpy
 from scipy import integrate
+
 from pennylane import numpy as np
 from pennylane.operation import AnyWires, Operation
 
@@ -27,7 +28,7 @@ class FirstQuantization(Operation):
     algorithm in first quantization using a plane-wave basis.
 
     To estimate the gate and qubit costs for implementing this method, the number of plane waves,
-    the number of electrons and the unit cell volume need to be defined. The costs can then be
+    the number of electrons and the lattice vectors need to be defined. The costs can then be
     computed using the functions :func:`~.pennylane.resource.FirstQuantization.gate_cost` and
     :func:`~.pennylane.resource.FirstQuantization.qubit_cost` with a target error that has the default
     value of 0.0016 Ha (chemical accuracy). Atomic units are used throughout the class.
@@ -39,18 +40,21 @@ class FirstQuantization(Operation):
         error (float): target error in the algorithm
         charge (int): total electric charge of the system
         br (int): number of bits for ancilla qubit rotation
+        vectors (array[float]): lattice vectors
 
     **Example**
 
     >>> n = 100000
     >>> eta = 156
-    >>> omega = 1145.166
-    >>> algo = FirstQuantization(n, eta, omega)
+    >>> vectors = np.array([[10.46219511,  0.00000000,  0.00000000],
+    ...                     [ 0.00000000, 10.46219511,  0.00000000],
+    ...                     [ 0.00000000,  0.00000000, 10.46219511]])
+    >>> algo = FirstQuantization(n, eta, vectors=vectors)
     >>> print(algo.lamb,  # the 1-Norm of the Hamiltonian
     >>>       algo.gates, # estimated number of non-Clifford gates
     >>>       algo.qubits # estimated number of logical qubits
     >>>       )
-    649912.4801542697 1.10e+13 4416
+    649912.4804278888 1.1e+13 4416
 
     .. details::
         :title: Theory
@@ -83,10 +87,11 @@ class FirstQuantization(Operation):
         self,
         n,
         eta,
-        omega,
+        omega=None,
         error=0.0016,
         charge=0,
         br=7,
+        vectors=None,
     ):
         self.n = n
         self.eta = eta
@@ -94,12 +99,39 @@ class FirstQuantization(Operation):
         self.error = error
         self.charge = charge
         self.br = br
+        self.vectors = vectors
+        self.cubic = True
+
+        if omega is None and vectors is None:
+            raise ValueError("The lattice vectors must be provided.")
+
+        if omega is not None and vectors is not None:
+            raise ValueError(
+                "The lattice vectors and the unit cell volume should not be provided"
+                " together. Please provide only one of them."
+            )
 
-        self.lamb = self.norm(self.n, self.eta, self.omega, self.error, self.br, self.charge)
+        if self.vectors is not None:
+            self.omega = np.abs(np.sum((np.cross(vectors[0], vectors[1]) * vectors[2])))
+            recip_vectors = (  # taken from Eq. (35) of arXiv:2302.07981v1 (2023)
+                2
+                * np.pi
+                / self.omega
+                * np.array([np.cross(vectors[i], vectors[j]) for i, j in [(1, 2), (2, 0), (0, 1)]])
+            )
+            bbt = np.matrix(recip_vectors) @ np.matrix(recip_vectors).T
+            self.cubic = np.linalg.norm(bbt - (recip_vectors**2).max() * np.identity(3)) < 1e-6
+
+        self.lamb = self.norm(
+            self.n, self.eta, self.omega, self.error, self.br, self.charge, self.cubic, self.vectors
+        )
+
+        self.gates = self.gate_cost(
+            self.n, self.eta, self.omega, self.error, self.br, self.charge, self.cubic, self.vectors
+        )
 
-        self.gates = self.gate_cost(self.n, self.eta, self.omega, self.error, self.br, self.charge)
         self.qubits = self.qubit_cost(
-            self.n, self.eta, self.omega, self.error, self.br, self.charge
+            self.n, self.eta, self.omega, self.error, self.br, self.charge, self.cubic, self.vectors
         )
 
         super().__init__(wires=range(self.qubits))
@@ -141,7 +173,7 @@ class FirstQuantization(Operation):
         return p
 
     @staticmethod
-    def norm(n, eta, omega, error, br=7, charge=0):
+    def norm(n, eta, omega, error, br=7, charge=0, cubic=True, vectors=None):
         r"""Return the 1-norm of a first-quantized Hamiltonian in the plane-wave basis.
 
         The expressions needed for computing the norm are taken from
@@ -155,6 +187,8 @@ class FirstQuantization(Operation):
             error (float): target error in the algorithm
             br (int): number of bits for ancilla qubit rotation
             charge (int): total electric charge of the system
+            cubic (bool): True if the unit cell is cubic
+            vectors (array[float]): lattice vectors
 
         Returns:
             float: 1-norm of a first-quantized Hamiltonian in the plane-wave basis
@@ -250,6 +284,9 @@ class FirstQuantization(Operation):
         if not isinstance(charge, int):
             raise ValueError("system charge must be an integer.")
 
+        if not cubic:
+            return FirstQuantization._norm_noncubic(n, eta, error, br, charge, vectors)[0]
+
         l_z = eta + charge
 
         # target error in the qubitization of U+V which we set to be 0.01 of the algorithm error
@@ -380,12 +417,12 @@ class FirstQuantization(Operation):
         if not isinstance(charge, int):
             raise ValueError("system charge must be an integer.")
 
-        lamb = FirstQuantization.norm(n, eta, omega, error, br=7, charge=charge)
+        lamb = FirstQuantization.norm(n, eta, omega, error, br, charge)
         alpha = 0.01
         l_z = eta + charge
         l_nu = 2 * np.pi * n ** (2 / 3)
 
-        # n_eta and n_etaz are defined in the third and second paragraphs of page 040332-15
+        # defined in the third and second paragraphs of page 15 of PRX Quantum 2, 040332 (2021)
         n_eta = np.ceil(np.log2(eta))
         n_etaz = np.ceil(np.log2(eta + 2 * l_z))
 
@@ -421,7 +458,7 @@ class FirstQuantization(Operation):
         return int(np.ceil(cost))
 
     @staticmethod
-    def estimation_cost(n, eta, omega, error, br=7, charge=0):
+    def estimation_cost(n, eta, omega, error, br=7, charge=0, cubic=True, vectors=None):
         r"""Return the number of calls to the unitary needed to achieve the desired error in quantum
         phase estimation.
 
@@ -435,6 +472,8 @@ class FirstQuantization(Operation):
             error (float): target error in the algorithm
             br (int): number of bits for ancilla qubit rotation
             charge (int): total electric charge of the system
+            cubic (bool): True if the unit cell is cubic
+            vectors (array[float]): lattice vectors
 
         Returns:
             int: number of calls to unitary
@@ -451,15 +490,23 @@ class FirstQuantization(Operation):
         if error <= 0.0:
             raise ValueError("The target error must be greater than zero.")
 
-        lamb = FirstQuantization.norm(n, eta, omega, error, br=br, charge=charge)
-        alpha = 0.01
-        # qpe_error obtained to satisfy inequality (131)
-        error_qpe = np.sqrt(error**2 * (1 - (3 * alpha) ** 2))
+        lamb = FirstQuantization.norm(
+            n, eta, omega, error, br=br, charge=charge, cubic=cubic, vectors=vectors
+        )
+
+        if cubic:
+            alpha = 0.01
+            # qpe_error obtained to satisfy inequality (131)
+            error_qpe = np.sqrt(error**2 * (1 - (3 * alpha) ** 2))
+        else:
+            alpha = 0.0248759298
+            # qpe_error obtained to satisfy inequality (131)
+            error_qpe = np.sqrt(error**2 * (1 - (4 * alpha) ** 2))
 
         return int(np.ceil(np.pi * lamb / (2 * error_qpe)))
 
     @staticmethod
-    def gate_cost(n, eta, omega, error, br=7, charge=0):
+    def gate_cost(n, eta, omega, error, br=7, charge=0, cubic=True, vectors=None):
         r"""Return the total number of Toffoli gates needed to implement the first quantization
         algorithm.
 
@@ -473,6 +520,8 @@ class FirstQuantization(Operation):
             error (float): target error in the algorithm
             br (int): number of bits for ancilla qubit rotation
             charge (int): total electric charge of the system
+            cubic (bool): True if the unit cell is cubic
+            vectors (array[float]): lattice vectors
 
         Returns:
             int: the number of Toffoli gates needed to implement the first quantization algorithm
@@ -504,13 +553,19 @@ class FirstQuantization(Operation):
         if br <= 0 or not isinstance(br, int):
             raise ValueError("br must be a positive integer.")
 
-        e_cost = FirstQuantization.estimation_cost(n, eta, omega, error, br=br, charge=charge)
-        u_cost = FirstQuantization.unitary_cost(n, eta, omega, error, br, charge)
+        e_cost = FirstQuantization.estimation_cost(
+            n, eta, omega, error, br=br, charge=charge, cubic=cubic, vectors=vectors
+        )
+
+        if cubic:
+            u_cost = FirstQuantization.unitary_cost(n, eta, omega, error, br, charge)
+        else:
+            u_cost = FirstQuantization._unitary_cost_noncubic(n, eta, error, br, charge, vectors)
 
         return e_cost * u_cost
 
     @staticmethod
-    def qubit_cost(n, eta, omega, error, br=7, charge=0):
+    def qubit_cost(n, eta, omega, error, br=7, charge=0, cubic=True, vectors=None):
         r"""Return the number of logical qubits needed to implement the first quantization
         algorithm.
 
@@ -524,6 +579,8 @@ class FirstQuantization(Operation):
             error (float): target error in the algorithm
             br (int): number of bits for ancilla qubit rotation
             charge (int): total electric charge of the system
+            cubic (bool): True if the unit cell is cubic
+            vectors (array[float]): lattice vectors
 
         Returns:
             int: number of logical qubits needed to implement the first quantization algorithm
@@ -552,6 +609,9 @@ class FirstQuantization(Operation):
         if not isinstance(charge, int):
             raise ValueError("system charge must be an integer.")
 
+        if not cubic:
+            return FirstQuantization._qubit_cost_noncubic(n, eta, error, br, charge, vectors)
+
         lamb = FirstQuantization.norm(n, eta, omega, error, br=br, charge=charge)
         alpha = 0.01
         l_z = eta + charge
@@ -589,3 +649,359 @@ class FirstQuantization(Operation):
         qubits += np.maximum(5 * n_p + 1, 5 * n_r - 4) + np.maximum(n_t, n_r + 1) + 33
 
         return int(np.ceil(qubits))
+
+    @staticmethod
+    def _norm_noncubic(
+        n,
+        eta,
+        error,
+        br,
+        charge,
+        vectors,
+    ):
+        r"""Return the 1-norm of a first-quantized Hamiltonian in the plane-wave basis
+        for non-cubic systems.
+
+        Args:
+            n (int): number of plane waves
+            eta (int): number of electrons
+            error (float): target error in the algorithm
+            br (int): number of bits for ancilla qubit rotation
+            charge (int): total electric charge of the system
+            vectors (array[float]): lattice vectors
+
+        Returns:
+            float: 1-norm of a first-quantized Hamiltonian in the plane-wave basis
+        """
+        omega = np.abs(np.sum((np.cross(vectors[0], vectors[1]) * vectors[2])))
+
+        recip_vectors = (  # taken from Eq. (35) of arXiv:2302.07981v1 (2023)
+            2
+            * np.pi
+            / omega
+            * np.array([np.cross(vectors[i], vectors[j]) for i, j in [(1, 2), (2, 0), (0, 1)]])
+        )
+
+        bbt = np.matrix(recip_vectors) @ np.matrix(recip_vectors).T
+
+        orthogonal = (
+            np.linalg.norm(bbt - np.array([np.max(b**2) for b in recip_vectors]) * np.identity(3))
+            < 1e-6
+        )
+
+        l_z = eta + charge
+
+        alpha = 0.0248759298  # optimal value for lower resource estimates
+        p_th = 0.95  # optimal value for lower resource estimates
+
+        error_uv = alpha * error
+
+        # taken from Eq. (22) of PRX Quantum 2, 040332 (2021)
+        n_p = int(np.ceil(np.log2(n ** (1 / 3) + 1)))
+
+        n0 = n ** (1 / 3)
+
+        # defined in Eq. (F3) of arXiv:2302.07981v1 (2023)
+        bmin = np.min(np.linalg.svd(recip_vectors)[1])
+
+        n_m = int(
+            np.ceil(
+                np.log2(  # taken from Eq. (132) of PRX Quantum 2, 040332 (2021) with
+                    # modifications taken from arXiv:2302.07981v1 (2023)
+                    (8 * np.pi * eta)
+                    / (error_uv * omega * bmin**2)
+                    * (eta - 1 + 2 * l_z)
+                    * (7 * 2 ** (n_p + 1) - 9 * n_p - 11 - 3 * 2 ** (-1 * n_p))
+                )
+            )
+        )
+
+        lambda_nu = (  # expression is taken from Eq. (F6) of PRX 8, 011044 (2018)
+            # bmin**2 is needed for non-cubic systems, from Eq. (F3) of arXiv:2302.07981v1 (2023)
+            4 * np.pi * (np.sqrt(3) * n ** (1 / 3) / 2 - 1)
+            + 3
+            - 3 / n ** (1 / 3)
+            + 3 * integrate.nquad(lambda x, y: 1 / (x**2 + y**2), [[1, n0], [1, n0]])[0]
+        ) / bmin**2
+
+        # computed using error term derived in Eq. (113) of PRX Quantum 2, 040332 (2021)
+        # bmin**2 is needed for non-cubic systems, from Eq. (F3) of arXiv:2302.07981v1 (2023)
+        lambda_nu_1 = lambda_nu + 4 / (2**n_m * bmin**2) * (
+            7 * 2 ** (n_p + 1) - 9 * n_p - 11 - 3 * 2 ** (-1 * n_p)
+        )
+
+        # taken from Eq. (I7) of arXiv:2302.07981v1 (2023)
+        p_nu = lambda_nu_1 * bmin**2 / 2 ** (n_p + 6)
+
+        # amplitude amplification steps, taken from Eq. (117) of PRX Quantum 2, 040332 (2021)
+        p_nu_amp = 0
+        aa_steps = 0
+        n_steps = 30  # maximum number of amplitude amplification steps
+        for i in reversed(range(n_steps)):
+            probability = (np.sin((2 * i + 1) * np.arcsin(np.sqrt(p_nu)))) ** 2
+            if probability > p_th:
+                aa_steps = i
+                p_nu_amp = probability
+
+        # lambda_u and lambda_v are taken from Eq. (25) of PRX Quantum 2, 040332 (2021) with
+        # modifications taken from arXiv:2302.07981v1 (2023)
+        lambda_u = 4 * np.pi * eta * l_z * lambda_nu / omega
+        lambda_v = 2 * np.pi * eta * (eta - 1) * lambda_nu / omega
+
+        # lambda_u_1 and lambda_v_1 are taken from Eq. (124) of PRX Quantum 2, 040332 (2021)
+        lambda_u_1 = lambda_u * lambda_nu_1 / lambda_nu
+        lambda_v_1 = lambda_v * lambda_nu_1 / lambda_nu
+
+        b_mat = np.matrix(recip_vectors)
+        abs_sum = np.abs(b_mat @ b_mat.T).flatten().sum()
+
+        # taken from Appendix I.1 of arXiv:2302.07981v1 (2023)
+        if orthogonal:
+            lambda_t_p = abs_sum * eta * 2 ** (2 * n_p - 2) / 4
+        else:
+            lambda_t_p = abs_sum * eta * 2 ** (2 * n_p - 2) / 2
+
+        # taken from Eq. (63) of PRX Quantum 2, 040332 (2021) with
+        # modifications taken from arXiv:2302.07981v1 (2023)
+        p_eq = (
+            FirstQuantization.success_prob(3 * eta + 2 * charge, br)
+            * FirstQuantization.success_prob(eta, br) ** 2
+        )
+
+        # final lambda value is computed from case checking derived in Section E and Appendix D of
+        # PRX Quantum 2, 040332 (2021)
+        if p_nu * lambda_t_p >= (1 - p_nu) * (lambda_u_1 + lambda_v_1):
+            raise ValueError("The computed 1-norm is zero.")  # return 0.0, 0.0
+        if p_nu_amp * lambda_t_p >= (1 - p_nu_amp) * (lambda_u_1 + lambda_v_1):
+            return (lambda_t_p + lambda_u_1 + lambda_v_1) / p_eq, aa_steps
+        return ((lambda_u_1 + lambda_v_1 / (1 - 1 / eta)) / p_nu_amp) / p_eq, aa_steps
+
+    @staticmethod
+    def _qubit_cost_noncubic(n, eta, error, br, charge, vectors):
+        r"""Return the number of logical qubits needed to implement the first quantization
+        algorithm for non-cubic systems.
+
+        Args:
+            n (int): number of plane waves
+            eta (int): number of electrons
+            error (float): target error in the algorithm
+            br (int): number of bits for ancilla qubit rotation
+            charge (int): total electric charge of the system
+            vectors (array[float]): lattice vectors
+
+        Returns:
+            int: number of logical qubits needed to implement the first quantization algorithm
+        """
+
+        lambda_total, _ = FirstQuantization._norm_noncubic(
+            n,
+            eta,
+            error,
+            br,
+            charge,
+            vectors,
+        )
+
+        omega = np.abs(np.sum((np.cross(vectors[0], vectors[1]) * vectors[2])))
+
+        recip_vectors = (  # taken from Eq. (35) of arXiv:2302.07981v1 (2023)
+            2
+            * np.pi
+            / omega
+            * np.array([np.cross(vectors[i], vectors[j]) for i, j in [(1, 2), (2, 0), (0, 1)]])
+        )
+
+        l_z = eta + charge
+        l_nu = 2 * np.pi * n ** (2 / 3)
+
+        # taken from Eq. (22) of PRX Quantum 2, 040332 (2021)
+        n_p = np.ceil(np.log2(n ** (1 / 3) + 1))
+
+        # defined in Eq. (F3) of arXiv:2302.07981v1 (2023)
+        bmin = np.min(np.linalg.svd(recip_vectors)[1])
+
+        alpha = 0.0248759298  # optimal value for lower resource estimates
+
+        # errors derived according to the strategy in Appendix M.1 of arXiv:2302.07981v1 (2023)
+        error_t, error_r, error_m, error_b = [alpha * error] * 4
+
+        # parameters taken from PRX Quantum 2, 040332 (2021)
+        n_t = int(np.ceil(np.log2(np.pi * lambda_total / error_t)))  # Eq. (134)
+        n_r = int(np.ceil(np.log2((eta * l_z * l_nu) / (error_r * omega ** (1 / 3)))))  # Eq. (133)
+        n_m = int(  # taken from Eq. (J13) arXiv:2302.07981v1 (2023)
+            np.ceil(
+                np.log2(
+                    (8 * np.pi * eta)
+                    / (error_m * omega * bmin**2)
+                    * (eta - 1 + 2 * l_z)
+                    * (7 * 2 ** (n_p + 1) - 9 * n_p - 11 - 3 * 2 ** (-1 * n_p))
+                )
+            )
+        )
+
+        # qpe error obtained to satisfy a modification of
+        # inequality Eq. (131) of PRX Quantum 2, 040332 (2021) including error_b
+        n_errors = 4
+        error_qpe = np.sqrt(error**2 * (1 - (n_errors * alpha) ** 2))
+
+        # adapted from equations (L1, L2) in Appendix L of arXiv:2302.07981v1 (2023)
+        clean_temp_H_cost = max([5 * n_r - 4, 5 * n_p + 1]) + max([5, n_m + 3 * n_p])
+        reflection_cost = (
+            np.ceil(np.log2(eta + 2 * l_z)) + 2 * np.ceil(np.log2(eta)) + 6 * n_p + n_m + 16 + 3
+        )
+        clean_temp_cost = max([clean_temp_H_cost, reflection_cost])
+
+        # the expression for computing the cost is taken from Appendix C of
+        # PRX Quantum 2, 040332 (2021) and adapting to non-cubic using Appendix L of
+        # arXiv:2302.07981v1 (2023)
+        clean_cost = 3 * eta * n_p
+        clean_cost += np.ceil(np.log2(np.ceil(np.pi * lambda_total / (2 * error_qpe))))
+        clean_cost += 1 + 1 + np.ceil(np.log2(eta + 2 * l_z)) + 3 + 3
+        clean_cost += 2 * np.ceil(np.log2(eta)) + 5 + 3 * (n_p + 1)
+        clean_cost += n_p + n_m + 3 * n_p + 2 + 2 * n_p + 1 + 1 + 2 + 2 * n_p + 6 + 1
+
+        clean_cost += clean_temp_cost
+
+        # taken from Eq. (J7) of arXiv:2302.07981v1 (2023)
+        n_b = np.ceil(
+            np.log2(
+                4
+                * np.pi
+                * eta
+                * 2 ** (2 * n_p - 2)
+                * np.abs(np.matrix(recip_vectors) @ np.matrix(recip_vectors).T).flatten().sum()
+                / error_b
+            )
+        )
+        clean_cost += np.max([n_r + 1, n_t, n_b]) + 6 + n_m + 1
+
+        return int(np.ceil(clean_cost))
+
+    @staticmethod
+    def _unitary_cost_noncubic(n, eta, error, br, charge, vectors):
+        r"""Return the number of Toffoli gates needed to implement the qubitization unitary
+        operator for non-cubic systems.
+
+        Args:
+            n (int): number of plane waves
+            eta (int): number of electrons
+            error (float): target error in the algorithm
+            br (int): number of bits for ancilla qubit rotation
+            charge (int): total electric charge of the system
+            vectors (array[float]): lattice vectors
+
+        Returns:
+            int: the number of Toffoli gates needed to implement the qubitization unitary operator
+        """
+
+        lambda_total, aa_steps = FirstQuantization._norm_noncubic(
+            n,
+            eta,
+            error,
+            br,
+            charge,
+            vectors,
+        )
+
+        omega = np.abs(np.sum((np.cross(vectors[0], vectors[1]) * vectors[2])))
+
+        recip_vectors = (  # taken from Eq. (35) of arXiv:2302.07981v1 (2023)
+            2
+            * np.pi
+            / omega
+            * np.array([np.cross(vectors[i], vectors[j]) for i, j in [(1, 2), (2, 0), (0, 1)]])
+        )
+
+        n_tof = 500  # optimal value for lower resource estimates
+
+        alpha = 0.0248759298  # optimal value for lower resource estimates
+        l_z = eta + charge
+        l_nu = 2 * np.pi * n ** (2 / 3)
+
+        # defined in the third and second paragraphs of page 15 of PRX Quantum 2, 040332 (2021)
+        n_eta = np.ceil(np.log2(eta))
+        n_etaz = np.ceil(np.log2(eta + 2 * l_z))
+
+        # taken from Eq. (22) of PRX Quantum 2, 040332 (2021)
+        n_p = int(np.ceil(np.log2(n ** (1 / 3) + 1)))
+
+        # errors in Eqs. (132-134) of PRX Quantum 2, 040332 (2021)
+        error_t, error_r, error_m = [alpha * error] * 3
+
+        # parameters taken from PRX Quantum 2, 040332 (2021)
+        n_t = int(np.ceil(np.log2(np.pi * lambda_total / error_t)))  # Eq. (134)
+        n_r = int(np.ceil(np.log2((eta * l_z * l_nu) / (error_r * omega ** (1 / 3)))))  # Eq. (133)
+
+        # defined in Eq. (F3) of arXiv:2302.07981v1 (2023)
+        bmin = np.min(np.linalg.svd(recip_vectors)[1])
+
+        # equivalent to Eq. (J13) of arXiv:2302.07981v1 (2023)
+        n_m = int(
+            np.ceil(
+                np.log2(
+                    (8 * np.pi * eta)
+                    / (error_m * omega * bmin**2)
+                    * (eta - 1 + 2 * l_z)
+                    * (7 * 2 ** (n_p + 1) - 9 * n_p - 11 - 3 * 2 ** (-1 * n_p))
+                )
+            )
+        )
+
+        e_r = FirstQuantization._cost_qrom(l_z)
+
+        # taken from Eq. (125) of PRX Quantum 2, 040332 (2021)
+        cost = 2 * (n_t + 4 * n_etaz + 2 * br - 12) + 14 * n_eta + 8 * br - 36
+        cost += (2 * aa_steps + 1) * (3 * n_p**2 + 15 * n_p - 7 + 4 * n_m * (n_p + 1))
+        cost += l_z + e_r + 2 * (2 * n_p + 2 * br - 7) + 12 * eta * n_p
+        cost += 5 * (n_p - 1) + 2 + 24 * n_p + 6 * n_p * n_r + 18
+        cost += n_etaz + 2 * n_eta + 6 * n_p + n_m + 16
+
+        # taken from Eq. (J7) of arXiv:2302.07981v1 (2023)
+        error_b = alpha * error
+        n_b = np.ceil(
+            np.log2(
+                2
+                * np.pi
+                * eta
+                * 2 ** (2 * n_p - 2)
+                * np.abs(np.matrix(recip_vectors) @ np.matrix(recip_vectors).T).flatten().sum()
+                / error_b
+            )
+        )
+
+        n_dirty = FirstQuantization._qubit_cost_noncubic(n, eta, error, br, charge, vectors)
+        ms_cost = FirstQuantization._momentum_state_qrom(n_p, n_m, n_dirty, n_tof, kappa=1)[0]
+
+        # adapted Eq. (125) of PRX Quantum 2, 040332 (2021) to the noncubic case using Appendix K of
+        # arXiv:2302.07981v1 (2023)
+        cost -= 2 * (3 * 2 + 2 * br - 9)
+        cost += 2 * (2 * (2 * (2 ** (4 + 1) - 1) + (n_b - 3) * 4 + 2**4 + (n_p - 2)))
+        cost += 8
+        cost -= (2 * aa_steps + 1) * (3 * n_p**2 + 15 * n_p - 7 + 4 * n_m * (n_p + 1))
+        cost += (2 * aa_steps + 1) * ms_cost
+
+        return int(np.ceil(cost))
+
+    @staticmethod
+    def _momentum_state_qrom(n_p, n_m, n_dirty, n_tof, kappa):
+        r"""Returns the Toffoli cost for preparing the momentum state superposition.
+
+        Derived from Section D.1 item (6) and Appendix K.1.f of arXiv:2302.07981v1 (2023)"""
+
+        x = 2 ** (3 * n_p)
+
+        beta_dirty = max([np.floor(n_dirty / n_m), 1])
+        beta_parallel = max([np.floor(n_tof / kappa), 1])
+
+        if n_tof == 1:
+            beta_gate = max([np.floor(np.sqrt(2 * x / (3 * n_m))), 1])
+            beta = np.min([beta_dirty, beta_gate])
+            ms_cost_qrom = 2 * np.ceil(x / beta) + 3 * n_m * beta
+        else:
+            beta_gate = max([np.floor(2 * x / (3 * n_m / kappa) * np.log(2)), 1])
+            beta = np.min([beta_dirty, beta_gate, beta_parallel])
+            ms_cost_qrom = 2 * np.ceil(x / beta) + 3 * np.ceil(n_m / kappa) * np.ceil(np.log2(beta))
+
+        ms_cost = 2 * ms_cost_qrom + n_m + 8 * (n_p - 1) + 6 * n_p + 2 + 2 * n_p + n_m + 2
+
+        return ms_cost, beta


Test patch:
diff --git a/tests/resource/test_first_quantization.py b/tests/resource/test_first_quantization.py
index b0b5a8433..818e497db 100644
--- a/tests/resource/test_first_quantization.py
+++ b/tests/resource/test_first_quantization.py
@@ -264,3 +264,140 @@ def test_norm_error(n, eta, omega, error, br, charge):
     r"""Test that norm raises an error with incorrect inputs."""
     with pytest.raises(ValueError, match="must be"):
         qml.resource.FirstQuantization.norm(n, eta, omega, error, br, charge)
+
+
+@pytest.mark.parametrize(
+    ("n", "eta", "omega", "vectors", "lamb", "g_cost", "q_cost"),
+    [
+        (
+            100000,
+            156,
+            None,
+            np.array(
+                [
+                    [9.44862994, 0.0, 0.0],
+                    [0.0, 10.39349294, 0.0],
+                    [0.94486299, 0.94486299, 11.33835593],
+                ]
+            ),
+            817051.632523202,
+            151664625909497,
+            3331,
+        ),
+        (
+            10000,
+            312,
+            None,
+            np.array(
+                [
+                    [18.89725988, 0.0, 0.0],
+                    [0.0, 10.39349294, 0.0],
+                    [0.94486299, 0.94486299, 11.33835593],
+                ]
+            ),
+            1793399.3143809892,
+            64986483274430,
+            5193,
+        ),
+        (
+            100000,
+            156,
+            None,
+            np.array(
+                [
+                    [9.44862994, 0.0, 0.0],
+                    [0.0, 10.39349294, 0.0],
+                    [0.0, 0.0, 11.33835593],
+                ]
+            ),
+            725147.0916537816,
+            134604911168852,
+            3331,
+        ),
+    ],
+)
+def test_fq_vals_non_qubic(n, eta, omega, vectors, lamb, g_cost, q_cost):
+    r"""Test that the FirstQuantization class computes correct attributes."""
+    est = qml.resource.FirstQuantization(n, eta, omega, vectors=vectors)
+
+    assert np.allclose(est.lamb, lamb)
+    assert np.allclose(est.gates, g_cost)
+    assert np.allclose(est.qubits, q_cost)
+
+
+@pytest.mark.parametrize(
+    ("n", "eta", "omega", "error", "br", "charge", "vectors"),
+    [
+        (10000, 156, None, 0.001, 7, 0, None),
+    ],
+)
+def test_init_error_1(n, eta, omega, error, br, charge, vectors):
+    r"""Test that init raises an error when volume and vectors are None."""
+    with pytest.raises(ValueError, match="The lattice vectors must be provided"):
+        qml.resource.FirstQuantization(n, eta, omega, error, br, charge, vectors)
+
+
+@pytest.mark.parametrize(
+    ("n", "eta", "omega", "error", "br", "charge", "vectors"),
+    [
+        (
+            10000,
+            156,
+            1113.47,
+            0.001,
+            7,
+            0,
+            np.array(
+                [
+                    [9.0, 0.0, 0.0],
+                    [0.0, 9.0, 0.0],
+                    [9.0, 9.0, 9.0],
+                ]
+            ),
+        ),
+    ],
+)
+def test_init_error_2(n, eta, omega, error, br, charge, vectors):
+    r"""Test that init raises an error when volume and vectors are None."""
+    with pytest.raises(ValueError, match="lattice vectors and the unit cell volume should not be"):
+        qml.resource.FirstQuantization(n, eta, omega, error, br, charge, vectors)
+
+
+@pytest.mark.parametrize(
+    ("n", "eta", "error", "br", "charge", "vectors"),
+    [
+        (
+            1e10,
+            100,
+            0.0016,
+            7,
+            0,
+            np.array(
+                [
+                    [10.0, 0.0, 0.0],
+                    [0.0, 10.0, 0.0],
+                    [1.0, 1.0, 10.0],
+                ]
+            ),
+        ),
+    ],
+)
+def test_norm_error_noncubic(n, eta, error, br, charge, vectors):
+    r"""Test that _norm_noncubic raises an error when the computed norm is zero."""
+    print(n, eta, error, br, charge, vectors)
+    with pytest.raises(ValueError, match="The computed 1-norm is zero"):
+        qml.resource.FirstQuantization._norm_noncubic(n, eta, error, br, charge, vectors)
+
+
+@pytest.mark.parametrize(
+    ("n_p, n_m, n_dirty, n_tof, kappa, ms_cost_ref, beta_ref"),
+    [(6, 37, 3331, 500, 1, 13372.0, 90.0), (6, 37, 3331, 1, 1, 30686.0, 68.0)],
+)
+def test_momentum_state_qrom(n_p, n_m, n_dirty, n_tof, kappa, ms_cost_ref, beta_ref):
+    r"""Test that the _momentum_state_qrom function returns correct values."""
+    ms_cost, beta = qml.resource.FirstQuantization._momentum_state_qrom(
+        n_p, n_m, n_dirty, n_tof, kappa
+    )
+
+    assert ms_cost == ms_cost_ref
+    assert beta == beta_ref