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)