diff --git a/src/deepquantum/photonic/__init__.py b/src/deepquantum/photonic/__init__.py index 69f95193..9f5a52f3 100644 --- a/src/deepquantum/photonic/__init__.py +++ b/src/deepquantum/photonic/__init__.py @@ -9,6 +9,7 @@ draw, gate, hafnian_, + kensingtonian_, mapper, measurement, operation, @@ -47,6 +48,7 @@ UAnyGate, ) from .hafnian_ import hafnian +from .kensingtonian_ import kensingtonian from .mapper import UnitaryMapper from .measurement import GeneralBosonic, Generaldyne, Homodyne, PhotonNumberResolvingBosonic from .qmath import ( diff --git a/src/deepquantum/photonic/circuit.py b/src/deepquantum/photonic/circuit.py index 23788eb1..16d03aa6 100644 --- a/src/deepquantum/photonic/circuit.py +++ b/src/deepquantum/photonic/circuit.py @@ -43,6 +43,7 @@ UAnyGate, ) from .hafnian_ import hafnian +from .kensingtonian_ import kensingtonian from .measurement import Generaldyne, Homodyne from .operation import Channel, Delay, Gate, Operation from .qmath import ( @@ -1077,10 +1078,14 @@ def _get_prob_fock_vmap(self, sub_mat: torch.Tensor, per_norm: torch.Tensor) -> prob = torch.abs(amplitude) ** 2 return prob - def _get_prob_gaussian(self, final_state: Any, state: Any = None) -> torch.Tensor: + def _get_prob_gaussian( + self, final_state: Any, state: Any = None, detector: str | None = None, num_detectors: int = 1 + ) -> torch.Tensor: """Get the batched probabilities of the final state for Gaussian backend.""" if not isinstance(final_state, torch.Tensor): final_state = torch.tensor(final_state, dtype=torch.long) + detector = self.detector if detector is None else detector + num_detectors = getattr(self, '_num_detectors', num_detectors) if state is None: cov = self._cov mean = self._mean @@ -1097,7 +1102,9 @@ def _get_prob_gaussian(self, final_state: Any, state: Any = None) -> torch.Tenso batch = cov.shape[0] probs = [] for i in range(batch): - prob = self._get_probs_gaussian_helper(final_state, cov=cov[i], mean=mean[i], detector=self.detector)[0] + prob = self._get_probs_gaussian_helper( + final_state, cov=cov[i], mean=mean[i], detector=detector, num_detectors=num_detectors + )[0] probs.append(prob) return torch.stack(probs).squeeze() @@ -1109,6 +1116,7 @@ def _get_probs_gaussian_helper( detector: str = 'pnrd', purity: bool | None = None, loop: bool | None = None, + num_detectors: int = 1, ) -> torch.Tensor: """Get the probabilities of the final states for Gaussian backend.""" if loop is None: @@ -1129,13 +1137,32 @@ def _get_probs_gaussian_helper( gamma = mean_ladder.mH @ q.inverse() if detector == 'pnrd': matrix = a_mat - elif detector == 'threshold': + elif detector in ('threshold', 'click_counting'): matrix = o_mat + else: + raise ValueError(f'Unsupported detector: {detector}') if purity is None: purity = GaussianState([cov, mean]).is_pure p_vac = torch.exp(-0.5 * mean_ladder.mH @ q.inverse() @ mean_ladder) / det_q.sqrt() - batch_get_prob = vmap(self._get_prob_gaussian_base, in_dims=(0, None, None, None, None, None, None)) - probs = batch_get_prob(final_states, matrix, gamma, p_vac, detector, purity, loop) + if detector == 'click_counting': + probs = torch.stack( + [ + self._get_prob_gaussian_base( + final_state, + matrix, + gamma, + p_vac, + detector, + purity, + loop, + num_detectors, + ) + for final_state in final_states + ] + ) + else: + batch_get_prob = vmap(self._get_prob_gaussian_base, in_dims=(0, None, None, None, None, None, None, None)) + probs = batch_get_prob(final_states, matrix, gamma, p_vac, detector, purity, loop, num_detectors) return probs def _get_prob_gaussian_base( @@ -1147,6 +1174,7 @@ def _get_prob_gaussian_base( detector: str = 'pnrd', purity: bool = True, loop: bool = False, + num_detectors: int = 1, ) -> torch.Tensor: """Get the probability of the final state for Gaussian backend.""" gamma = gamma.squeeze() @@ -1175,6 +1203,9 @@ def _get_prob_gaussian_base( final_state_double = torch.cat([final_state, final_state]) sub_mat = sub_matrix(matrix, final_state_double, final_state_double) prob = p_vac * torontonian(sub_mat, sub_gamma) + elif detector == 'click_counting': + loop_gamma = gamma if loop else None + prob = p_vac * kensingtonian(matrix, final_state, num_detectors, gamma=loop_gamma) return abs(prob.real).squeeze() def _get_prob_mps(self, final_state: Any, wires: int | list[int] | None = None) -> torch.Tensor: @@ -1203,6 +1234,7 @@ def measure( wires: int | list[int] | None = None, detector: str | None = None, mcmc: bool = False, + num_detectors: int = 1, ) -> dict | list[dict] | None: """Measure the final state. @@ -1215,6 +1247,7 @@ def measure( detector: For Gaussian backend, use ``'pnrd'`` for the photon-number-resolving detector or ``'threshold'`` for the threshold detector. Default: ``None`` mcmc: Whether to use MCMC sampling method. Default: ``False`` + num_detectors: The number of threshold detectors in each click-counting detector. Default: 1 See https://arxiv.org/pdf/2108.01622 for MCMC. """ @@ -1228,7 +1261,7 @@ def measure( results = self._measure_fock(shots, with_prob, wires, mcmc) elif self.backend == 'gaussian': detector = self.detector if detector is None else detector.lower() - results = self._measure_gaussian(shots, with_prob, wires, detector, mcmc) + results = self._measure_gaussian(shots, with_prob, wires, detector, mcmc, num_detectors) if len(results) == 1: results = results[0] return results @@ -1428,10 +1461,18 @@ def _sample_mcmc_fock(self, shots: int, init_state: torch.Tensor, unitary: torch ) return merged_samples - def _measure_gaussian(self, shots: int, with_prob: bool, wires: list[int], detector: str, mcmc: bool) -> list[dict]: + def _measure_gaussian( + self, + shots: int, + with_prob: bool, + wires: list[int], + detector: str, + mcmc: bool, + num_detectors: int = 1, + ) -> list[dict]: """Measure the final state for Gaussian backend.""" if isinstance(self.state, list): - return self._measure_gaussian_state(shots, with_prob, wires, detector, mcmc) + return self._measure_gaussian_state(shots, with_prob, wires, detector, mcmc, num_detectors) elif isinstance(self.state, dict): assert not mcmc, "Final states have been calculated, we don't need mcmc!" print('Automatically using the default detector!') @@ -1440,7 +1481,13 @@ def _measure_gaussian(self, shots: int, with_prob: bool, wires: list[int], detec raise ValueError('Check your forward function or input!') def _measure_gaussian_state( - self, shots: int, with_prob: bool, wires: list[int], detector: str, mcmc: bool + self, + shots: int, + with_prob: bool, + wires: list[int], + detector: str, + mcmc: bool, + num_detectors: int = 1, ) -> list[dict]: """Measure the final state according to Gaussian state for Gaussian backend. @@ -1455,14 +1502,19 @@ def _measure_gaussian_state( print('Using MCMC method to sample the final states!') for i in range(batch): samples_i = self._sample_mcmc_gaussian( - shots=shots, cov=cov[i], mean=mean[i], detector=detector, num_chain=5 + shots=shots, + cov=cov[i], + mean=mean[i], + detector=detector, + num_chain=5, + num_detectors=num_detectors, ) all_samples.append(samples_i) else: # chain-rule method with small number of shots print('Using chain-rule method to sample the final states!') samples = [] for _ in range(shots): - sample = self._generate_chain_sample(wires, detector) + sample = self._generate_chain_sample(wires, detector, num_detectors) samples.append(sample) samples = torch.stack(samples).permute(1, 0, 2) # (batch, shots, wires) for i in range(batch): @@ -1475,11 +1527,16 @@ def _measure_gaussian_state( if with_prob: for k in samples_i: if mcmc: - prob = self._get_prob_gaussian(k, [cov[i], mean[i]]) + prob = self._get_prob_gaussian(k, [cov[i], mean[i]], num_detectors=num_detectors) else: wires_ = torch.tensor(wires, device=cov.device) idx = torch.cat([wires_, wires_ + self.nmode]) - prob = self._get_prob_gaussian(k, [cov[i][idx[:, None], idx], mean[i][idx, :]]) + prob = self._get_prob_gaussian( + k, + [cov[i][idx[:, None], idx], mean[i][idx, :]], + detector=detector, + num_detectors=num_detectors, + ) samples_i[k] = samples_i[k], prob for key in samples_i: state_b = [key[wire] for wire in wires] if mcmc else list(key) @@ -1495,11 +1552,20 @@ def _measure_gaussian_state( all_results.append(results) return all_results - def _sample_mcmc_gaussian(self, shots: int, cov: torch.Tensor, mean: torch.Tensor, detector: str, num_chain: int): + def _sample_mcmc_gaussian( + self, + shots: int, + cov: torch.Tensor, + mean: torch.Tensor, + detector: str, + num_chain: int, + num_detectors: int = 1, + ): """Sample the output states for Gaussian backend via SC-MCMC method.""" self._cov = cov self._mean = mean self.detector = detector + self._num_detectors = num_detectors if detector == 'threshold' and not torch.allclose(mean, torch.zeros_like(mean)): # For the displaced state, aggregate PNRD detector samples to derive threshold detector results self.detector = 'pnrd' @@ -1530,25 +1596,31 @@ def _proposal_sampler(self): assert self.basis, 'Currently NOT supported.' sample = self._out_fock_basis[torch.randint(0, len(self._out_fock_basis), (1,))[0]] elif self.backend == 'gaussian': - sample = self._generate_rand_sample(self.detector) + num_detectors = getattr(self, '_num_detectors', 1) + sample = self._generate_rand_sample(self.detector, num_detectors) return tuple(sample.tolist()) - def _generate_rand_sample(self, detector: str = 'pnrd'): + def _generate_rand_sample(self, detector: str = 'pnrd', num_detectors: int = 1): """Generate a random sample according to uniform proposal distribution.""" nmode = self._nmode_tdm if self._with_delay else self.nmode if detector == 'threshold': sample = torch.randint(0, 2, [nmode]) + elif detector == 'click_counting': + sample = torch.randint(0, num_detectors + 1, [nmode]) elif detector == 'pnrd': sample = torch.randint(0, self.cutoff, [nmode]) + else: + raise ValueError(f'Unsupported detector: {detector}') return sample - def _generate_chain_sample(self, wires: list[int], detector: str) -> torch.Tensor: + def _generate_chain_sample(self, wires: list[int], detector: str, num_detectors: int = 1) -> torch.Tensor: """Generate batched random samples via chain rule. Args: wires: The wires to measure. It can be a list of integers specifying the indices of the wires. - detector: For Gaussian backend, use ``'pnrd'`` for the photon-number-resolving detector or - ``'threshold'`` for the threshold detector. + detector: For Gaussian backend, use ``'pnrd'`` for the photon-number-resolving detector, + ``'threshold'`` for the threshold detector, or ``'click_counting'`` for the click-counting detector. + num_detectors: The number of threshold detectors in each click-counting detector. Default: 1 Returns: Tensor of shape (batch, nwire). @@ -1567,10 +1639,10 @@ def _generate_chain_sample(self, wires: list[int], detector: str) -> torch.Tenso mps[i] = torch.gather(mps[i], dim=2, index=index) sample = torch.stack(sample, dim=-1).squeeze(1) elif self.backend == 'gaussian': # chain rule for GBS - sample = self._generate_chain_sample_gaussian(wires, detector) + sample = self._generate_chain_sample_gaussian(wires, detector, num_detectors) return sample - def _generate_chain_sample_gaussian(self, wires: list[int], detector: str) -> torch.Tensor: + def _generate_chain_sample_gaussian(self, wires: list[int], detector: str, num_detectors: int = 1) -> torch.Tensor: """Generate batched random samples via chain rule for Gaussian backend. See @@ -1581,7 +1653,16 @@ def _generate_chain_sample_gaussian(self, wires: list[int], detector: str) -> to def _sample_wire(sample, cov_sub, mean_sub, cutoff, detector): """Sample for a wire.""" states = [torch.tensor(sample + [i], device=cov_sub.device) for i in range(cutoff)] - probs = [self._get_probs_gaussian_helper(s, cov_sub, mean_sub, detector) for s in states] + probs = [ + self._get_probs_gaussian_helper( + s, + cov_sub, + mean_sub, + detector, + num_detectors=num_detectors, + ) + for s in states + ] sample_wire = torch.multinomial(torch.cat(probs), num_samples=1) return sample_wire @@ -1633,7 +1714,14 @@ def _sample_mixed(cov, mean, wires, nmode, cutoff, detector, eps=5e-5): purity = GaussianState(self.state).is_pure cov, mean = self.state batch = cov.shape[0] - cutoff = 2 if detector == 'threshold' else self.cutoff + if detector == 'threshold': + cutoff = 2 + elif detector == 'click_counting': + cutoff = num_detectors + 1 + elif detector == 'pnrd': + cutoff = self.cutoff + else: + raise ValueError(f'Unsupported detector: {detector}') if purity: for i in range(batch): sample.append(_sample_pure(cov[i], mean[i], wires, self.nmode, cutoff, detector)) diff --git a/src/deepquantum/photonic/kensingtonian_.py b/src/deepquantum/photonic/kensingtonian_.py new file mode 100644 index 00000000..a7654a70 --- /dev/null +++ b/src/deepquantum/photonic/kensingtonian_.py @@ -0,0 +1,209 @@ +"""Functions for real- and complex-representation Kensingtonians.""" + +from collections import defaultdict +from functools import cache +from itertools import product +from math import comb + +import torch + + +@cache +def _d_term_data( + clicks: tuple[int, ...], num_detectors: int +) -> tuple[tuple[tuple[int, ...], tuple[int, ...], tuple[float, ...], float], ...]: + """Return matrix-independent data for all mixed-radix ``d`` terms.""" + size = 2 * len(clicks) + ranges = [range(click + 1) for click in clicks] + data = [] + for d_tuple in product(*ranges): + sign = -1 if sum(click - d for click, d in zip(clicks, d_tuple, strict=True)) % 2 else 1 + coeff = 1 + for click, d in zip(clicks, d_tuple, strict=True): + coeff *= comb(num_detectors, click) * comb(click, d) + keep_modes = tuple(i for i, d in enumerate(d_tuple) if d > 0) + prefactor = 1.0 + for i in keep_modes: + prefactor *= num_detectors / d_tuple[i] + idx_tuple = tuple(sorted(keep_modes + tuple(i + size // 2 for i in keep_modes))) + dz_values = tuple((num_detectors - d_tuple[i]) / d_tuple[i] for i in keep_modes) + data.append((d_tuple, idx_tuple, dz_values + dz_values, sign * coeff * prefactor)) + return tuple(data) + + +@cache +def _grouped_term_data( + clicks: tuple[int, ...], num_detectors: int +) -> tuple[ + tuple[ + int, + tuple[tuple[tuple[int, ...], tuple[float, ...], float], ...], + ], + ..., +]: + """Group terms by submatrix size for batched linear algebra.""" + groups = defaultdict(list) + for _, idx, dz_values, scalar in _d_term_data(clicks, num_detectors): + groups[len(idx)].append((idx, dz_values, scalar)) + return tuple((size, tuple(groups[size])) for size in sorted(groups)) + + +def _kensingtonian_same_clicks( + matrix: torch.Tensor, + clicks: tuple[int, ...], + num_detectors: int, + gamma: torch.Tensor | None, + chunk_size: int | None, +) -> torch.Tensor: + """Evaluate a matrix batch whose samples share one click pattern.""" + batch_size, size, _ = matrix.shape + identity = torch.eye(size, dtype=matrix.dtype, device=matrix.device) + sigma_inv = identity - matrix + + result = matrix.new_zeros(batch_size) + for submatrix_size, terms in _grouped_term_data(clicks, num_detectors): + if submatrix_size == 0: + result = result + sum(term[2] for term in terms) + continue + + group_chunk_size = len(terms) if chunk_size is None else chunk_size + for start in range(0, len(terms), group_chunk_size): + term_chunk = terms[start : start + group_chunk_size] + idx = torch.tensor( + [term[0] for term in term_chunk], + dtype=torch.long, + device=matrix.device, + ) + dz_values = torch.tensor( + [term[1] for term in term_chunk], + dtype=matrix.dtype, + device=matrix.device, + ) + scalars = torch.tensor( + [term[2] for term in term_chunk], + dtype=matrix.dtype, + device=matrix.device, + ) + + det_mat = sigma_inv[:, idx[:, :, None], idx[:, None, :]] + det_mat.diagonal(dim1=-2, dim2=-1).add_(dz_values.unsqueeze(0)) + det_term = torch.linalg.det(det_mat) + + if gamma is None: + numerator = scalars.unsqueeze(0) + else: + rhs = gamma[:, idx] + solution = torch.linalg.solve(det_mat, rhs.conj().unsqueeze(-1)).squeeze(-1) + quad = (rhs * solution).sum(dim=-1) / 2 + numerator = torch.exp(quad) * scalars.unsqueeze(0) + + result = result + (numerator / torch.sqrt(det_term)).sum(dim=-1) + return result + + +def kensingtonian( + matrix: torch.Tensor, + clicks: torch.Tensor | list[int], + num_detectors: int, + gamma: torch.Tensor | None = None, + chunk_size: int | None = None, +) -> torch.Tensor: + r"""Calculate the Kensingtonian or loop Kensingtonian. + + See https://arxiv.org/abs/2305.00853 Eq. (27)-(29) and Appendix C. + + Args: + matrix: The input matrix :math:`O=I-\Sigma^{-1}` in either + DeepQuantum's real ``xxpp`` ordering or complex ladder ordering. + clicks: The click-counting pattern :math:`k`. + num_detectors: The number of threshold detectors in each + click-counting detector. + gamma: The loop vector in the same convention as + :func:`torontonian`, namely + :math:`\gamma=\sqrt{2}\alpha^T\Sigma^{-1}`. + Default: ``None`` + chunk_size: The maximum number of ``d`` terms evaluated together + within one same-size group. Default: ``None`` + """ + assert matrix.dim() == 2, 'Input matrix should be 2D' + assert matrix.shape[-2] == matrix.shape[-1], 'Input matrix should be square' + assert matrix.shape[-1] % 2 == 0, 'Input matrix dimension should be even' + assert num_detectors >= 1, 'num_detectors should be positive' + assert chunk_size is None or chunk_size > 0, 'chunk_size should be positive' + + nmode = matrix.shape[-1] // 2 + clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device) + clicks = clicks.reshape(-1) + assert clicks.shape == (nmode,), f'Click pattern should have shape ({nmode},)' + assert (clicks >= 0).all(), 'Click numbers should be non-negative' + assert (clicks <= num_detectors).all(), 'Click numbers should not exceed num_detectors' + + if gamma is not None: + gamma = torch.as_tensor(gamma, device=matrix.device) + assert matrix.is_complex() or not gamma.is_complex(), 'Complex gamma requires a complex matrix' + gamma = gamma.to(matrix.dtype).reshape(-1) + assert gamma.shape == (2 * nmode,), f'gamma should have shape ({2 * nmode},)' + + clicks_tuple = tuple(clicks.tolist()) + return _kensingtonian_same_clicks( + matrix.unsqueeze(0), + clicks_tuple, + num_detectors, + None if gamma is None else gamma.unsqueeze(0), + chunk_size, + )[0] + + +def kensingtonian_batch( + matrix: torch.Tensor, + clicks: torch.Tensor | list[int], + num_detectors: int, + gamma: torch.Tensor | None = None, + chunk_size: int | None = None, +) -> torch.Tensor: + """Calculate batched Kensingtonians for one or more click patterns.""" + assert matrix.dim() == 3, 'Input tensor should be in batched size' + assert matrix.shape[-2] == matrix.shape[-1], 'Input matrices should be square' + assert matrix.shape[-1] % 2 == 0, 'Input matrix dimension should be even' + assert num_detectors >= 1, 'num_detectors should be positive' + assert chunk_size is None or chunk_size > 0, 'chunk_size should be positive' + + batch_size = matrix.shape[0] + nmode = matrix.shape[-1] // 2 + clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device) + if clicks.dim() == 1: + clicks = clicks.expand(batch_size, -1) + assert clicks.shape == (batch_size, nmode), f'Click patterns should have shape ({batch_size}, {nmode})' + assert (clicks >= 0).all(), 'Click numbers should be non-negative' + assert (clicks <= num_detectors).all(), 'Click numbers should not exceed num_detectors' + + if gamma is not None: + gamma = torch.as_tensor(gamma, device=matrix.device) + assert matrix.is_complex() or not gamma.is_complex(), 'Complex gamma requires complex matrices' + gamma = gamma.to(matrix.dtype) + if gamma.dim() == 1: + gamma = gamma.expand(batch_size, -1) + assert gamma.shape == (batch_size, 2 * nmode), f'gamma should have shape ({batch_size}, {2 * nmode})' + + pattern_groups = defaultdict(list) + for sample, pattern in enumerate(clicks.tolist()): + pattern_groups[tuple(pattern)].append(sample) + + indices = [] + values = [] + for pattern, sample_indices in pattern_groups.items(): + index = torch.tensor(sample_indices, dtype=torch.long, device=matrix.device) + gamma_group = None if gamma is None else gamma.index_select(0, index) + values.append( + _kensingtonian_same_clicks( + matrix.index_select(0, index), + pattern, + num_detectors, + gamma_group, + chunk_size, + ) + ) + indices.append(index) + + order = torch.cat(indices).argsort() + return torch.cat(values)[order]