Merge pull request #797 from vprusso/678-infix-vector-mat-mult

Replacing '*' with '@' where appropriate.

docs/intro_tutorial.rst

@@ -93,7 +93,7 @@ In :code:`toqito`, that can be obtained as
     >>> import numpy as np
     >>> e_0, e_1 = basis(2, 0), basis(2, 1)
     >>> u_0 = 1/np.sqrt(2) * (np.kron(e_0, e_0) + np.kron(e_1, e_1))
-    >>> rho_0 = u_0 * u_0.conj().T
+    >>> rho_0 = u_0 @ u_0.conj().T
     >>> rho_0
     array([[0.5, 0. , 0. , 0.5],
            [0. , 0. , 0. , 0. ],
@@ -194,7 +194,7 @@ Any one of the Bell states serve as an example of a pure state
 
     >>> from toqito.states import bell
     >>> from toqito.state_props import is_pure
-    >>> rho = bell(0) * bell(0).conj().T
+    >>> rho = bell(0) @ bell(0).conj().T
     >>> is_pure(rho)
     True
 
@@ -215,7 +215,7 @@ PPT criterion.
 
     >>> from toqito.states import bell
     >>> from toqito.state_props import is_ppt
-    >>> rho = bell(2) * bell(2).conj().T
+    >>> rho = bell(2) @ bell(2).conj().T
     >>> is_ppt(rho)
     False
 
@@ -237,7 +237,7 @@ For instance, the following bound-entangled tile state is found to be entangled
     >>> from toqito.states import tile
     >>> rho = np.identity(9)
     >>> for i in range(5):
-    ...    rho = rho - tile(i) * tile(i).conj().T
+    ...    rho = rho - tile(i) @ tile(i).conj().T
     >>> rho = rho / 4
     >>> is_separable(rho)
     False

docs/tutorials.extended_nonlocal_games.rst

@@ -291,13 +291,13 @@ arrays where :code:`prob_mat` corresponds to the probability distribution
     >>> bb84_pred_mat = np.zeros([dim, dim, a_out, b_out, a_in, b_in])
     >>>
     >>> # V(0,0|0,0) = |0><0|
-    >>> bb84_pred_mat[:, :, 0, 0, 0, 0] = e_0 * e_0.conj().T
+    >>> bb84_pred_mat[:, :, 0, 0, 0, 0] = e_0 @ e_0.conj().T
     >>> # V(1,1|0,0) = |1><1|
-    >>> bb84_pred_mat[:, :, 1, 1, 0, 0] = e_1 * e_1.conj().T
+    >>> bb84_pred_mat[:, :, 1, 1, 0, 0] = e_1 @ e_1.conj().T
     >>> # V(0,0|1,1) = |+><+|
-    >>> bb84_pred_mat[:, :, 0, 0, 1, 1] = e_p * e_p.conj().T
+    >>> bb84_pred_mat[:, :, 0, 0, 1, 1] = e_p @ e_p.conj().T
     >>> # V(1,1|1,1) = |-><-|
-    >>> bb84_pred_mat[:, :, 1, 1, 1, 1] = e_m * e_m.conj().T
+    >>> bb84_pred_mat[:, :, 1, 1, 1, 1] = e_m @ e_m.conj().T
     >>>
     >>> # The probability matrix encode \pi(0,0) = \pi(1,1) = 1/2
     >>> bb84_prob_mat = 1/2*np.identity(2)
@@ -666,21 +666,21 @@ Taking the description of :math:`G_{MUB}`, we can encode this as follows.
     >>> num_out = 3
     >>> pred_mat = np.zeros([dim, dim, num_out, num_out, num_in, num_in], dtype=complex)
     >>>
-    >>> pred_mat[:, :, 0, 0, 0, 0] = mubs[0][0] * mubs[0][0].conj().T
-    >>> pred_mat[:, :, 1, 1, 0, 0] = mubs[0][1] * mubs[0][1].conj().T
-    >>> pred_mat[:, :, 2, 2, 0, 0] = mubs[0][2] * mubs[0][2].conj().T
+    >>> pred_mat[:, :, 0, 0, 0, 0] = mubs[0][0] @ mubs[0][0].conj().T
+    >>> pred_mat[:, :, 1, 1, 0, 0] = mubs[0][1] @ mubs[0][1].conj().T
+    >>> pred_mat[:, :, 2, 2, 0, 0] = mubs[0][2] @ mubs[0][2].conj().T
     >>>
-    >>> pred_mat[:, :, 0, 0, 1, 1] = mubs[1][0] * mubs[1][0].conj().T
-    >>> pred_mat[:, :, 1, 1, 1, 1] = mubs[1][1] * mubs[1][1].conj().T
-    >>> pred_mat[:, :, 2, 2, 1, 1] = mubs[1][2] * mubs[1][2].conj().T
+    >>> pred_mat[:, :, 0, 0, 1, 1] = mubs[1][0] @ mubs[1][0].conj().T
+    >>> pred_mat[:, :, 1, 1, 1, 1] = mubs[1][1] @ mubs[1][1].conj().T
+    >>> pred_mat[:, :, 2, 2, 1, 1] = mubs[1][2] @ mubs[1][2].conj().T
     >>>
-    >>> pred_mat[:, :, 0, 0, 2, 2] = mubs[2][0] * mubs[2][0].conj().T
-    >>> pred_mat[:, :, 1, 1, 2, 2] = mubs[2][1] * mubs[2][1].conj().T
-    >>> pred_mat[:, :, 2, 2, 2, 2] = mubs[2][2] * mubs[2][2].conj().T
+    >>> pred_mat[:, :, 0, 0, 2, 2] = mubs[2][0] @ mubs[2][0].conj().T
+    >>> pred_mat[:, :, 1, 1, 2, 2] = mubs[2][1] @ mubs[2][1].conj().T
+    >>> pred_mat[:, :, 2, 2, 2, 2] = mubs[2][2] @ mubs[2][2].conj().T
     >>>
-    >>> pred_mat[:, :, 0, 0, 3, 3] = mubs[3][0] * mubs[3][0].conj().T
-    >>> pred_mat[:, :, 1, 1, 3, 3] = mubs[3][1] * mubs[3][1].conj().T
-    >>> pred_mat[:, :, 2, 2, 3, 3] = mubs[3][2] * mubs[3][2].conj().T
+    >>> pred_mat[:, :, 0, 0, 3, 3] = mubs[3][0] @ mubs[3][0].conj().T
+    >>> pred_mat[:, :, 1, 1, 3, 3] = mubs[3][1] @ mubs[3][1].conj().T
+    >>> pred_mat[:, :, 2, 2, 3, 3] = mubs[3][2] @ mubs[3][2].conj().T
 
 Now that we have encoded :math:`G_{MUB}`, we can calculate the unentangled value.
 

docs/tutorials.state_distinguishability.rst

@@ -123,8 +123,8 @@ Using :code:`toqito`, we can calculate this probability directly as follows:
     >>>
     >>> # Define the corresponding density matrices of |0> and |1> 
     >>> # given as |0><0| and |1><1|, respectively.
-    >>> e_00 = e_0 * e_0.conj().T
-    >>> e_11 = e_1 * e_1.conj().T
+    >>> e_00 = e_0 @ e_0.conj().T
+    >>> e_11 = e_1 @ e_1.conj().T
     >>>
     >>> # Define a list of states and a corresponding list of 
     >>> # probabilities with which those states are selected.
@@ -222,10 +222,10 @@ via PPT measurements in the following manner.
     >>> x_4 = np.kron(psi_3, psi_3)
     >>>
     >>> # YDY density matrices:
-    >>> rho_1 = x_1 * x_1.conj().T
-    >>> rho_2 = x_2 * x_2.conj().T
-    >>> rho_3 = x_3 * x_3.conj().T
-    >>> rho_4 = x_4 * x_4.conj().T
+    >>> rho_1 = x_1 @ x_1.conj().T
+    >>> rho_2 = x_2 @ x_2.conj().T
+    >>> rho_3 = x_3 @ x_3.conj().T
+    >>> rho_4 = x_4 @ x_4.conj().T
     >>>
     >>> states = [rho_1, rho_2, rho_3, rho_4]
     >>> probs = [1 / 4, 1 / 4, 1 / 4, 1 / 4]

toqito/channel_ops/kraus_to_choi.py

@@ -63,7 +63,7 @@ def kraus_to_choi(kraus_ops: list[list[np.ndarray]], sys: int = 2) -> np.ndarray
     dim_op_1, dim_op_2 = dim_in
 
     choi_mat = partial_channel(
-        max_entangled(dim_op_1, False, False) * max_entangled(dim_op_2, False, False).conj().T,
+        max_entangled(dim_op_1, False, False) @ max_entangled(dim_op_2, False, False).conj().T,
         kraus_ops,
         sys,
         np.array([[dim_op_1, dim_op_1], [dim_op_2, dim_op_2]]),

toqito/channel_ops/partial_channel.py

@@ -181,7 +181,7 @@ def partial_channel(
         psi_c2 = max_entangled(prod_dim_c2, False, False)
 
         phi_map = permute_systems(
-            np.kron(np.kron(psi_r1 * psi_c1.conj().T, phi_map), psi_r2 * psi_c2.conj().T),
+            np.kron(np.kron(psi_r1 @ psi_c1.conj().T, phi_map), psi_r2 @ psi_c2.conj().T),
             [0, 2, 4, 1, 3, 5],
             dim,
         )

toqito/channels/choi.py

@@ -103,4 +103,4 @@ def choi(a_var: int = 1, b_var: int = 1, c_var: int = 0) -> np.ndarray:
 
     """
     psi = max_entangled(3, False, False)
-    return np.diag([a_var + 1, c_var, b_var, b_var, a_var + 1, c_var, c_var, b_var, a_var + 1]) - psi * psi.conj().T
+    return np.diag([a_var + 1, c_var, b_var, b_var, a_var + 1, c_var, c_var, b_var, a_var + 1]) - psi @ psi.conj().T

toqito/channels/dephasing.py

@@ -86,4 +86,4 @@ def dephasing(dim: int, param_p: float = 0) -> np.ndarray:
 
     # Gives a sparse non-normalized state.
     psi = max_entangled(dim=dim, is_sparse=False, is_normalized=False)
-    return (1 - param_p) * np.diag(np.diag(psi * psi.conj().T)) + param_p * (psi * psi.conj().T)
+    return (1 - param_p) * np.diag(np.diag(psi @ psi.conj().T)) + param_p * (psi @ psi.conj().T)

toqito/channels/depolarizing.py

@@ -92,4 +92,4 @@ def depolarizing(dim: int, param_p: float = 0) -> np.ndarray:
 
     # Gives a sparse non-normalized state.
     psi = max_entangled(dim=dim, is_sparse=False, is_normalized=False)
-    return (1 - param_p) * np.identity(dim**2) / dim + param_p * (psi * psi.conj().T)
+    return (1 - param_p) * np.identity(dim**2) / dim + param_p * (psi @ psi.conj().T)

toqito/channels/tests/test_partial_transpose.py

@@ -688,7 +688,7 @@ def test_partial_transpose_16_by_16_subsystems_2_2_2_2():
 
 def test_partial_transpose_bell_state():
     """Test partial transpose on a Bell state."""
-    rho = bell(2) * bell(2).conj().T
+    rho = bell(2) @ bell(2).conj().T
     expected_res = np.array([[0, 0, 0, 1 / 2], [0, 1 / 2, 0, 0], [0, 0, 1 / 2, 0], [1 / 2, 0, 0, 0]])
     res = partial_transpose(rho)
     np.testing.assert_equal(np.allclose(res, expected_res), True)

toqito/matrix_props/is_density.py

@@ -36,7 +36,7 @@ def is_density(mat: np.ndarray) -> bool:
     >>> from toqito.matrix_props import is_density
     >>> from toqito.states import bell
     >>> import numpy as np
-    >>> rho = bell(0) * bell(0).conj().T
+    >>> rho = bell(0) @ bell(0).conj().T
     >>> is_density(rho)
     np.True_
 

toqito/matrix_props/is_unitary.py

@@ -83,5 +83,5 @@ def is_unitary(mat: np.ndarray, rtol: float = 1e-05, atol: float = 1e-08) -> boo
     u_uc_mat = mat @ mat.conj().T
     id_mat = np.eye(len(mat))
 
-    # If U^* * U = I U * U^*, the matrix "U" is unitary.
+    # If U^* @ U = I U @ U^*, the matrix "U" is unitary.
     return np.allclose(uc_u_mat, id_mat, rtol=rtol, atol=atol) and np.allclose(u_uc_mat, id_mat, rtol=rtol, atol=atol)

toqito/matrix_props/majorizes.py

@@ -36,7 +36,7 @@ def majorizes(a_var: np.ndarray | list[int], b_var: np.ndarray | list[int]) -> b
     >>> from toqito.channels import partial_trace
     >>>
     >>> v_vec = max_entangled(3)
-    >>> rho = v_vec * v_vec.conj().T
+    >>> rho = v_vec @ v_vec.conj().T
     >>> majorizes(partial_trace(rho, [1]), rho)
     False
 

toqito/matrix_props/tests/test_majorizes.py

@@ -15,12 +15,12 @@ def test_majorizes_simple_example():
 def test_majorizes_max_entangled():
     """Test that max entangled partial trace returns False."""
     v_vec = max_entangled(3)
-    rho = v_vec * v_vec.conj().T
+    rho = v_vec @ v_vec.conj().T
     np.testing.assert_equal(majorizes(partial_trace(rho, [1], [3, 3]), rho), False)
 
 
 def test_majorizes_max_entangled_flip():
     """Test that max entangled partial trace returns True (flipped args)."""
     v_vec = max_entangled(3)
-    rho = v_vec * v_vec.conj().T
+    rho = v_vec @ v_vec.conj().T
     np.testing.assert_equal(majorizes(rho, partial_trace(rho, [1], [3, 3])), True)

toqito/matrix_props/tests/test_trace_norm.py

@@ -13,7 +13,7 @@ def test_trace_norm():
     e_11 = np.kron(e_1, e_1)
 
     u_vec = 1 / np.sqrt(2) * (e_00 + e_11)
-    rho = u_vec * u_vec.conj().T
+    rho = u_vec @ u_vec.conj().T
 
     res = trace_norm(rho)
     _, singular_vals, _ = np.linalg.svd(rho)

toqito/matrix_props/trace_norm.py

@@ -34,7 +34,7 @@ def trace_norm(rho: np.ndarray) -> float:
 
     >>> from toqito.states import bell
     >>> from toqito.matrix_props import trace_norm
-    >>> rho = bell(0) * bell(0).conj().T
+    >>> rho = bell(0) @ bell(0).conj().T
     >>> trace_norm(rho)
     np.float64(0.9999999999999999)
 

toqito/measurement_ops/measure.py

@@ -40,10 +40,10 @@ def measure(measurement: np.ndarray, state: np.ndarray) -> float:
     >>> e_0, e_1 = basis(2, 0), basis(2, 1)
     >>>
     >>> u = 1/np.sqrt(3) * e_0 + np.sqrt(2/3) * e_1
-    >>> rho = u * u.conj().T
+    >>> rho = u @ u.conj().T
     >>>
-    >>> proj_0 = e_0 * e_0.conj().T
-    >>> proj_1 = e_1 * e_1.conj().T
+    >>> proj_0 = e_0 @ e_0.conj().T
+    >>> proj_1 = e_1 @ e_1.conj().T
 
     Then the probability of obtaining outcome :math:`0` is given by
 
@@ -73,4 +73,4 @@ def measure(measurement: np.ndarray, state: np.ndarray) -> float:
              the variable :code:`measurement` to the variable :code:`state`.
 
     """
-    return float(np.trace(measurement.conj().T * state))
+    return float(np.trace(measurement.conj().T @ state))

toqito/measurement_ops/tests/test_measure.py

@@ -10,9 +10,9 @@ def test_measure_state():
     """Test measure on quantum state."""
     e_0, e_1 = basis(2, 0), basis(2, 1)
     psi = 1 / np.sqrt(3) * e_0 + np.sqrt(2 / 3) * e_1
-    rho = psi * psi.conj().T
+    rho = psi @ psi.conj().T
 
-    proj_0 = e_0 * e_0.conj().T
-    proj_1 = e_1 * e_1.conj().T
+    proj_0 = e_0 @ e_0.conj().T
+    proj_1 = e_1 @ e_1.conj().T
     np.testing.assert_equal(np.isclose(measure(proj_0, rho), 1 / 3), True)
     np.testing.assert_equal(np.isclose(measure(proj_1, rho), 2 / 3), True)

toqito/nonlocal_games/extended_nonlocal_game.py

@@ -112,7 +112,7 @@ def unentangled_value(self) -> float:
 
                 rho = cvxpy.Variable((dim_x, dim_y), hermitian=True)
 
-                objective = cvxpy.Maximize(cvxpy.real(cvxpy.trace(p_win.conj().T @ rho)))
+                objective = cvxpy.Maximize(cvxpy.real(cvxpy.trace(p_win.conj().T * rho)))
 
                 constraints = [cvxpy.trace(rho) == 1, rho >> 0]
                 problem = cvxpy.Problem(objective, constraints)
@@ -194,7 +194,7 @@ def nonsignaling_value(self) -> float:
                 for x_in in range(alice_in):
                     for y_in in range(bob_in):
                         p_win += self.prob_mat[x_in, y_in] * cvxpy.trace(
-                            self.pred_mat[:, :, a_out, b_out, x_in, y_in].conj().T @ k_var[a_out, b_out, x_in, y_in]
+                            self.pred_mat[:, :, a_out, b_out, x_in, y_in].conj().T * k_var[a_out, b_out, x_in, y_in]
                         )
 
         objective = cvxpy.Maximize(cvxpy.real(p_win))
@@ -328,7 +328,7 @@ def __optimize_alice(self, bob_povms) -> tuple[dict, float]:
                                 )
                                 .conj()
                                 .T
-                                @ rho[x_ques, a_ans]
+                                * rho[x_ques, a_ans]
                             )
                         if isinstance(
                             bob_povms[y_ques, b_ans],
@@ -343,7 +343,7 @@ def __optimize_alice(self, bob_povms) -> tuple[dict, float]:
                                 )
                                 .conj()
                                 .T
-                                @ rho[x_ques, a_ans]
+                                * rho[x_ques, a_ans]
                             )
         objective = cvxpy.Maximize(cvxpy.real(win))
         constraints = []
@@ -396,7 +396,7 @@ def __optimize_bob(self, rho) -> tuple[dict, float]:
                                     bob_povms[y_ques, b_ans],
                                 )
                             )
-                            @ rho[x_ques, a_ans].value
+                            * rho[x_ques, a_ans].value
                         )
         objective = cvxpy.Maximize(cvxpy.real(win))
 
@@ -455,7 +455,7 @@ def commuting_measurement_value_upper_bound(self, k: int | str = 1) -> float:
                     for y_in in range(bob_in):
                         p_win += self.prob_mat[x_in, y_in] * cvxpy.trace(
                             self.pred_mat[:, :, a_out, b_out, x_in, y_in].conj().T
-                            @ mat[x_in, y_in][
+                            * mat[x_in, y_in][
                                 a_out * referee_dim : (a_out + 1) * referee_dim,
                                 b_out * referee_dim : (b_out + 1) * referee_dim,
                             ]

toqito/nonlocal_games/nonlocal_game.py

@@ -372,7 +372,7 @@ def __optimize_alice(self, dim, bob_povms) -> tuple[dict, float]:
                             win += (
                                 self.prob_mat[x_ques, y_ques]
                                 * self.pred_mat[a_ans, b_ans, x_ques, y_ques]
-                                * cvxpy.trace(bob_povms[y_ques, b_ans].conj().T @ alice_povms[x_ques, a_ans])
+                                * cvxpy.trace(bob_povms[y_ques, b_ans].conj().T * alice_povms[x_ques, a_ans])
                             )
                         if isinstance(
                             bob_povms[y_ques, b_ans],
@@ -382,7 +382,7 @@ def __optimize_alice(self, dim, bob_povms) -> tuple[dict, float]:
                             win += (
                                 self.prob_mat[x_ques, y_ques]
                                 * self.pred_mat[a_ans, b_ans, x_ques, y_ques]
-                                * cvxpy.trace(bob_povms[y_ques, b_ans].value.conj().T @ alice_povms[x_ques, a_ans])
+                                * cvxpy.trace(bob_povms[y_ques, b_ans].value.conj().T * alice_povms[x_ques, a_ans])
                             )
 
         objective = cvxpy.Maximize(cvxpy.real(win))
@@ -430,7 +430,7 @@ def __optimize_bob(self, dim, alice_povms) -> tuple[dict, float]:
                         win += (
                             self.prob_mat[x_ques, y_ques]
                             * self.pred_mat[a_ans, b_ans, x_ques, y_ques]
-                            * cvxpy.trace(bob_povms[y_ques, b_ans].H @ alice_povms[x_ques, a_ans].value)
+                            * cvxpy.trace(bob_povms[y_ques, b_ans].H * alice_povms[x_ques, a_ans].value)
                         )
 
         objective = cvxpy.Maximize(cvxpy.real(win))

toqito/nonlocal_games/quantum_hedging.py

@@ -48,8 +48,8 @@ class QuantumHedging:
     >>> l_2 = alpha * sin(theta) * e_10
     >>> l_3 = sqrt(1 - alpha ** 2) * cos(theta) * e_01
     >>>
-    >>> q_1 = w_var * w_var.conj().T
-    >>> q_0 = l_1 * l_1.conj().T + l_2 * l_2.conj().T + l_3 * l_3.conj().T
+    >>> q_1 = w_var @ w_var.conj().T
+    >>> q_0 = l_1 @ l_1.conj().T + l_2 @ l_2.conj().T + l_3 @ l_3.conj().T
     >>> molina_watrous = QuantumHedging(q_0, 1)
     >>>
     >>> # cos(pi/8)**2 \approx 0.8536

toqito/nonlocal_games/tests/test_extended_nonlocal_game.py

@@ -23,10 +23,10 @@ def bb84_extended_nonlocal_game():
         num_alice_in, num_bob_in = 2, 2
 
         pred_mat = np.zeros([dim, dim, num_alice_out, num_bob_out, num_alice_in, num_bob_in])
-        pred_mat[:, :, 0, 0, 0, 0] = e_0 * e_0.conj().T
-        pred_mat[:, :, 0, 0, 1, 1] = e_p * e_p.conj().T
-        pred_mat[:, :, 1, 1, 0, 0] = e_1 * e_1.conj().T
-        pred_mat[:, :, 1, 1, 1, 1] = e_m * e_m.conj().T
+        pred_mat[:, :, 0, 0, 0, 0] = e_0 @ e_0.conj().T
+        pred_mat[:, :, 0, 0, 1, 1] = e_p @ e_p.conj().T
+        pred_mat[:, :, 1, 1, 0, 0] = e_1 @ e_1.conj().T
+        pred_mat[:, :, 1, 1, 1, 1] = e_m @ e_m.conj().T
 
         prob_mat = 1 / 2 * np.identity(2)
 
@@ -88,21 +88,21 @@ def moe_mub_4_in_3_out_game():
         num_out = 3
         pred_mat = np.zeros([dim, dim, num_out, num_out, num_in, num_in], dtype=complex)
 
-        pred_mat[:, :, 0, 0, 0, 0] = mubs[0][0] * mubs[0][0].conj().T
-        pred_mat[:, :, 1, 1, 0, 0] = mubs[0][1] * mubs[0][1].conj().T
-        pred_mat[:, :, 2, 2, 0, 0] = mubs[0][2] * mubs[0][2].conj().T
+        pred_mat[:, :, 0, 0, 0, 0] = mubs[0][0] @ mubs[0][0].conj().T
+        pred_mat[:, :, 1, 1, 0, 0] = mubs[0][1] @ mubs[0][1].conj().T
+        pred_mat[:, :, 2, 2, 0, 0] = mubs[0][2] @ mubs[0][2].conj().T
 
-        pred_mat[:, :, 0, 0, 1, 1] = mubs[1][0] * mubs[1][0].conj().T
-        pred_mat[:, :, 1, 1, 1, 1] = mubs[1][1] * mubs[1][1].conj().T
-        pred_mat[:, :, 2, 2, 1, 1] = mubs[1][2] * mubs[1][2].conj().T
+        pred_mat[:, :, 0, 0, 1, 1] = mubs[1][0] @ mubs[1][0].conj().T
+        pred_mat[:, :, 1, 1, 1, 1] = mubs[1][1] @ mubs[1][1].conj().T
+        pred_mat[:, :, 2, 2, 1, 1] = mubs[1][2] @ mubs[1][2].conj().T
 
-        pred_mat[:, :, 0, 0, 2, 2] = mubs[2][0] * mubs[2][0].conj().T
-        pred_mat[:, :, 1, 1, 2, 2] = mubs[2][1] * mubs[2][1].conj().T
-        pred_mat[:, :, 2, 2, 2, 2] = mubs[2][2] * mubs[2][2].conj().T
+        pred_mat[:, :, 0, 0, 2, 2] = mubs[2][0] @ mubs[2][0].conj().T
+        pred_mat[:, :, 1, 1, 2, 2] = mubs[2][1] @ mubs[2][1].conj().T
+        pred_mat[:, :, 2, 2, 2, 2] = mubs[2][2] @ mubs[2][2].conj().T
 
-        pred_mat[:, :, 0, 0, 3, 3] = mubs[3][0] * mubs[3][0].conj().T
-        pred_mat[:, :, 1, 1, 3, 3] = mubs[3][1] * mubs[3][1].conj().T
-        pred_mat[:, :, 2, 2, 3, 3] = mubs[3][2] * mubs[3][2].conj().T
+        pred_mat[:, :, 0, 0, 3, 3] = mubs[3][0] @ mubs[3][0].conj().T
+        pred_mat[:, :, 1, 1, 3, 3] = mubs[3][1] @ mubs[3][1].conj().T
+        pred_mat[:, :, 2, 2, 3, 3] = mubs[3][2] @ mubs[3][2].conj().T
 
         return prob_mat, pred_mat