def get_unitary_and_grad(
self,
params: Sequence[float] = [],
) -> tuple[UnitaryMatrix, np.ndarray]:
"""Returns the unitary and gradient for this gate."""
self.check_parameters(params)
a, l = self.split_params(params)
l = softmax(l, 10)
P = np.sum([a * s.get_numpy() for a, s in zip(l, self.perms)], 0)
G = self.gate.get_unitary(a).get_numpy() # type: ignore
G = np.kron(G, self.I)
PG = P @ G
GPT = G @ P.T
PGPT = P @ GPT
dG = self.gate.get_grad(a) # type: ignore
dG = np.kron(dG, self.I)
dG = P @ dG @ P.T
perm_array = np.array([perm.get_numpy() for perm in self.perms])
dP = perm_array @ GPT + PG @ perm_array.transpose((0, 2, 1)) - 2 * PGPT
dP = np.array([10 * x * y for x, y in zip(l, dP)])
U = UnitaryMatrix.closest_to(PGPT, self.get_radixes())
return U, np.concatenate([dG, dP])
def get_unitary(self, params: Sequence[float] = []) -> UnitaryMatrix:
"""Returns the unitary for this gate, see Unitary for more info."""
self.check_parameters(params)
a, l = self.split_params(params)
l = softmax(l, 10)
P = np.sum([a * s.get_numpy() for a, s in zip(l, self.perms)], 0)
G = self.gate.get_unitary(a) # type: ignore
# TODO: Change get_unitary params to be union with np.ndarray
PGPT = P @ np.kron(G.get_numpy(), self.I) @ P.T
return UnitaryMatrix.closest_to(PGPT, self.get_radixes())
def get_unitary(self, params: Sequence[float] = []) -> UnitaryMatrix:
"""
Returns the unitary for this gate, see Unitary for more info.
Note:
Ideally, params form a unitary matrix when reshaped,
however, params are unconstrained so we return the closest
UnitaryMatrix to the given matrix.
"""
self.check_parameters(params)
mid = len(params) // 2
real = np.array(params[:mid], dtype=np.complex128)
imag = 1j * np.array(params[mid:], dtype=np.complex128)
x = real + imag
return UnitaryMatrix.closest_to(np.reshape(x, self.shape))