Improve efficiency and fix bugs (device and dtype) for Gaussian backend

src/deepquantum/photonic/circuit.py

@@ -490,7 +490,8 @@ def _forward_gaussian_prob_helper(self, cov, mean, even_basis, odd_basis, basis,
                 probs_i = torch.cat(probs_half)
             else:
                 probs_half = torch.cat(probs_half)
-                probs_i = torch.cat([probs_half.squeeze(), torch.zeros(len(torch.cat(odd_basis)))])
+                probs_i = torch.cat([probs_half.squeeze(),
+                                     torch.zeros(len(torch.cat(odd_basis)), device=probs_half.device)])
             probs_i = probs_i.squeeze()
         if detector == 'threshold':
             probs_i = []
@@ -843,14 +844,13 @@ def _get_probs_gaussian_helper(
         assert final_states.ndim == 2
         nmode = final_states.shape[-1]
         final_states = final_states.to(cov.device)
-        identity = torch.eye(nmode, dtype=cov.dtype, device=cov.device)
-        identity2 = torch.eye(2 * nmode, dtype=cov.dtype, device=cov.device)
+        identity = torch.eye(2 * nmode, dtype=cov.dtype, device=cov.device)
         cov_ladder = quadrature_to_ladder(cov)
         mean_ladder = quadrature_to_ladder(mean)
-        q = cov_ladder + identity2 / 2
+        q = cov_ladder + identity / 2
         det_q = torch.det(q)
-        x_mat = torch.block_diag(identity.fliplr(), identity.fliplr()).fliplr() + 0j
-        o_mat = identity2 - torch.inverse(q)
+        x_mat = identity.reshape(2, nmode, 2 * nmode).flip(0).reshape(2 * nmode, 2 * nmode) + 0j
+        o_mat = identity - torch.inverse(q)
         a_mat = x_mat @ o_mat
         gamma = mean_ladder.conj().mT @ torch.inverse(q)
         if detector == 'pnrd':
@@ -897,7 +897,7 @@ def _get_prob_gaussian_base(
                 haf = abs(hafnian(sub_mat, loop=loop)) ** 2
             else:
                 haf = hafnian(sub_mat, loop=loop)
-            prob = p_vac * haf / product_factorial(final_state)
+            prob = p_vac * haf / product_factorial(final_state).to(device=haf.device, dtype=haf.dtype)
         elif detector == 'threshold':
             final_state_double = torch.cat([final_state, final_state])
             sub_mat = sub_matrix(matrix, final_state_double, final_state_double)

src/deepquantum/photonic/hafnian_.py

@@ -54,15 +54,16 @@ def get_submat_haf(a: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
 
 def poly_lambda(submat: torch.Tensor, int_partition: List, power: int, loop: bool = False) -> torch.Tensor:
     """Get the coefficient of the polynomial."""
-    sigma_x_list = [torch.tensor([[0, 1], [1, 0]], dtype=submat.dtype, device=submat.device)] * (submat.shape[-1] // 2)
-    x_mat = torch.block_diag(*sigma_x_list)
+    size = submat.shape[-1]
+    identity = torch.eye(size, dtype=submat.dtype, device=submat.device)
+    x_mat = identity.reshape(size // 2, 2, size).flip(1).reshape(size, size)
     xaz = x_mat @ submat
     eigen = torch.linalg.eigvals(xaz) # eigen decomposition
     trace_list = torch.stack([(eigen ** i).sum() for i in range(0, power + 1)])
     coeff = 0
     if loop: # loop hafnian case
         v = torch.diag(submat)
-        diag_term = torch.stack([v @ torch.linalg.matrix_power(xaz, i - 1) @ x_mat @ v / 2 for i in range(1, power + 1)])
+        diag_term = torch.stack([v @ torch.linalg.matrix_power(xaz, i - 1) @ x_mat @ v / 2 for i in range(1, power+1)])
     for orders in int_partition:
         ncount = count_unique_permutations(orders)
         orders = torch.tensor(orders, device=submat.device)

src/deepquantum/photonic/operation.py

@@ -115,8 +115,8 @@ def get_unitary(self) -> torch.Tensor:
         """Get the global unitary matrix acting on creation operators."""
         matrix = self.update_matrix()
         assert matrix.shape[-2] == matrix.shape[-1] == len(self.wires), 'The matrix may not act on creation operators.'
-        u = matrix.new_zeros(self.nmode, self.nmode)
-        u[torch.arange(self.nmode), torch.arange(self.nmode)] = 1
+        u = matrix.new_ones(self.nmode)
+        u = torch.diag(u)
         u[np.ix_(self.wires, self.wires)] = matrix
         return u
 
@@ -151,8 +151,8 @@ def get_symplectic(self) -> torch.Tensor:
         """Get the global symplectic matrix acting on quadrature operators in xxpp order."""
         matrix, _ = self.update_transform_xp()
         assert matrix.shape[-2] == matrix.shape[-1] == 2 * len(self.wires), 'The matrix may not act on xxpp operators.'
-        s = matrix.new_zeros(2 * self.nmode, 2 * self.nmode)
-        s[torch.arange(2 * self.nmode), torch.arange(2 * self.nmode)] = 1
+        s = matrix.new_ones(2 * self.nmode)
+        s = torch.diag(s)
         wires = self.wires + [wire + self.nmode for wire in self.wires]
         s[np.ix_(wires, wires)] = matrix
         return s