pqcprep.training_tools

  1"""
  2Collection of useful functions for network training purposes.
  3"""
  4
  5import numpy as np 
  6from qiskit_machine_learning.neural_networks import SamplerQNN
  7from qiskit_machine_learning.connectors import TorchConnector
  8from qiskit_algorithms.utils import algorithm_globals 
  9from qiskit.primitives import Sampler
 10from torch.optim import Adam
 11from torch import Tensor, no_grad
 12import sys, time, os, warnings, torch  
 13
 14from .binary_tools import bin_to_dec, dec_to_bin  
 15from .phase_tools import full_encode, phase_from_state
 16from .pqc_tools import generate_network, binary_to_encode_param, A_generate_network, get_state_vec  
 17from .file_tools import compress_args,compress_args_ampl, vars_to_name_str, vars_to_name_str_ampl 
 18from .psi_tools import psi, A 
 19
 20#---------------------------------------------------------------------------------------------------
 21
 22def set_loss_func(loss_str, arg_dict, ampl=False):
 23    r"""
 24    Set the loss function to be used in training a network. 
 25
 26    Arguments:
 27    ----
 28
 29    - **loss_str** : *str*
 30
 31        String specifying the loss function to use. Options are 
 32
 33        - `'MSE'` : mean squared error loss. Using [pytorch's implementation](https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html) with default settings. 
 34          `criterion` takes two pytorch `Tensors` as inputs, corresponding to the network output and the desired output. 
 35
 36        - `'L1'` : mean absolute error loss. Using [pytorch's implementation](https://pytorch.org/docs/stable/generated/torch.nn.L1Loss.html) with default settings. 
 37          `criterion` takes two pytorch `Tensors` as inputs, corresponding to the network output and the desired output.    
 38
 39        - `'KLD'` : Kullback-Leibler divergence loss. Using [pytorch's implementation](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html) with default settings. 
 40          `criterion` takes two pytorch `Tensors` as inputs, corresponding to the network output and the desired output.   
 41
 42        - `'CE'` : cross entropy loss. Using [pytorch's implementation](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) with default settings. 
 43          `criterion` takes two pytorch `Tensors` as inputs, corresponding to the network output and the desired output.  
 44
 45        - `'SAM'` : sign-adjusted mismatch. Defined as $$\text{SAM}(\ket{x}, \ket{y}) =  1 - \sum_k |x_k| |y_k|,$$ where $\ket{x}$, $\ket{y}$ are the 
 46            network output and desired output, respectively, and $x_k$, $y_k$ are the coefficients w.r.t the two-register computational basis states.  `criterion` takes 
 47            two pytorch `Tensors` as inputs, corresponding to the network output and the desired output.       
 48
 49        - `'WIM'` : weighted mismatch. Defined analogously to SAM but with additional weights: $$\text{WIM}(\ket{x}, \ket{y}) =1 - \sum_k w_k |x_k||y_k|.$$
 50           `criterion` takes three pytorch `Tensors` as inputs, corresponding to the network output, the desired output and the weights. See `set_WIM_weights()`
 51           for information on how the weights are calculated. This loss function is not an option if `ampl` is True. 
 52
 53        - `'WILL'` : weighted Lp loss. Defined as $$\text{WILL}(\ket{x}, \ket{y}; p,q) = \sum_k |x_k -y_k|^p + |x_k| |[k]_m - \Psi([k]_n)|^q,$$ where $p$ and $q$ are coefficients stored in 
 54           `arg_dict['WILL_p']`, `arg_dict['WILL_q']`, $[k]_n$ is the target register bit-string associated with the basis state $\ket{k}$, $[k]_n$ is the input register bit-string 
 55           associated with the basis state $\ket{k}$, $\Psi$ is the function to be evaluated for the network, and $x_k$, $y_k$ have the same meaning as above.  `criterion` takes 
 56            two pytorch `Tensors` as inputs, corresponding to the network output and the desired output. This loss function is not an option if `ampl` is True.               
 57
 58    - **arg_dict** : *dict* 
 59
 60        A dictionary containing information on training variables, created with `pqcprep.file_tools.compress_args()` (or created 
 61        with `pqcprep.file_tools.compress_args_ampl()` in the case of `ampl` being True). 
 62
 63    - **ampl** : *boolean* 
 64
 65        If True, the loss function is defined for an amplitude-encoding network, as opposed to a function evaluation network. Default is 
 66        False.     
 67        
 68    Returns: 
 69    ----
 70
 71    - **criterion** : *callable* 
 72
 73        The loss function as a callable object. Number and type of arguments depend on the chosen `loss_str` option (see above).    
 74
 75    """
 76    
 77    if loss_str=="MSE":
 78        from torch.nn import MSELoss
 79        criterion=MSELoss() 
 80    elif loss_str=="L1":
 81        from torch.nn import L1Loss
 82        criterion=L1Loss() 
 83    elif loss_str=="KLD":
 84        from torch.nn import KLDivLoss
 85        criterion=KLDivLoss()
 86    elif loss_str=="CE":
 87        from torch.nn import CrossEntropyLoss
 88        criterion=CrossEntropyLoss()
 89    elif loss_str=="SAM":
 90        def criterion(output, target):
 91            return  torch.abs(1. -torch.sum(torch.mul(output, target))) 
 92    elif loss_str=="WIM":  
 93        if arg_dict["train_superpos"]==False:
 94            raise ValueError(f"The loss function {loss_str} requires training in superposition, i.e. 'train_superpos==True'.") 
 95        if ampl:
 96            raise ValueError(f"The loss function {loss_str} is not available for amplitude training. Allowed options are 'CE', 'MSE', 'L1', 'KLD', 'SAM'.")
 97
 98        def criterion(output, target, weights):
 99            output = torch.mul(output, weights)  
100            output = output / torch.sum(torch.mul(output, output)) 
101            return  torch.abs(1. -torch.sum(torch.mul(output, target)))     
102    elif loss_str=="WILL":  
103        if arg_dict["train_superpos"]==False:
104             raise ValueError(f"The loss function {loss_str} requires training in superposition, i.e. 'train_superpos==True'.")
105        if ampl:
106            raise ValueError(f"The loss function {loss_str} is not available for amplitude training. Allowed options are 'CE', 'MSE', 'L1', 'KLD', 'SAM'.")
107        fx_arr = [psi(i, mode=arg_dict["func_str"]) for i in np.arange(0, 2**arg_dict["n"])]
108        if arg_dict["phase_reduce"]:
109            fx_arr = [np.modf(i/ (2* np.pi))[0] for i in fx_arr]
110        
111        fx_arr_rounded = [bin_to_dec(dec_to_bin(i,arg_dict["m"],'unsigned mag', nint=arg_dict["mint"]),'unsigned mag',nint=arg_dict["mint"]) for i in fx_arr]
112        distance_arr = np.empty(2**(arg_dict["n"]+arg_dict["m"]))
113        
114        for i in np.arange(2**arg_dict["n"]):
115            bin_i=dec_to_bin(i,arg_dict["n"],'unsigned mag') 
116            for j in np.arange(2**arg_dict["m"]):
117                bin_j=dec_to_bin(j,arg_dict["m"],'unsigned mag')
118                ind=int(bin_j + bin_i,2) 
119                distance_arr[ind] = np.abs(bin_to_dec(dec_to_bin(j,arg_dict["m"],'unsigned mag'),'unsigned mag',nint=arg_dict["mint"]) - fx_arr_rounded[i])  
120        distance=Tensor(distance_arr)   
121        
122        def criterion(output, target):
123            loss =torch.pow(torch.abs(output-target),arg_dict["WILL_p"]) + torch.mul(torch.abs(output),torch.pow(distance,arg_dict["WILL_q"])) 
124            return torch.sum(loss)**(1/arg_dict["WILL_p"]) / torch.numel(loss)
125    else:
126        raise ValueError("Unrecognised loss function. Options are: 'CE', 'MSE', 'L1', 'KLD', 'SAM', 'WIM', 'WILL'.")
127
128    return criterion 
129
130def set_WIM_weights(generated_weights, arg_dict):
131    """ 
132    
133    Determine the weight coefficients for the WIM loss function. 
134
135    Arguments:
136    ----
137    - **generated_weights** : *array_like*
138
139        The network weights generated in a given epoch. 
140
141    - **arg_dict** : *dict* 
142
143        A dictionary containing information on training variables, created with `pqcprep.file_tools.compress_args()`. 
144
145    Returns:
146    ----
147    - **WIM_weights_arr** : *array_like*
148
149        The calculated coefficients for the WIM loss function. 
150
151    """
152    
153    # initialise arrays to store results 
154    WIM_weights_arr= np.empty(2**(arg_dict["n"]+arg_dict["m"]))
155    
156    # iterate over input states 
157    x_arr_temp=np.arange(2**arg_dict["n"])
158    fx_arr_temp = [psi(k, mode=arg_dict["func_str"]) for k in x_arr_temp]
159
160    if arg_dict["phase_reduce"]: fx_arr_temp = [np.modf(k/ (2* np.pi))[0] for k in fx_arr_temp]
161
162    for q in x_arr_temp:
163        
164        # prepare circuit 
165        enc=binary_to_encode_param(np.binary_repr(q,arg_dict["n"]))
166        params=np.concatenate((enc, generated_weights))  
167
168        qc = generate_network(arg_dict["n"],arg_dict["m"],arg_dict["L"], encode=True,toggle_IL=True, real=arg_dict["real"],repeat_params=arg_dict["repeat_params"])
169        qc = qc.assign_parameters(params) 
170
171        # get target array 
172        target_arr_temp = np.zeros(2**(arg_dict["n"]+arg_dict["m"]))
173
174        index = int(dec_to_bin(fx_arr_temp[q],arg_dict["m"],'unsigned mag',nint=arg_dict["mint"])+dec_to_bin(x_arr_temp[q],arg_dict["n"],'unsigned mag',nint=arg_dict["nint"]),2)
175        target_arr_temp[index]=1 
176
177        # get statevector 
178        state_vector_temp = get_state_vec(qc)
179
180        # for each output state, calculate the "binary difference" to the target as well as the "coefficient difference"
181        sum = 0 
182        for j in np.arange(2**arg_dict["m"]):
183            ind = int(dec_to_bin(j,arg_dict["m"],'unsigned mag',nint=arg_dict["m"])+dec_to_bin(x_arr_temp[q],arg_dict["n"],'unsigned mag',nint=arg_dict["nint"]),2)
184            
185            num_dif =np.abs(fx_arr_temp[q] -bin_to_dec(dec_to_bin(j,arg_dict["m"],'unsigned mag', nint=arg_dict["mint"]),'unsigned mag',nint=arg_dict["mint"]))
186            coeff_dif = np.abs(state_vector_temp[ind]- target_arr_temp[ind])
187            sum += num_dif * coeff_dif
188
189        # add to weights_arr 
190        for j in np.arange(2**arg_dict["m"]):
191            ind = int(dec_to_bin(j,arg_dict["m"],'unsigned mag',nint=arg_dict["m"])+dec_to_bin(x_arr_temp[q],arg_dict["n"],'unsigned mag',nint=arg_dict["nint"]),2) 
192            WIM_weights_arr[ind]= sum    
193
194    # focus on outliers: double-weight on states 0.5 sigma or more above the mean
195    WIM_weights_arr += (WIM_weights_arr >= np.mean(WIM_weights_arr)+0.5 *np.std(WIM_weights_arr) ).astype(int) 
196    
197    # smoothen WIM weights 
198    WIM_weights_arr=np.exp(0.8 * WIM_weights_arr)
199
200    return WIM_weights_arr
201
202def train_QNN(n,m,L, seed, epochs,func_str,loss_str,meta, recover_temp, nint, mint, phase_reduce, train_superpos, real, repeat_params, WILL_p, WILL_q, delta, DIR):
203    r"""
204    Train a QCNN to perform function evaluation $\ket{j}\ket{0} \mapsto \ket{j}\ket{\Psi(j)}$.
205
206    The QCNN is generated using `pqcprep.pqc_tools.generate_network()`. 
207
208    Arguments:
209    ---
210    - **n** : *int*
211
212        Number of qubits in the input register. 
213
214    - **m** : *int*
215
216        Number of qubits in the target register. 
217
218    - **L** : *int*
219
220        Number of layers in the network. 
221
222    - **seed** : *int* 
223
224        Seed for random number generation. 
225
226    - **epochs** : *int* 
227
228        Number of training runs. 
229
230    - **func_str** : *str*
231
232        String specifying the function $\Psi$ to be evaluated. Must be a valid option for the argument `mode` of `pqcprep.psi_tools.psi()`. 
233
234    - **loss_str** : *str* 
235
236        String specifying the loss function minimised by the optimiser. Must be a valid option for the argument `loss_str` of `set_loss_func()`. 
237
238    - **meta** : *str*
239
240        String containing meta information to be included in output file names. 
241
242    - **recover_temp** : *boolean* 
243
244        If True, continue training from TEMP files (should they exist). If False and TEMP files exist they will be overwritten. 
245
246    - **nint** : *int*
247
248        Number of integer qubits in input register. 
249
250    - **mint** : *int*
251
252        Number of integer qubits in target register.  
253
254    - **phase_reduce** : *boolean* 
255
256        If True, reduce $\Psi(j)$ to the interval $[0, 2 \pi)$ i.e. perform the mapping $\Psi \to \Psi (\text{mod} \; 2 \pi)$. 
257
258    - **train_superpos** : *boolean*
259
260        If True, train on a superposition of input states. If False, train on randomly sampled individual input states. 
261
262    - **real** : *boolean*
263
264        If True, generate a network only involving CX and Ry rotations, resulting in real amplitudes. 
265
266    - **repeat_params** : *str*, *optional* 
267
268        Keep parameters fixed for different layer types, i.e. use the same parameter values for each instance of a layer type. 
269        Options are `None` (do not keep parameters fixed), `'CL'` (keep parameters fixed for convolutional layers), 
270        `'IL'` (keep parameters fixed for input layers), `'both'` (keep parameters fixed for both convolutional and input layers).    
271
272    - **WILL_p** : *float* 
273
274        The $p$ parameter of the WILL loss function, as described in `set_loss_func()`.
275
276    - **WILL_q** : *float* 
277
278        The $q$ parameter of the WILL loss function, as described in `set_loss_func()`.    
279
280    - **delta** : *float* 
281
282        Hyper-parameter controlling the sampling of input state coefficients when training in superposition (`train_superpos = True`). Must be 
283        between 0 and 1. `delta = 0` corresponds to coefficients fixed at $\\frac{1}{\\sqrt{2^n}}$ while `delta = 1` corresponds to coefficients randomly assuming values on the range $(0,1)$. 
284        Intermediate values of `delta` result in coefficinets being randomly sampled on an interval around $\\frac{1}{\\sqrt{2^n}}$, with the range of the interval 
285        determined by the value of `delta`.    
286
287    - **DIR** : *str*
288
289        Parent directory for output files.      
290
291    Returns:
292    ----
293
294    The output produced by the training is saved in binary `.npy` files in the directory `DIR/outputs` using naming convention `<TYPE>_<NAME_STR>.npy`
295    where `<NAME_STR>` is the name string produced by `pqcprep.file_tools.vars_to_name_str()` and `<TYPE>` is one of: 
296
297    - `weights` : file containing the QCNN weights determined by the optimiser;
298
299    - `loss` : file containing the loss value after each epoch; 
300
301    - `mismatch` : file containg the mismatch value after each epoch; 
302
303    - `grad` : file containing the  squared weight gradient norm after each epoch;  
304
305    - `vargrad` : file containing the variance of the weight gradients after each epoch.      
306
307    """
308    
309    # compress arguments into dictionary 
310    args =compress_args(n,m,L, seed, epochs,func_str,loss_str,meta,nint, mint, phase_reduce, train_superpos, real, repeat_params, WILL_p, WILL_q, delta)
311
312    # set precision strings 
313    if nint==None: nint=n
314    if mint==None: mint=m  
315    if phase_reduce: mint=0
316
317    # set seed for PRNG 
318    algorithm_globals.random_seed= seed
319    rng = np.random.default_rng(seed=seed)
320
321    # generate circuit and set up as QNN 
322    qc = generate_network(n,m,L, encode=not train_superpos, toggle_IL=True, initial_IL=True,input_Ry=train_superpos, real=real,repeat_params=repeat_params, wrap=False)
323                                                                                        
324    qnn = SamplerQNN(
325                    circuit=qc.decompose(),           
326                    sampler=Sampler(options={"shots": 10000, "seed": algorithm_globals.random_seed}),
327                    input_params=qc.parameters[:n], 
328                    weight_params=qc.parameters[n:],
329                    input_gradients=True
330                )
331              
332    # choose initial weights
333    recovered_k =0
334    if recover_temp:    
335        recover_labels=["weights", "mismatch", "loss", "grad", "vargrad"]
336        recover_paths={}
337        for k in np.arange(100,epochs, step=100):
338            for e in np.arange(len(recover_labels)):
339                file=os.path.join(DIR,"outputs", f"__TEMP{k}_{recover_labels[e]}{vars_to_name_str(args)}.npy")
340                recover_paths[recover_labels[e]]= (file if os.path.isfile(file) else None)
341
342                if recover_labels[e]=="weights" and os.path.isfile(file):
343                    recovered_k=k+1
344                     
345        if not None in list(recover_paths.values):
346            initial_weights=np.load(recover_paths["weights"])
347        else:
348            initial_weights =np.zeros(len(qc.parameters[n:])) # initialise parameters to zero   
349    else:
350        initial_weights =np.zeros(len(qc.parameters[n:])) # initialise parameters to zero   
351           
352    # initialise TorchConnector
353    model = TorchConnector(qnn, initial_weights)
354
355    # choose optimiser 
356    optimizer = Adam(model.parameters(), lr=0.01, betas=(0.7, 0.999), weight_decay=0.005) 
357                
358    # set up arrays to store training outputs 
359    if recover_temp and not None in list(recover_paths.values):
360        mismatch_vals=np.load(recover_paths["mismatch"])
361        loss_vals=np.load(recover_paths["loss"])
362        grad_vals=np.load(recover_paths["grad"])
363        var_grad_vals=np.load(recover_paths["var_grad"])
364    else:    
365        mismatch_vals = np.empty(epochs)
366        loss_vals = np.empty(epochs)
367        grad_vals = np.empty(epochs)
368        var_grad_vals = np.empty(epochs)
369
370    # generate x and f(x) values
371    pn =n - nint
372    pm =m - mint
373
374    if train_superpos:
375
376        # sample all basis states of input register and convert to binary 
377        x_arr = np.arange(0, 2**n)
378        x_arr_bin =[dec_to_bin(i,n,encoding="unsigned mag") for i in x_arr]
379
380        # apply function and reduce to phase value between 0 and 1 
381        fx_arr = [psi(i, mode=func_str) for i in x_arr]
382
383        if phase_reduce:
384            fx_arr = [np.modf(i/ (2* np.pi))[0] for i in fx_arr]
385
386        # convert fx_arr to binary at available target register precision 
387        fx_arr_bin = [dec_to_bin(i,m,nint=mint,encoding="unsigned mag") for i in fx_arr]
388
389        if np.max(fx_arr)> 2.**mint - 2.**(-pm) and mint != 0:
390            raise ValueError(f"Insufficient number of target (integer) qubits.")
391        
392        # get bit strings corresponding to target arrays and convert to indices
393        target_bin = [fx_arr_bin[i]+x_arr_bin[i] for i in x_arr]
394        target_ind = [bin_to_dec(i, encoding='unsigned mag') for i in target_bin]
395
396        # prepare target array 
397        target_arr = np.zeros(2**(n+m))
398        for k in target_ind:
399            target_arr[int(k)]=1
400
401    else:        
402        x_min = 0
403        x_max = 2.**nint - 2.**(-pn) 
404        x_arr = np.array(x_min + (x_max - x_min) *rng.random(size=epochs))
405        fx_arr = [psi(i, mode=func_str) for i in x_arr]
406
407        # reduce to phase value between 0 and 1:
408        if phase_reduce: 
409            fx_arr = [np.modf(i/ (2* np.pi))[0] for i in fx_arr]
410        
411        if np.max(fx_arr)> 2.**mint - 2.**(-pm) and mint != 0:
412            raise ValueError(f"Insufficient number of target (integer) qubits.")
413
414    # choose loss function 
415    criterion=set_loss_func(loss_str, args)
416    
417    # start training 
418    print(f"\n\nTraining started. Epochs: {epochs}. Input qubits: {n}. Target qubits: {m}. QCNN layers: {L}. \n")
419    start = time.time() 
420
421    warnings.filterwarnings("ignore", category=UserWarning)
422
423    for i in np.arange(epochs)[recovered_k:]:
424
425        if train_superpos == False:
426            # get input data
427            input = Tensor(binary_to_encode_param(dec_to_bin(x_arr[i],n,'unsigned mag',nint=nint))) 
428
429            # get target data 
430            target_arr = np.zeros(2**(n+m))
431            index = int(dec_to_bin(fx_arr[i],m,'unsigned mag',nint=mint)+dec_to_bin(x_arr[i],n,'unsigned mag',nint=nint),2)
432            target_arr[index]=1 
433            target=Tensor(target_arr)
434        else:
435
436            # generate random coefficients 
437            coeffs = np.array(np.pi / 2 * (1+delta *(2 *rng.random(size=n)-1)))
438            
439            # get input data 
440            input=Tensor(coeffs)
441
442            # get target data 
443            
444            target_ampl = np.empty(2**(n+m))
445
446            for j in np.arange(2**n):
447                for k in np.arange(2**m):
448                    ind =int(dec_to_bin(k,m)+dec_to_bin(j,n),2) 
449                    bin_n = dec_to_bin(j,n)[:: -1] # reverse bit order !
450                    val=1 
451
452                    for l in np.arange(n):
453                        val *= np.cos(coeffs[l]/2) if bin_n[l]=='0' else np.sin(coeffs[l]/2)
454
455                    target_ampl[ind] = val * target_arr[ind] 
456
457            target=Tensor(target_ampl**2)
458
459        # train model  
460        optimizer.zero_grad()
461
462        # apply loss function          
463        if loss_str=="WIM":
464            WIM_weights_arr=np.ones(2**(n+m)) if i==recovered_k else WIM_weights_arr
465            WIM_weights_tensor=Tensor(WIM_weights_arr)
466            loss =criterion(torch.sqrt(torch.abs(model(input))+1e-10), torch.sqrt(target), WIM_weights_tensor)    # add small number in sqrt to avoid zero grad !
467        else: 
468            loss = criterion(torch.sqrt(torch.abs(model(input))+1e-10), torch.sqrt(target))                       # add small number in sqrt to avoid zero grad !
469
470        # propagate gradients and recompute weights
471        loss.backward()
472        optimizer.step()
473
474        # save loss and grad for plotting 
475        loss_vals[i]=loss.item()
476        grad_vals[i]=np.sum(model.weight.grad.numpy()**2)
477        var_grad_vals[i]=np.std(model.weight.grad.numpy())**2
478        
479        # set up circuit with calculated weights
480        circ = generate_network(n,m,L, encode=not train_superpos, toggle_IL=True, initial_IL=True,input_Ry=train_superpos, real=real,repeat_params=repeat_params, wrap=False)
481        
482        with no_grad():
483            generated_weights = model.weight.detach().numpy()   
484        if train_superpos:
485            input_params = coeffs
486            params=np.concatenate((input_params, generated_weights))  
487        else:
488            input_params = binary_to_encode_param(dec_to_bin(x_arr[i],n,'unsigned mag',nint=nint))
489            params = np.concatenate((input_params, generated_weights))   
490        circ = circ.assign_parameters(params)
491
492        # get statevector 
493        state_vector = get_state_vec(circ)
494
495        # calculate fidelity and mismatch
496        target_state = target_ampl**2 if train_superpos else target_arr 
497        fidelity = np.abs(np.dot(np.sqrt(target_state),np.conjugate(state_vector)))**2
498        mismatch = 1. - np.sqrt(fidelity)
499        mismatch_vals[i]=mismatch
500
501        # set loss func weights for WIM
502        if loss_str=="WIM" and (i % 10 ==0) and (i >=1): WIM_weights_arr= set_WIM_weights(generated_weights, args)
503            
504        # temporarily save outputs every hundred iterations
505        temp_ind = epochs - 100 
506        
507        if recover_temp:
508            temp_ind = recovered_k -1
509
510        if (i % 100 ==0) and (i != 0) and (i != epochs-1): 
511
512            temp_labels=["weights", "mismatch", "loss", "grad", "vargrad"] 
513            temp_arrs=[generated_weights, mismatch_vals, loss_vals, grad_vals, var_grad_vals] 
514
515            for e in np.arange(len(temp_labels)):
516                # save temp file 
517                file=os.path.join(DIR,"outputs",f"__TEMP{i}_{temp_labels[e]}{vars_to_name_str(args)}")
518                np.save(file,temp_arrs[e])
519
520                # delete previous temp file 
521                old_file=os.path.join(DIR,"outputs",f"__TEMP{i-100}_{temp_labels[e]}{vars_to_name_str(args)}.npy")
522                os.remove(old_file) if os.path.isfile(old_file) else None
523
524            # make note of last created temp files
525            temp_ind = i   
526        
527        # print status
528        a = int(20*(i+1)/epochs)
529       
530        if i==recovered_k:
531            time_str="--:--:--.--"
532        elif i==epochs-1:
533            time_str="00:00:00.00"    
534        else:
535            if recover_temp:
536                    remaining = ((time.time() - start) / (i-recovered_k)) * (epochs - i)
537            else:
538                remaining = ((time.time() - start) / i) * (epochs - i)
539            mins, sec = divmod(remaining, 60)
540            hours, mins = divmod(mins, 60)
541            time_str = f"{int(hours):02}:{int(mins):02}:{sec:05.2f}"
542
543        prefix="\t" 
544        print(f"{prefix}[{u'█'*a}{('.'*(20-a))}] {100.*((i+1)/epochs):.2f}% ; Loss {loss_vals[i]:.2e} ; Mismatch {mismatch:.2e} ; ETA {time_str}", end='\r', file=sys.stdout, flush=True)
545        
546        
547    print(" ", flush=True, file=sys.stdout)
548    
549    warnings.filterwarnings("default", category=UserWarning)
550
551    elapsed = time.time()-start
552    mins, sec = divmod(elapsed, 60)
553    hours, mins = divmod(mins, 60)
554    time_str = f"{int(hours):02}:{int(mins):02}:{sec:05.2f}" 
555
556    # decompose circuit for gate count 
557    num_CX = dict(circ.decompose(reps=4).count_ops())["cx"]
558    num_gates = num_CX + dict(circ.decompose(reps=4).count_ops())["u"]
559    print(f"\nTraining completed in {time_str}. Number of weights: {len(generated_weights)}. Number of gates: {num_gates} (of which CX gates: {num_CX}). \n\n")
560
561    # delete temp files
562    temp_labels=["weights", "mismatch", "loss", "grad", "vargrad"]  
563    for i in np.arange(len(temp_labels)):
564        file=os.path.join(DIR,"outputs",f"__TEMP{temp_ind}_{temp_labels[i]}{vars_to_name_str(args)}.npy")
565        os.remove(file) if os.path.isfile(file) else None
566                            
567    # save outputs 
568    with no_grad():
569            generated_weights = model.weight.detach().numpy()
570    outputs= [generated_weights, mismatch_vals, loss_vals, grad_vals, var_grad_vals]
571    output_labels=["weights", "mismatch", "loss", "grad", "vargrad"]  
572    for i in np.arange(len(outputs)):
573        np.save(os.path.join(DIR,"outputs", f"{output_labels[i]}{vars_to_name_str(args)}"), outputs[i])      
574
575    return 0 
576
577def test_QNN(n,m,L,seed,epochs, func_str,loss_str,meta,nint,mint,phase_reduce,train_superpos,real,repeat_params,WILL_p, WILL_q,delta,DIR,verbose=True):   
578    """
579    Test performance of a QCNN trained for function evaluation with respect to different metrics. 
580
581    This requires the existence of an appropriate `weights_<NAME_STR>.npy` file (as produced by `train_QNN()`) in the directory `DIR/outputs`, where `<NAME_STR>` is the name string produced by `pqcprep.file_tools.vars_to_name_str()`. 
582
583    Arguments: 
584    ---
585
586    Same arguments as `train_QNN()`. See there for a description. 
587
588    Returns:
589    ---
590
591    The testing output produced is saved in binary `.npy` files in the directory `DIR/ampl_outputs` using naming convention `<TYPE>_<NAME_STR>.npy`
592    where `<NAME_STR>` is the name string produced by `pqcprep.file_tools.vars_to_name_str()` and `<TYPE>` is one of: 
593 
594    - `mismatch_by_state` : file containing the mismatch between the state produced by the network and the desired output state for each of the possible input 
595        register states. Contains a dictionary with the input bit strings as keys and the associated mismatch as values. 
596
597    - `phase` : file containing the phase function encoded by the network when the input register is in an equal superposition of input states. 
598
599    - `metrics` : file containing a dictionary with several metrics evaluating the performance of the network. These metrics are 
600        * `mu` : mean mismatch (mean of the data contained in `mismatch_by_state_<NAME_STR>.npy`); should be zero for ideal performance
601        * `sigma` : mismatch standard deviation (standard deviation of the data contained in `mismatch_by_state_<NAME_STR>.npy`);  should be zero for ideal performance
602        * `eps` : normalisation error on the state vector associated with the data contained in `phase_<NAME_STR>.npy`;  should be zero for ideal performance
603        * `chi` : mean absolute difference between the phase function contained in `phase_<NAME_STR>.npy` and the rounded desired phase function;  should be zero for ideal performance
604        * `omega` : a super-metric defined as `1/(mu + sigma + eps + chi)` ; should be maximal for ideal performance
605
606    """
607    # compress arguments into dictionary 
608    args =compress_args(n,m,L, seed, epochs,func_str,loss_str,meta,nint, mint, phase_reduce, train_superpos, real, repeat_params, WILL_p, WILL_q, delta)
609    name_str=vars_to_name_str(args)  
610
611    # set precision strings 
612    if nint==None: nint=n
613    if mint==None: mint=m  
614    if phase_reduce: mint=0                  
615
616    # load weights 
617    if os.path.isfile(os.path.join(DIR,"outputs",f"weights{name_str}.npy")):
618        weights = np.load(os.path.join(DIR,"outputs",f"weights{name_str}.npy"))
619    else:
620        raise ValueError("No appropriate QCNN weights could be found. Check the network configuration as well as the relevant directory.")    
621
622    # initialise array to store results 
623    mismatch = np.empty(2**n)
624    signs = np.empty(2**n)
625    
626    # iterate over input states 
627    x_arr = np.arange(2**n)
628    fx_arr = [psi(i, mode=func_str) for i in x_arr]
629
630    if phase_reduce: 
631        fx_arr = [np.modf(i/ (2* np.pi))[0] for i in fx_arr]
632
633    for i in x_arr:
634        
635        # prepare circuit 
636        enc=binary_to_encode_param(np.binary_repr(i,n))
637        params=np.concatenate((enc, weights))  
638
639        circ = generate_network(n,m,L, encode=True,toggle_IL=True, real=real,repeat_params=repeat_params)
640        circ = circ.assign_parameters(params) 
641
642        # get target array 
643        target_arr = np.zeros(2**(n+m))
644
645        index = int(dec_to_bin(fx_arr[i],m,'unsigned mag',nint=mint)+dec_to_bin(x_arr[i],n,'unsigned mag',nint=nint),2)
646        target_arr[index]=1 
647
648        # get statevector 
649        state_vector = get_state_vec(circ)
650
651        signs[i]=np.sign(np.sum(np.real(state_vector)*np.sqrt(target_arr)))
652
653        # calculate fidelity and mismatch 
654        fidelity = np.abs(np.dot(np.sqrt(target_arr),np.conjugate(state_vector)))**2
655        mismatch[i] = 1. - np.sqrt(fidelity) 
656            
657    # get phase target 
658    fx_arr_bin = [dec_to_bin(i,m, "unsigned mag", 0) for i in fx_arr]
659    phase_target =  np.array([bin_to_dec(i,"unsigned mag", mint) for i in fx_arr_bin])
660    if phase_reduce:
661        phase_target *= 2 * np.pi 
662
663    # get state vector for full phase extraction 
664    state_vec= full_encode(n,m, weights_A_str=None, weights_p_str=weights, L_A=None, L_p=L,real_p=real,repeat_params=repeat_params,full_state_vec=False, no_UA=True, operators="QRQ")
665    phase = phase_from_state(state_vec)
666
667    # calculate metrics
668    mu = np.mean(mismatch)
669    sigma = np.std(mismatch) 
670    eps = 1 - np.sum(np.abs(state_vec)**2)
671    chi = np.mean(np.abs(phase - phase_target))
672    omega= 1/(mu+sigma+eps+chi) 
673
674    # save outputs
675    full_dic = dict(zip(x_arr, mismatch)) 
676    np.save(os.path.join(DIR,"outputs",f"mismatch_by_state{name_str}.npy"), full_dic) 
677
678    metric_dic = {} 
679    metric_dic["mu"]=mu 
680    metric_dic["sigma"]=sigma   
681    metric_dic["eps"]=eps 
682    metric_dic["chi"]=chi
683    metric_dic["omega"]=omega  
684    np.save(os.path.join(DIR,"outputs",f"metrics{name_str}.npy"), metric_dic) 
685
686    np.save(os.path.join(DIR,"outputs",f"phase{name_str}.npy"), phase) 
687
688    if verbose:
689        print("Mismatch by input state:")
690        for i in x_arr:
691            print(f"\t{np.binary_repr(i,n)}:  {mismatch[i]:.2e} ({signs[i]})")
692        print("-----------------------------------")
693        print(f"Mu: \t{mu:.3e}") 
694        print(f"Sigma: \t{sigma:.3e}") 
695        print(f"Eps: \t{eps:.3e}")
696        print(f"Chi: \t{chi:.3e}") 
697        print(f"Omega: \t{omega:.3f}")
698        print("-----------------------------------")
699        print("")
700        print("")    
701
702    return 0 
703
704def ampl_train_QNN(n,L,x_min,x_max,seed, epochs,func_str,loss_str,meta, recover_temp, nint, repeat_params, DIR):
705    """
706
707    Train a QCNN to prepare an amplitude distribution: $\ket{0} \mapsto \sum_j A(j) \ket{j}$.
708
709    The QCNN is generated using `pqcprep.pqc_tools.A_generate_network()`. 
710
711    Arguments:
712    ---
713
714    - **n** : *int* 
715
716        Number of qubits in the register. 
717
718    - **L** : *int* 
719
720        Number of layers in the network. 
721
722    - **x_min** : *float* 
723
724        Minimum of function domain.         
725
726    - **x_max** : *float* 
727
728        Maximum of function domain.   
729
730    - **seed** : *int* 
731
732        Seed for random number generation. 
733
734    - **epochs** : *int* 
735
736        Number of training runs. 
737
738    - **func_str** : *str*
739
740        String specifying the function $A$ to be prepared. Must be a valid option for the argument `mode` of `pqcprep.psi_tools.A()`. 
741
742    - **loss_str** : *str* 
743
744        String specifying the loss function minimised by the optimiser. Must be a valid option for the argument `loss_str` of `set_loss_func()`. 
745
746    - **meta** : *str*
747
748        String containing meta information to be included in output file names. 
749
750    - **recover_temp** : *boolean* 
751
752        If True, continue training from TEMP files (should they exist). If False and TEMP files exist they will be overwritten. 
753
754    - **nint** : *int*
755
756        Number of integer qubits in the register.  
757
758    - **repeat_params** : *boolean* 
759
760        If True, keep parameters fixed for different layer types, i.e. use the same parameter values for each instance of a layer type.
761
762    - **DIR** : *str*
763
764        Parent directory for output files.     
765          
766
767    Returns:
768    ---
769
770    The output produced by the training is saved in binary `.npy` files in the directory `DIR/ampl_outputs` using naming convention `<TYPE>_<NAME_STR>.npy`
771    where `<NAME_STR>` is the name string produced by `pqcprep.file_tools.vars_to_name_str_ampl()` and `<TYPE>` is one of: 
772
773    - `weights` : file containing the QCNN weights determined by the optimiser;
774
775    - `state_vec` : file containing the statevector corresponding to the register after applying the QCNN; 
776
777    - `loss` : file containing the loss value after each training run; 
778
779    - `mismatch` : file containg the mismatch value after each training run.   
780 
781        
782    """
783
784    # compress arguments into dictionary 
785    args=compress_args_ampl(n,L,x_min,x_max,seed, epochs,func_str,loss_str,meta, nint, repeat_params)
786
787    # set seed for PRNG 
788    algorithm_globals.random_seed= seed
789    rng = np.random.default_rng(seed=seed)
790    
791    # generate circuit and set up as QNN 
792    qc = A_generate_network(n,L,repeat_params)
793    qnn = SamplerQNN(
794            circuit=qc.decompose(),            # decompose to avoid data copying (?)
795            sampler=Sampler(options={"shots": 10000, "seed": algorithm_globals.random_seed}),
796            weight_params=qc.parameters, 
797            input_params=[],  
798            input_gradients=False 
799        )
800    
801    # set precision 
802    if nint==None: nint=n   
803    
804    # choose initial weights
805    recovered_k =0
806
807    if recover_temp:    
808        recover_labels=["weights", "mismatch", "loss"]
809        recover_paths={}
810
811        for k in np.arange(100,epochs, step=100):
812            for e in np.arange(len(recover_labels)):
813                file=os.path.join(DIR,"ampl_outputs", f"__TEMP{k}_{recover_labels[e]}{vars_to_name_str_ampl(args)}.npy")
814                recover_paths[recover_labels[e]]= (file if os.path.isfile(file) else None)
815
816                if recover_labels[e]=="weights" and os.path.isfile(file):
817                    recovered_k=k+1      
818        
819        if not None in list(recover_paths.values):
820            initial_weights=np.load(recover_paths["weights"])
821        else:
822            initial_weights =np.zeros(len(qc.parameters)) # initialise parameters to zero   
823    
824    else:
825        initial_weights = np.zeros(len(qc.parameters)) # initialise parameters to zero   
826    
827    # initialise TorchConnector
828    model = TorchConnector(qnn, initial_weights)
829
830    # choose optimiser and loss function 
831    optimizer = Adam(model.parameters(), lr=0.01, betas=(0.7, 0.999), weight_decay=0.005) # Adam optimizer 
832    criterion=set_loss_func(loss_str, args, ampl=True)
833                    
834    # set up arrays to store training outputs 
835    if recover_temp and not None in list(recover_paths.values):
836        mismatch_vals=np.load(recover_paths["mismatch"])
837        loss_vals=np.load(recover_paths["loss"])
838    else:    
839        mismatch_vals = np.empty(epochs)
840        loss_vals = np.empty(epochs)
841
842    # calculate target and normalise 
843    dx = (x_max-x_min)/(2**n)
844    target_arr = np.array([A(i, mode=func_str) for i in np.arange(x_min,x_max, dx)])**2
845    target_arr = target_arr / np.sum(target_arr)
846    
847    # start training 
848    print(f"\n\nTraining started. Epochs: {epochs}. Input qubits: {n}. Function range: [{x_min},{x_max}]. QCNN layers: {L}. \n")
849    start = time.time() 
850
851    warnings.filterwarnings("ignore", category=UserWarning)
852
853    for i in np.arange(epochs)[recovered_k:]:
854
855        # get input data
856        input = Tensor([]) 
857
858        # get target data 
859        target=Tensor(target_arr)
860
861        # train model  
862        optimizer.zero_grad()
863        loss = criterion(torch.sqrt(torch.abs(model(input))+1e-10), torch.sqrt(target)) # adding 1e-10 to prevent taking sqrt(0) ??!!
864        loss.backward()
865        optimizer.step()
866
867        # save loss for plotting 
868        loss_vals[i]=loss.item()
869
870        # set up circuit with calculated weights
871        circ = A_generate_network(n,L, repeat_params)
872
873        with no_grad():
874            generated_weights = model.weight.detach().numpy()
875        
876        circ = circ.assign_parameters(generated_weights)    
877
878        # get statevector 
879        state_vector = get_state_vec(circ)
880        
881        # calculate fidelity and mismatch 
882        fidelity = np.abs(np.dot(np.sqrt(target_arr),np.conjugate(state_vector)))**2
883        mismatch = 1. - np.sqrt(fidelity)
884
885        # save mismatch for plotting 
886        mismatch_vals[i]=mismatch
887
888        # temporarily save outputs every hundred iterations
889        temp_ind = epochs - 100 
890        
891        if recover_temp:
892            temp_ind = recovered_k -1
893
894        if (i % 100 ==0) and (i != 0) and (i != epochs-1): 
895
896            temp_labels=["weights", "mismatch", "loss"] 
897            temp_arrs=[generated_weights, mismatch_vals, loss_vals] 
898
899            for e in np.arange(len(temp_labels)):
900                # save temp file 
901                file=os.path.join(DIR,"ampl_outputs",f"__TEMP{i}_{temp_labels[e]}{vars_to_name_str_ampl(args)}")
902                np.save(file,temp_arrs[e])
903
904                # delete previous temp file 
905                old_file=os.path.join(DIR,"ampl_outputs",f"__TEMP{i-100}_{temp_labels[e]}{vars_to_name_str_ampl(args)}.npy")
906                os.remove(old_file) if os.path.isfile(old_file) else None
907
908            # make note of last created temp files
909            temp_ind = i 
910         
911        # print status
912        a = int(20*(i+1)/epochs)
913
914        if i==recovered_k:
915            time_str="--:--:--.--"
916        elif i==epochs-1:
917            time_str="00:00:00.00"    
918        else:
919            if recover_temp:
920                remaining = ((time.time() - start) / (i-recovered_k)) * (epochs - i)
921            else:
922                remaining = ((time.time() - start) / i) * (epochs - i)
923            mins, sec = divmod(remaining, 60)
924            hours, mins = divmod(mins, 60)
925            time_str = f"{int(hours):02}:{int(mins):02}:{sec:05.2f}"
926
927        prefix="\t" 
928        print(f"{prefix}[{u'█'*a}{('.'*(20-a))}] {100.*((i+1)/epochs):.2f}% ; Loss {loss_vals[i]:.2e} ; Mismatch {mismatch:.2e} ; ETA {time_str}", end='\r', file=sys.stdout, flush=True)
929        
930    warnings.filterwarnings("default", category=UserWarning)
931
932    print(" ", flush=True, file=sys.stdout)
933
934    elapsed = time.time()-start
935    mins, sec = divmod(elapsed, 60)
936    hours, mins = divmod(mins, 60)
937    time_str = f"{int(hours):02}:{int(mins):02}:{sec:05.2f}" 
938
939    # decompose circuit for gate count 
940    num_CX = dict(circ.decompose(reps=4).count_ops())["cx"]
941    num_gates = num_CX + dict(circ.decompose(reps=4).count_ops())["u"]
942
943    print(f"\nTraining completed in {time_str}. Number of weights: {len(generated_weights)}. Number of gates: {num_gates} (of which CX gates: {num_CX}). \n\n")
944
945    # delete temp files
946    temp_labels=["weights", "mismatch", "loss"]  
947    for i in np.arange(len(temp_labels)):
948        file=os.path.join(DIR,"ampl_outputs",f"__TEMP{temp_ind}_{temp_labels[i]}{vars_to_name_str_ampl(args)}.npy")
949        os.remove(file) if os.path.isfile(file) else None             
950
951    # save outputs 
952    with no_grad():
953            generated_weights = model.weight.detach().numpy()
954    outputs= [generated_weights, mismatch_vals, loss_vals, state_vector]
955    output_labels=["weights", "mismatch", "loss", "statevec"]  
956    for i in np.arange(len(outputs)):
957        np.save(os.path.join(DIR,"ampl_outputs", f"{output_labels[i]}{vars_to_name_str_ampl(args)}"), outputs[i])         
958
959    return 0