def reconstruct_model(config): ham_matrix = qnnops.ising_hamiltonian( n_qubits=config.n_qubits, g=config.g, h=config.h) ev_str = re.sub(r' +', r' ', config.eigenvalues.replace('\n', '')).replace(' ', ',') try: ev = ast.literal_eval(ev_str) ground_energy = ev[0] except ValueError: print(f'Malformed string: {ev_str}') print(f'Parse the ground state energy manually.') ground_energy = float(ev_str.split(',')[0].lstrip('[')) def loss_fn(params): ansatz_state = qnnops.alternating_layer_ansatz( params, config.n_qubits, config.n_qubits, config.n_layers, config.rot_axis) return qnnops.energy(ham_matrix, ansatz_state) - ground_energy return loss_fn
def main(): args = get_arguments() project = args.project n_qubits, g, h = args.n_qubits, 2, 0 filters = dict(n_qubits=n_qubits, g=g, h=h) if args.n_layers_list is not None: print(f'Add [n_layers] filter: {args.n_layers_list}') filters['n_layers'] = {'$in': args.n_layers_list} if args.lr is not None: print(f'Add [lr] filter: {args.lr}') filters['lr'] = args.lr resdir = Path('results_hessian') opt_circuits = download_circuits(project, filters, resdir) ham_matrix = qnnops.ising_hamiltonian(n_qubits, g, h) print('Hessian spectrum') hess_fns = {} for cfg, fpath in opt_circuits: print(f'| computing hessian of {fpath}') params = jnp.load(fpath) circuit_name = f'Q{n_qubits}-L{cfg["n_layers"]}-R{cfg["rot_axis"]}' if circuit_name not in hess_fns: loss_fn = get_loss_fn( ham_matrix, n_qubits, cfg['n_layers'], cfg['rot_axis']) hess_fns[circuit_name] = jax.hessian(loss_fn) _, hess_eigvals = compute_hessian_eigenvalues(hess_fns[circuit_name], params) name = get_normalized_name(cfg) jnp.savez( resdir / f'{name}_all.npz', params=params, ham_matrix=ham_matrix, hess_spectrum=hess_eigvals, ) print(f'| ...plotting hessian spectrum and histogram') plot_spectrum(hess_eigvals, resdir / f'{name}_spectrum.pdf') plot_histogram(hess_eigvals, resdir / f'{name}_histogram.pdf')
choices=[1, 2], help='qnnops version (Default: 1)') args = parser.parse_args() seed = args.seed n_qubits, max_n_layers, rot_axis = args.n_qubits, args.max_n_layers, args.rot_axis block_size = args.n_qubits sample_size = args.sample_size g, h = args.g, args.h if not args.exp_name: args.exp_name = f'Q{n_qubits}R{rot_axis}BS{block_size} - g{g}h{h} - S{seed} - SN{sample_size}' expmgr.init(project='IsingBP', name=args.exp_name, config=args) # Construct the hamiltonian matrix of Ising model. ham_matrix = qnnops.ising_hamiltonian(n_qubits, g, h) rng = jax.random.PRNGKey(seed) # Set of random seeds for parameter sampling M = int(log2(max_n_layers)) for i in range(1, M + 1): n_layers = 2**i print(f'{n_qubits} Qubits & {n_layers} Layers ({i}/{M})') def loss(_params): ansatz_state = qnnops.alternating_layer_ansatz(_params, n_qubits, block_size, n_layers, rot_axis) return qnnops.energy(ham_matrix, ansatz_state) grad_fn = jax.grad(loss)
def main(): parser = argparse.ArgumentParser('Mode Connectivity') parser.add_argument('--n-qubits', type=int, metavar='N', required=True, help='Number of qubits') parser.add_argument('--n-layers', type=int, metavar='N', required=True, help='Number of alternating layers') parser.add_argument('--rot-axis', type=str, metavar='R', required=True, choices=['x', 'y', 'z'], help='Direction of rotation gates.') parser.add_argument('--g', type=float, metavar='M', required=True, help='Transverse magnetic field') parser.add_argument('--h', type=float, metavar='M', required=True, help='Longitudinal magnetic field') parser.add_argument('--n-bends', type=int, metavar='N', default=3, help='Number of bends between endpoints') parser.add_argument('--train-steps', type=int, metavar='N', default=int(1e3), help='Number of training steps. (Default: 1000)') parser.add_argument('--batch-size', type=int, metavar='N', default=8, help='Batch size. (Default: 8)') parser.add_argument('--lr', type=float, metavar='LR', default=0.05, help='Initial value of learning rate. (Default: 0.05)') parser.add_argument('--log-every', type=int, metavar='N', default=1, help='Logging every N steps. (Default: 1)') parser.add_argument( '--seed', type=int, metavar='N', required=True, help='Random seed. For reproducibility, the value is set explicitly.') parser.add_argument('--model-seeds', type=int, metavar='N', nargs=2, required=True, help='Random seed used for model training.') parser.add_argument('--exp-name', type=str, metavar='NAME', default=None, help='Experiment name.') parser.add_argument( '--scheduler-name', type=str, metavar='NAME', default='exponential_decay', help=f'Scheduler name. Supports: {qnnops.supported_schedulers()} ' f'(Default: constant)') parser.add_argument( '--params-start', type=str, metavar='PATH', help='A file path of a checkpoint where the curve starts from') parser.add_argument( '--params-end', type=str, metavar='PATH', help='A file path of a checkpoint where the curve ends') parser.add_argument('--no-jit', dest='use_jit', action='store_false', help='Disable jit option to loss function.') parser.add_argument('--no-time-tag', dest='time_tag', action='store_false', help='Omit the time tag from experiment name.') parser.add_argument('--quiet', action='store_true', help='Quite mode (No training logs)') args = parser.parse_args() n_qubits, n_layers = args.n_qubits, args.n_layers if args.exp_name is None: args.exp_name = f'Q{n_qubits}L{n_layers}_nB{args.n_bends}' if args.model_seeds is None: params_start, params_end = args.params_start, args.params_end else: params_start, params_end = download_checkpoints( 'checkpoints', n_qubits=n_qubits, n_layers=n_layers, lr=0.05, # fixed lr that we used in training seed={'$in': args.model_seeds}) print('Initializing project') expmgr.init('ModeConnectivity', args.exp_name, args) print('Loading pretrained models') w1 = jnp.load(params_start) w2 = jnp.load(params_end) expmgr.save_array('endpoint_begin.npy', w1) expmgr.save_array('endpoint_end.npy', w2) print('Constructing Hamiltonian matrix') ham_matrix = qnnops.ising_hamiltonian(n_qubits=n_qubits, g=args.g, h=args.h) print('Define the loss function') def loss_fn(params): ansatz_state = qnnops.alternating_layer_ansatz(params, n_qubits, n_qubits, n_layers, args.rot_axis) return qnnops.energy(ham_matrix, ansatz_state) find_connected_curve(w1, w2, loss_fn, n_bends=args.n_bends, train_steps=args.train_steps, lr=args.lr, scheduler_name=args.scheduler_name, batch_size=args.batch_size, log_every=1, seed=args.seed, use_jit=args.use_jit)