Coverage for qml_essentials/entanglement.py: 86%

1from typing import Optional, Any, Tuple

2import jax

3import jax.numpy as jnp

4import numpy as np

6import jaqsi as js

7from jaqsi import operations as op

8from jaqsi import gateset

9from jaqsi.math import logm_v

10from qml_essentials.model import Model

11import logging

13log = logging.getLogger(__name__)

16class Entanglement:

17 @classmethod

18 def meyer_wallach(

19 cls,

20 model: Model,

21 n_samples: Optional[int | None],

22 random_key: Optional[jax.random.PRNGKey] = None,

23 scale: bool = False,

24 **kwargs: Any,

25 ) -> float:

26 """

27 Calculates the entangling capacity of a given quantum circuit

28 using Meyer-Wallach measure.

30 Args:

31 model (Model): The quantum circuit model.

32 n_samples (Optional[int]): Number of samples per qubit.

33 If None or < 0, the current parameters of the model are used.

34 random_key (Optional[jax.random.PRNGKey]): JAX random key for

35 parameter initialization. If None, uses the model's internal

36 random key.

37 scale (bool): Whether to scale the number of samples.

38 kwargs (Any): Additional keyword arguments for the model function.

40 Returns:

41 float: Entangling capacity of the given circuit, guaranteed

42 to be between 0.0 and 1.0.

43 """

44 if "noise_params" in kwargs:

45 log.warning(

46 "Meyer-Wallach measure not suitable for noisy circuits. "

47 "Consider 'mean_negativity' instead."

48 )

50 if scale:

51 n_samples = jnp.power(2, model.n_qubits) * n_samples

53 if n_samples is not None and n_samples > 0:

54 random_key = model.initialize_params(random_key, repeat=n_samples)

56 # implicitly set input to none in case it's not needed

57 kwargs.setdefault("inputs", None)

58 # explicitly set execution type because everything else won't work

59 rhos = model(execution_type="density", **kwargs).reshape(

60 -1, 2**model.n_qubits, 2**model.n_qubits

61 )

63 ent = cls._compute_meyer_wallach_meas(rhos, model.n_qubits)

65 log.debug(f"Variance of measure: {ent.var()}")

67 return ent.mean()

69 @classmethod

70 def _compute_meyer_wallach_meas(

71 cls, rhos: jnp.ndarray, n_qubits: int

72 ) -> jnp.ndarray:

73 """

74 Computes the Meyer-Wallach entangling capability measure for a given

75 set of density matrices.

77 Args:

78 rhos (jnp.ndarray): Density matrices of the sample quantum states.

79 The shape is (B_s, 2^n, 2^n), where B_s is the number of samples

80 (batch) and n the number of qubits

81 n_qubits (int): The number of qubits

83 Returns:

84 jnp.ndarray: Entangling capability for each sample, array with

85 shape (B_s,)

86 """

87 qb = list(range(n_qubits))

89 def _f(rhos):

90 entropy = 0

91 for j in range(n_qubits):

92 # Formula 6 in https://doi.org/10.48550/arXiv.quant-ph/0305094

93 # Trace out qubit j, keep all others

94 keep = qb[:j] + qb[j + 1 :]

95 density = js.partial_trace(rhos, n_qubits, keep)

96 # only real values, because imaginary part will be separate

97 # in all following calculations anyway

98 # entropy should be 1/2 <= entropy <= 1

99 entropy += jnp.trace((density @ density).real, axis1=-2, axis2=-1)

100

101 # inverse averaged entropy and scale to [0, 1]

102 return 2 * (1 - entropy / n_qubits)

103

104 return jax.vmap(_f)(rhos)

105

106 @classmethod

107 def mean_negativity(

108 cls,

109 model: Model,

110 n_samples: Optional[int],

111 random_key: Optional[jax.Array] = None,

112 scale: bool = False,

113 log: bool = False,

114 **kwargs: Any,

115 ) -> float:

116 r"""

117 Average normalized or logarithmic negativity over qubit-versus-rest partitions,

118 following [Vidal and Werner](https://doi.org/10.1103/PhysRevA.65.032314).

119

120 For pure or mixed states, this computes

121 :math:`\frac{2}{n}\sum_k \mathcal{N}_k(\rho)`, where

122 :math:`\mathcal{N}_k(\rho)=(\|\rho^{T_k}\|_1-1)/2`.

123 With log=True, average :math:`\log_2(1+2\mathcal{N}_k(\rho))`

124 over partitions and output states instead.

125 Zero does not exclude PPT entanglement, and a positive value does not

126 certify genuine multipartite entanglement.

127

128 Args:

129 model (Model): The quantum circuit model.

130 n_samples (Optional[int]): Number of parameter samples. If None or

131 nonpositive, use the current parameters without resampling.

132 random_key (Optional[jax.Array]): JAX key for parameter sampling.

133 If None, use the model's internal random key.

134 scale (bool): Multiply a positive sample count by 2**n_qubits.

135 log (bool): Use base-2 logarithmic negativity. Defaults to False.

136 kwargs (Any): Additional keyword arguments for the model function.

137

138 Returns:

139 float: Mean normalized or logarithmic negativity in [0, 1].

140 """

141 if n_samples is not None and n_samples > 0:

142 if scale:

143 n_samples *= 2**model.n_qubits

144 model.initialize_params(random_key, repeat=n_samples)

145

146 kwargs.setdefault("inputs", None)

147 rhos = model(execution_type="density", **kwargs).reshape(

148 -1, 2**model.n_qubits, 2**model.n_qubits

149 )

150 ent = cls._compute_mean_negativity(rhos, model.n_qubits, log=log)

151 logging.getLogger(__name__).debug(f"Variance of measure: {ent.var()}")

152 return float(ent.mean())

153

154 @classmethod

155 def _compute_mean_negativity(

156 cls, rhos: jnp.ndarray, n_qubits: int, log: bool = False

157 ) -> jnp.ndarray:

158 """

159 Compute mean normalized or logarithmic negativity for density matrices.

160

161 Args:

162 rhos (jnp.ndarray): Density matrices with shape (B, 2**n, 2**n).

163 n_qubits (int): Number of qubits.

164 log (bool): Apply the base-2 logarithm before averaging partitions.

165

166 Returns:

167 jnp.ndarray: Mean negativity with shape (B,).

168 """

169 tensor = rhos.reshape((-1,) + (2,) * (2 * n_qubits))

170 negativity = jnp.zeros(rhos.shape[0])

171 for k in range(n_qubits):

172 # Transpose only qubit k's row and column indices.

173 partial_transpose = jnp.swapaxes(tensor, k + 1, n_qubits + k + 1)

174 eigenvalues = jnp.linalg.eigvalsh(partial_transpose.reshape(rhos.shape))

175 value = 2 * jnp.maximum(-eigenvalues, 0).sum(axis=-1)

176 negativity += jnp.log1p(value) / jnp.log(2.0) if log else value

177 return negativity / n_qubits

178

179 @classmethod

180 def bell_measurements(

181 cls,

182 model: Model,

183 n_samples: int,

184 random_key: Optional[jax.random.PRNGKey] = None,

185 scale: bool = False,

186 **kwargs: Any,

187 ) -> float:

188 """

189 Compute the Bell measurement for a given model.

190

191 Constructs a ``2 * n_qubits`` circuit that prepares two copies of

192 the model state (on disjoint qubit registers), applies CNOTs and

193 Hadamards, and measures probabilities on the first register.

194

195 Args:

196 model (Model): The quantum circuit model.

197 n_samples (int): The number of samples to compute the measure for.

198 random_key (Optional[jax.random.PRNGKey]): JAX random key for

199 parameter initialization. If None, uses the model's internal

200 random key.

201 scale (bool): Whether to scale the number of samples

202 according to the number of qubits.

203 **kwargs (Any): Additional keyword arguments for the model function.

204

205 Returns:

206 float: The Bell measurement value.

207 """

208 if "noise_params" in kwargs:

209 log.warning(

210 "Bell Measurements not suitable for noisy circuits. "

211 "Consider 'mean_negativity' instead."

212 )

213

214 if scale:

215 n_samples = jnp.power(2, model.n_qubits) * n_samples

216

217 n = model.n_qubits

218

219 def _bell_circuit(params, inputs, pulse_params=None, random_key=None, **kw):

220 """Bell measurement circuit on 2*n qubits."""

221 from jaqsi.tape import copy_to_tape

222

223 def vari():

224 model._variational(

225 params,

226 inputs,

227 pulse_params=pulse_params,

228 random_key=random_key,

229 **kw,

230 )

231

232 # First copy on wires 0..n-1

233 vari()

234 # Second copy on wires n..2n-1

235 copy_to_tape(vari, offset=n)

236

237 # Bell measurement: CNOT + H

238 for q in range(n):

239 gateset.CX(wires=[q, q + n])

240 gateset.H(wires=q)

241

242 bell_script = js.Script(f=_bell_circuit, n_qubits=2 * n)

243

244 if n_samples is not None and n_samples > 0:

245 random_key = model.initialize_params(random_key, repeat=n_samples)

246 params = model.params

247 else:

248 if len(model.params.shape) <= 2:

249 params = model.params.reshape(1, *model.params.shape)

250 else:

251 log.info(f"Using sample size of model params: {model.params.shape[0]}")

252 params = model.params

253

254 n_samples = params.shape[0]

255 inputs = model._inputs_validation(kwargs.get("inputs", None))

256

257 # Execute: vmap over batch dimension of params (axis 0)

258 if n_samples > 1:

259 from jaqsi.utils import safe_random_split

260

261 random_keys = safe_random_split(random_key, num=n_samples)

262 result = bell_script.execute(

263 type="probs",

264 args=(params, inputs, model.pulse_params, random_keys),

265 kwargs=kwargs,

266 in_axes=(0, None, None, 0),

267 )

268 else:

269 result = bell_script.execute(

270 type="probs",

271 args=(params, inputs, model.pulse_params, random_key),

272 kwargs=kwargs,

273 )

274

275 # Marginalize: for each qubit q, keep wires [q, q+n] from the 2n-qubit probs

276 # The last probability in each pair gives P(|11⟩) for that qubit pair

277 per_qubit = []

278 for q in range(n):

279 marg = js.marginalize_probs(result, 2 * n, [q, q + n])

280 per_qubit.append(marg)

281 # per_qubit[q] has shape (n_samples, 4) or (4,)

282 exp = jnp.stack(per_qubit, axis=-2) # (..., n, 4)

283 exp = 1 - 2 * exp[..., -1] # (..., n)

284

285 if not jnp.isclose(jnp.sum(exp.imag), 0, atol=1e-6):

286 log.warning("Imaginary part of probabilities detected")

287 exp = jnp.abs(exp)

288

289 measure = 2 * (1 - exp.mean(axis=0))

290 entangling_capability = min(max(float(measure.mean()), 0.0), 1.0)

291 log.debug(f"Variance of measure: {measure.var()}")

292

293 return entangling_capability

294

295 @classmethod

296 def relative_entropy(

297 cls,

298 model: Model,

299 n_samples: int,

300 n_sigmas: int,

301 random_key: Optional[jax.random.PRNGKey] = None,

302 scale: bool = False,

303 **kwargs: Any,

304 ) -> float:

305 """

306 Calculates the relative entropy of entanglement of a given quantum

307 circuit. This measure is also applicable to mixed state, albeit it

308 might me not fully accurate in this simplified case.

309

310 As the relative entropy is generally defined as the smallest relative

311 entropy from the state in question to the set of separable states.

312 However, as computing the nearest separable state is NP-hard, we select

313 n_sigmas of random separable states to compute the distance to, which

314 is not necessarily the nearest. Thus, this measure of entanglement

315 presents an upper limit of entanglement.

316

317 As the relative entropy is not necessarily between zero and one, this

318 function also normalises by the relative entroy to the GHZ state.

319

320 Args:

321 model (Model): The quantum circuit model.

322 n_samples (int): Number of samples per qubit.

323 If <= 0, the current parameters of the model are used.

324 n_sigmas (int): Number of random separable pure states to compare against.

325 random_key (Optional[jax.random.PRNGKey]): JAX random key for

326 parameter initialization. If None, uses the model's internal

327 random key.

328 scale (bool): Whether to scale the number of samples.

329 kwargs (Any): Additional keyword arguments for the model function.

330

331 Returns:

332 float: Entangling capacity of the given circuit, guaranteed

333 to be between 0.0 and 1.0.

334 """

335 dim = jnp.power(2, model.n_qubits)

336 if scale:

337 n_samples = dim * n_samples

338 n_sigmas = dim * n_sigmas

339

340 if random_key is None:

341 random_key = model.random_key

342

343 # Random separable states

344 log_sigmas = sample_random_separable_states(

345 model.n_qubits, n_samples=n_sigmas, random_key=random_key, take_log=True

346 )

347

348 random_key, _ = jax.random.split(random_key)

349

350 if n_samples is not None and n_samples > 0:

351 model.initialize_params(random_key, repeat=n_samples)

352 else:

353 if len(model.params.shape) <= 2:

354 model.params = model.params.reshape(1, *model.params.shape)

355 else:

356 log.info(f"Using sample size of model params: {model.params.shape[0]}")

357

358 rhos, log_rhos = cls._compute_log_density(model, **kwargs)

359

360 rel_entropies = jnp.zeros((n_sigmas, model.params.shape[0]))

361

362 for i, log_sigma in enumerate(log_sigmas):

363 rel_entropies = rel_entropies.at[i].set(

364 cls._compute_rel_entropies(rhos, log_rhos, log_sigma)

365 )

366

367 # Entropy of GHZ states should be maximal

368 ghz_model = Model(model.n_qubits, 1, "GHZ", data_reupload=False)

369 rho_ghz, log_rho_ghz = cls._compute_log_density(ghz_model, **kwargs)

370 ghz_entropies = cls._compute_rel_entropies(rho_ghz, log_rho_ghz, log_sigmas)

371

372 normalised_entropies = rel_entropies / ghz_entropies

373

374 # Average all iterated states

375 entangling_capability = normalised_entropies.T.min(axis=1)

376 log.debug(f"Variance of measure: {entangling_capability.var()}")

377

378 return entangling_capability.mean()

379

380 @classmethod

381 def _compute_log_density(

382 cls, model: Model, **kwargs

383 ) -> Tuple[jnp.ndarray, jnp.ndarray]:

384 """

385 Obtains the density matrix of a model and computes its logarithm.

386

387 Args:

388 model (Model): The model for which to compute the density matrix.

389

390 Returns:

391 Tuple[jnp.ndarray, jnp.ndarray]:

392 - jnp.ndarray: density matrix.

393 - jnp.ndarray: logarithm of the density matrix.

394 """

395 # implicitly set input to none in case it's not needed

396 kwargs.setdefault("inputs", None)

397 # explicitly set execution type because everything else won't work

398 rho = model(execution_type="density", **kwargs)

399 rho = rho.reshape(-1, 2**model.n_qubits, 2**model.n_qubits)

400 log_rho = logm_v(rho) / jnp.log(2)

401 return rho, log_rho

402

403 @classmethod

404 def _compute_rel_entropies(

405 cls,

406 rhos: jnp.ndarray,

407 log_rhos: jnp.ndarray,

408 log_sigmas: jnp.ndarray,

409 ) -> jnp.ndarray:

410 """

411 Compute the relative entropy for a given model.

412

413 Args:

414 rhos (jnp.ndarray): Density matrix result of the circuit, has shape

415 (R, 2^n, 2^n), with the batch size R and number of qubits n

416 log_rhos (jnp.ndarray): Corresponding logarithm of the density

417 matrix, has shape (R, 2^n, 2^n).

418 log_sigmas (jnp.ndarray): Density matrix of next separable state,

419 has shape (2^n, 2^n) if it's a single sigma or (S, 2^n, 2^n),

420 with the batch size S (number of sigmas).

421

422 Returns:

423 jnp.ndarray: Relative Entropy for each sample

424 """

425 n_rhos = rhos.shape[0]

426 if len(log_sigmas.shape) == 3:

427 n_sigmas = log_sigmas.shape[0]

428 rhos = jnp.tile(rhos, (n_sigmas, 1, 1))

429 log_rhos = jnp.tile(log_rhos, (n_sigmas, 1, 1))

430 einsum_subscript = "ij,jk->ik"

431 else:

432 n_sigmas = 1

433 log_sigmas = log_sigmas[jnp.newaxis, ...].repeat(n_rhos, axis=0)

434

435 einsum_subscript = "ij,jk->ik"

436

437 def _f(rhos, log_rhos, log_sigmas):

438 prod = jnp.einsum(einsum_subscript, rhos, log_rhos - log_sigmas)

439 rel_entropies = jnp.abs(jnp.trace(prod, axis1=-2, axis2=-1))

440 return rel_entropies

441

442 rel_entropies = jax.vmap(_f, in_axes=(0, 0, 0))(rhos, log_rhos, log_sigmas)

443

444 if n_sigmas > 1:

445 rel_entropies = rel_entropies.reshape(n_sigmas, n_rhos)

446 return rel_entropies

447

448 @classmethod

449 def entanglement_of_formation(

450 cls,

451 model: Model,

452 n_samples: int,

453 random_key: Optional[jax.random.PRNGKey] = None,

454 scale: bool = False,

455 always_decompose: bool = False,

456 **kwargs: Any,

457 ) -> float:

458 """

459 This function implements the entanglement of formation for mixed

460 quantum systems.

461 In that a mixed state gets decomposed into pure states with respective

462 probabilities using the eigendecomposition of the density matrix.

463 Then, the Meyer-Wallach measure is computed for each pure state,

464 weighted by the eigenvalue.

465 See e.g. https://doi.org/10.48550/arXiv.quant-ph/0504163

466

467 Note that the decomposition is *not unique*! Therefore, this measure

468 presents the entanglement for *some* decomposition into pure states,

469 not necessarily the one that is anticipated when applying the Kraus

470 channels.

471 If a pure state is provided, this results in the same value as the

472 Entanglement.meyer_wallach function if `always_decompose` flag is not set.

473

474 Args:

475 model (Model): The quantum circuit model.

476 n_samples (int): Number of samples per qubit.

477 random_key (Optional[jax.random.PRNGKey]): JAX random key for

478 parameter initialization. If None, uses the model's internal

479 random key.

480 scale (bool): Whether to scale the number of samples.

481 always_decompose (bool): Whether to explicitly compute the

482 entantlement of formation for the eigendecomposition of a pure

483 state.

484 kwargs (Any): Additional keyword arguments for the model function.

485

486 Returns:

487 float: Entangling capacity of the given circuit, guaranteed

488 to be between 0.0 and 1.0.

489 """

490

491 if scale:

492 n_samples = jnp.power(2, model.n_qubits) * n_samples

493

494 if n_samples is not None and n_samples > 0:

495 model.initialize_params(random_key, repeat=n_samples)

496 else:

497 if len(model.params.shape) <= 2:

498 model.params = model.params.reshape(1, *model.params.shape)

499 else:

500 log.info(f"Using sample size of model params: {model.params.shape[0]}")

501

502 # implicitly set input to none in case it's not needed

503 kwargs.setdefault("inputs", None)

504 rhos = model(execution_type="density", **kwargs)

505 rhos = rhos.reshape(-1, 2**model.n_qubits, 2**model.n_qubits)

506 ent = cls._compute_entanglement_of_formation(

507 rhos, model.n_qubits, always_decompose

508 )

509 return ent.mean()

510

511 @classmethod

512 def _compute_entanglement_of_formation(

513 cls,

514 rhos: jnp.ndarray,

515 n_qubits: int,

516 always_decompose: bool,

517 ) -> jnp.ndarray:

518 """

519 Computes the entanglement of formation for a given batch of density

520 matrices.

521

522 Args:

523 rho (jnp.ndarray): The density matrices, has shape (B_s, 2^n, 2^n),

524 where B_s is the batch size and n the number of qubits.

525 n_qubits (int): Number of qubits

526 always_decompose (bool): Whether to explicitly compute the

527 entantlement of formation for the eigendecomposition of a pure

528 state.

529

530 Returns:

531 jnp.ndarray: Entanglement for the provided density matrices.

532 """

533 eigenvalues, eigenvectors = jnp.linalg.eigh(rhos)

534 if not always_decompose and jnp.isclose(eigenvalues, 1.0).any(axis=-1).all():

535 return cls._compute_meyer_wallach_meas(rhos, n_qubits)

536

537 rhos = np.einsum("sji,ski->sijk", eigenvectors, eigenvectors.conjugate())

538 measures = cls._compute_meyer_wallach_meas(

539 rhos.reshape(-1, 2**n_qubits, 2**n_qubits), n_qubits

540 )

541 ent = np.einsum("si,si->s", measures.reshape(-1, 2**n_qubits), eigenvalues)

542 return ent

543

544 @classmethod

545 def concentratable_entanglement(

546 cls,

547 model: Model,

548 n_samples: int,

549 random_key: Optional[jax.random.PRNGKey] = None,

550 scale: bool = False,

551 **kwargs: Any,

552 ) -> float:

553 """

554 Computes the concentratable entanglement of a given model.

555

556 This method utilizes the Concentratable Entanglement measure from

557 https://arxiv.org/abs/2104.06923. The swap test is implemented

558 directly in jaqsi using a ``3 * n_qubits`` circuit.

559

560 Args:

561 model (Model): The quantum circuit model.

562 n_samples (int): The number of samples to compute the measure for.

563 random_key (Optional[jax.random.PRNGKey]): JAX random key for

564 parameter initialization. If None, uses the model's internal

565 random key.

566 scale (bool): Whether to scale the number of samples according to

567 the number of qubits.

568 **kwargs (Any): Additional keyword arguments for the model function.

569

570 Returns:

571 float: Entangling capability of the given circuit, guaranteed

572 to be between 0.0 and 1.0.

573 """

574 n = model.n_qubits

575 N = 2**n

576

577 if scale:

578 n_samples = N * n_samples

579

580 def _swap_test_circuit(

581 params, inputs, pulse_params=None, random_key=None, **kw

582 ):

583 """Swap-test circuit on 3*n qubits."""

584 from jaqsi.tape import copy_to_tape

585

586 def vari():

587 model._variational(

588 params,

589 inputs,

590 pulse_params=pulse_params,

591 random_key=random_key,

592 **kw,

593 )

594

595 # First copy on wires n..2n-1

596 copy_to_tape(vari, offset=n)

597 # Second copy on wires 2n..3n-1

598 copy_to_tape(vari, offset=2 * n)

599

600 # Swap test: H on ancilla register (wires 0..n-1)

601 for i in range(n):

602 gateset.H(wires=i)

603

604 for i in range(n):

605 gateset.CSWAP(wires=[i, i + n, i + 2 * n])

606

607 for i in range(n):

608 gateset.H(wires=i)

609

610 swap_script = js.Script(f=_swap_test_circuit, n_qubits=3 * n)

611

612 if n_samples is not None and n_samples > 0:

613 random_key = model.initialize_params(random_key, repeat=n_samples)

614 else:

615 if len(model.params.shape) <= 2:

616 model.params = model.params.reshape(1, *model.params.shape)

617 else:

618 log.info(f"Using sample size of model params: {model.params.shape[0]}")

619

620 params = model.params

621 inputs = model._inputs_validation(kwargs.get("inputs", None))

622 n_batch = params.shape[0]

623

624 marg_probs = jax.jit(js.marginalize_probs, static_argnums=(1, 2))

625

626 if n_batch > 1:

627 from jaqsi.utils import safe_random_split

628

629 random_keys = safe_random_split(random_key, num=n_batch)

630 probs = swap_script.execute(

631 type="probs",

632 args=(params, inputs, model.pulse_params, random_keys),

633 in_axes=(0, None, None, 0),

634 kwargs=kwargs,

635 )

636 else:

637 probs = swap_script.execute(

638 type="probs",

639 args=(params, inputs, model.pulse_params, random_key),

640 kwargs=kwargs,

641 )

642

643 # Marginalize to the ancilla register (wires 0..n-1)

644 probs = marg_probs(probs, 3 * n, tuple(range(n)))

645

646 ent = 1 - probs[..., 0]

647

648 log.debug(f"Variance of measure: {ent.var()}")

649

650 return float(ent.mean())

651

652 @classmethod

653 def concentratable_entanglement_estimation(

654 cls,

655 model: Model,

656 n_samples: int,

657 random_key: Optional[jax.random.PRNGKey] = None,

658 scale: bool = False,

659 **kwargs: Any,

660 ) -> float:

661 """

662 Computes the concentratable entanglement of a given model.

663

664 This method utilizes the Concentratable Entanglement measure from

665 https://arxiv.org/abs/2104.06923. The swap test is implemented

666 directly in jaqsi using a ``3 * n_qubits`` circuit.

667

668 Args:

669 model (Model): The quantum circuit model.

670 n_samples (int): The number of samples to compute the measure for.

671 random_key (Optional[jax.random.PRNGKey]): JAX random key for

672 parameter initialization. If None, uses the model's internal

673 random key.

674 scale (bool): Whether to scale the number of samples according to

675 the number of qubits.

676 **kwargs (Any): Additional keyword arguments for the model function.

677

678 Returns:

679 float: Entangling capability of the given circuit, guaranteed

680 to be between 0.0 and 1.0.

681 """

682 n = model.n_qubits

683 N = 2**n

684

685 if scale:

686 n_samples = N * n_samples

687

688 def _bell_basis_measurement(

689 params, inputs, pulse_params=None, random_key=None, **kw

690 ):

691 """Bell-basis measurement circuit on 3*n qubits."""

692 from jaqsi.tape import copy_to_tape

693

694 def vari():

695 model._variational(

696 params,

697 inputs,

698 pulse_params=pulse_params,

699 random_key=random_key,

700 **kw,

701 )

702

703 # First copy on wires 0..n-1

704 copy_to_tape(vari, offset=0)

705 # Second copy on wires n..2n-1

706 copy_to_tape(vari, offset=n)

707

708 for i in range(n):

709 gateset.CX(wires=[i, i + n])

710 gateset.H(wires=i)

711

712 bell_basis_script = js.Script(f=_bell_basis_measurement, n_qubits=2 * n)

713

714 if n_samples is not None and n_samples > 0:

715 random_key = model.initialize_params(random_key, repeat=n_samples)

716 else:

717 if len(model.params.shape) <= 2:

718 model.params = model.params.reshape(1, *model.params.shape)

719 else:

720 log.info(f"Using sample size of model params: {model.params.shape[0]}")

721

722 params = model.params

723 inputs = model._inputs_validation(kwargs.get("inputs", None))

724 n_batch = params.shape[0]

725

726 # SWAP operator in Bell-basis

727 SWAP = jnp.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, -1]])

728 # Construct observable for measuring CE

729 CE_observable = gateset.Id([0, n]) + op.Operation([0, n], SWAP)

730 for i in range(1, n):

731 CE_observable = CE_observable @ (

732 gateset.Id([i, i + n]) + op.Operation([i, i + n], SWAP)

733 )

734 CE_observable = (1 / N) * CE_observable

735

736 expvals = []

737 if n_batch > 1:

738 from jaqsi.utils import safe_random_split

739

740 random_keys = safe_random_split(random_key, num=n_batch)

741 expvals = bell_basis_script.execute(

742 type="expval",

743 obs=[CE_observable],

744 args=(params, inputs, model.pulse_params, random_keys),

745 in_axes=(0, None, None, 0),

746 kwargs=kwargs,

747 )

748 else:

749 expvals = bell_basis_script.execute(

750 type="expval",

751 obs=[CE_observable],

752 args=(params, inputs, model.pulse_params, random_key),

753 kwargs=kwargs,

754 )

755

756 ent = 1 - expvals

757 log.debug(f"Variance of measure: {ent.var()}")

758 return float(ent.mean())

759

760

761def sample_random_separable_states(

762 n_qubits: int,

763 n_samples: int,

764 random_key: jax.random.PRNGKey,

765 take_log: bool = False,

766) -> jnp.ndarray:

767 """

768 Sample random separable states (density matrix).

769

770 Args:

771 n_qubits (int): number of qubits in the state

772 n_samples (int): number of states

773 random_key (random.PRNGKey): JAX random key

774 take_log (bool): if the matrix logarithm of the density matrix should be taken.

775

776 Returns:

777 jnp.ndarray: Density matrices of shape (n_samples, 2**n_qubits, 2**n_qubits)

778 """

779 model = Model(n_qubits, 1, "No_Entangling", data_reupload=False)

780 model.initialize_params(random_key, repeat=n_samples)

781 # explicitly set execution type because anything else won't work

782 sigmas = model(execution_type="density", inputs=None)

783 if take_log:

784 sigmas = logm_v(sigmas) / jnp.log(2.0 + 0j)

785

786 return sigmas

Coverage for qml_essentials / entanglement.py: 86%

258 statements