Coverage for qml_essentials / trainability.py: 72%

1"""Empirical loss-variance (barren-plateau) diagnostics over random circuits. 

2 

3Estimates ``Var_W[<O>]`` for a fixed input state evolved by a random 

4parameterised circuit, sampling the circuit parameters ``W ~ U[0, 2pi)``. This 

5is the empirical left-hand side of the Ragone/Fontana LASA identity 

6``Var_W[<O>] = sum_j P_j(rho) P_j(O) / dim g_j`` used to certify trainability of 

7polynomial-DLA models; it complements the analytic vehicle 

8:func:`qml_essentials.algebra.g_purity_from_basis`. 

9 

10The input is either an angle-encoded product state ``prod_k R_y(theta_k)|0>`` (pass 

11``theta``) or an explicit statevector (pass ``init_state``, e.g. a Haar-random 

12input), and the ansatz is any layer that writes onto the active jaqsi tape. 

13""" 

14 

15from __future__ import annotations 

16 

17import jax 

18import jax.numpy as jnp 

19import numpy as np 

20 

21from jaqsi.gates import Gates 

22from jaqsi import gateset 

23import jaqsi as js 

24from qml_essentials.ansaetze import Ansaetze 

25 

26 

27def ansatz_layer(circuit_type: str): 

28 """Return ``(layer_fn, n_params_per_layer)`` for a built-in ``Ansaetze`` template. 

29 

30 ``layer_fn(p, n)`` applies one layer of ``circuit_type`` (e.g. ``"Matchgate"``, 

31 ``"Strongly_Entangling"``) onto the active jaqsi tape, and 

32 ``n_params_per_layer(n)`` gives the length of ``p``. 

33 """ 

34 cls = getattr(Ansaetze, circuit_type) 

35 

36 def layer(params, n): 

37 cls.build(params, n) 

38 

39 return layer, (lambda n: cls.n_params_per_layer(n)) 

40 

41 

42def loss_variance( 

43 layer_fn, 

44 n_params_per_layer: int, 

45 theta: np.ndarray | None, 

46 depth: int, 

47 n_samples: int, 

48 key, 

49 obs=None, 

50 out_qubit: int | None = None, 

51 shots: int | None = None, 

52 init_state: np.ndarray | None = None, 

53): 

54 """Empirical ``Var_W[<O>]`` and the raw loss values. 

55 

56 Args: 

57 layer_fn: applies one ansatz layer onto the tape, ``layer_fn(p, n)``. 

58 n_params_per_layer: length of ``p`` consumed per layer. 

59 theta: fixed angle configuration, shape ``(n,)``, angle-encoded as a 

60 product state ``prod_k R_y(theta_k)|0>``. Ignored if ``init_state`` 

61 is given. 

62 depth: number of ansatz layers (circuit-side mixing / 2-design depth). 

63 n_samples: number of random ``W`` draws. 

64 key: JAX PRNG key. 

65 obs: list of observables whose summed expectation is the loss; defaults 

66 to a single bulk-qubit ``[PauliZ(out_qubit)]`` (a matchgate readout). 

67 out_qubit: measured qubit for the default observable (default ``n // 2``). 

68 shots: finite shots for the expectation (``None`` -> exact). 

69 init_state: optional input statevector ``(2**n,)``; bypasses the ``R_y`` 

70 product encoding, so ``theta`` is ignored when given (e.g. to inject 

71 a Haar-random small-g-purity input). 

72 

73 Returns: 

74 ``(variance: float, losses: np.ndarray of shape (n_samples,))``. 

75 """ 

76 if init_state is not None: 

77 init_state = jnp.asarray(init_state) 

78 n = int(round(float(np.log2(init_state.shape[-1])))) 

79 else: 

80 theta = jnp.asarray(theta, dtype=float) 

81 n = int(theta.shape[0]) 

82 i_out = n // 2 if out_qubit is None else out_qubit 

83 if obs is None: 

84 obs = [gateset.PauliZ(wires=i_out)] 

85 

86 # Separate streams for the parameter draws and the shot noise, so the two 

87 # sources of randomness stay independent. 

88 param_key, shot_key = jax.random.split(key) 

89 W = jax.random.uniform( 

90 param_key, (n_samples, depth, n_params_per_layer), minval=0.0, maxval=2 * np.pi 

91 ) 

92 

93 if init_state is None: 

94 

95 def circ(params, th): 

96 for q in range(n): 

97 Gates.RY(th[q], wires=q) 

98 for d in range(depth): 

99 layer_fn(params[d], n) 

100 

101 vals = js.Script(f=circ, n_qubits=n).execute( 

102 type="expval", 

103 obs=obs, 

104 args=(W, theta), 

105 in_axes=(0, None), 

106 shots=shots, 

107 key=shot_key, 

108 ) 

109 else: 

110 

111 def circ(params): # init_state already encodes the input; apply only U(W) 

112 for d in range(depth): 

113 layer_fn(params[d], n) 

114 

115 vals = js.Script(f=circ, n_qubits=n).execute( 

116 type="expval", 

117 obs=obs, 

118 args=(W,), 

119 in_axes=(0,), 

120 shots=shots, 

121 key=shot_key, 

122 initial_state=init_state, 

123 ) 

124 vals = np.asarray(vals) 

125 if vals.ndim > 1 and len(obs) > 1: # loss = <sum_i O_i> 

126 vals = vals.sum(axis=-1) 

127 vals = vals.reshape(-1) 

128 return float(np.var(vals, ddof=1)), vals