Test

cgs_vmc/normalizer.py

@@ -9,11 +9,13 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+from typing import Dict, Any
 
 import tensorflow as tf
 import numpy as np
 
 import wavefunctions
+import graph_builders
 
 
 def normalize_wavefunction(
@@ -33,7 +35,7 @@ def normalize_wavefunction(
       wavefunction: Wavefunction model to normalize.
       system_configs: Configurations on which system can be evaluated.
       update_config: Operation that samples new `system_configs`.
-      session: Active session in which th wavefunction is normalized.
+      session: Active session in which the wavefunction is normalized.
       value: What is the max value to aim for.
       normalization_iterations: Number of batches to process for normalization.
   """
@@ -47,3 +49,73 @@ def normalize_wavefunction(
   for _ in range(normalization_iterations):
     session.run(update_config)
     session.run(update_norm_on_batch)
+
+
+def build_normalization_ops(
+    wavefunction: wavefunctions.Wavefunction,
+    hparams: tf.contrib.training.HParams,
+    shared_resources: Dict[graph_builders.ResourceName, tf.Tensor],
+) -> Any:
+  """Generates operations that the fix range of wf amplitudes.
+
+  Construct computational graph to normalize the wave function magnitude by
+  dividing the largest magnitude from monte_carlo sampling.
+
+  Args:
+      wavefunction: Wavefunction model to normalize.
+      hparams: Class holding hyperparameters of the wavefunction ansatzs.
+      shared_resources: Resources sharable among different modules.
+
+  Returns:
+      Max magnitude of wave function from monte carlo sampling and monte carlo
+      step mc_step.
+  """
+
+  batch_size = hparams.batch_size
+  n_sites = hparams.num_sites
+  configs = graph_builders.get_configs(shared_resources, batch_size, n_sites)
+  mc_step = graph_builders.get_monte_carlo_sampling(
+      shared_resources, configs, wavefunction)[0]
+  psi_value = wavefunction(configs)
+  max_value = tf.reduce_max(tf.square(psi_value))
+
+  return max_value, mc_step
+
+
+def run_normalization_ops(
+    max_value: np.float,
+    mc_step: tf.Tensor,
+    wavefunction: wavefunctions.Wavefunction,
+    session: tf.Session,
+    hparams: tf.contrib.training.HParams,
+) -> wavefunctions.Wavefunction:
+  """Executes operations that the fix range of wf amplitudes.
+
+  Run computational graph to normalize the wave function magnitude by
+  dividing the largest magnitude from monte_carlo sampling.
+
+  Args:
+      max_value: Max value of wavefunction magnitude square estimated,
+      mc_step: Monte Carlo step,
+      wavefunction: Wavefunction model to normalize.
+      session: Active session in which the wavefunction is normalized.
+      hparams: Class holding hyperparameters of the wavefunction ansatzs.
+
+  Returns:
+      Wavefunction normalized by the max magnitude from monte carlo sampling
+      and monte carlo.
+  """
+
+  n_sites = hparams.num_sites
+  normalization = hparams.normalization
+
+  if normalization == 'monte_carlo':
+    current_max = 0.
+    for _ in range(hparams.num_equilibration_sweeps * n_sites):
+      session.run(mc_step)
+      current_max = max(current_max, session.run(max_value))
+    scale = 1.0 / np.sqrt(current_max)
+  elif normalization == 'space_scale':
+    scale = np.sqrt(2 ** n_sites)
+
+  return wavefunction * scale

cgs_vmc/run_energy_evaluation.py

@@ -25,6 +25,10 @@
     'checkpoint_dir', '',
     'Full path to the checkpoint directory.')
 
+flags.DEFINE_string(
+    'bonds_file', '',
+    'Name of the file that contains bonds of the Hamiltonian.')
+
 flags.DEFINE_string(
     'output_file', '',
     'Full path to the file where to save the results.')
@@ -48,7 +52,7 @@ def main(argv):
   n_sites = hparams.num_sites
 
   # TODO(dkochkov) make a more comprehensive Hamiltonian construction method
-  bonds_file_path = os.path.join(FLAGS.checkpoint_dir, 'J.txt')
+  bonds_file_path = os.path.join(FLAGS.checkpoint_dir, FLAGS.bonds_file)
   heisenberg_jx = FLAGS.heisenberg_jx
   if os.path.exists(bonds_file_path):
     heisenberg_data = np.genfromtxt(bonds_file_path, dtype=int)

cgs_vmc/run_supervised_training.py

@@ -36,7 +36,7 @@
     'Indicator to resotre variables from the latest checkpoint')
 
 flags.DEFINE_string(
-    'wavefunction_type', 'fully_connected',
+    'wavefunction_type', '',
     'Network architecture to train. Available architectures are listed in '
     'wavefunctions.WAVEFUNCTION_TYPES dict. and '
     'wavefunctions.build_wavefunction() function.')
@@ -88,6 +88,7 @@ def main(argv):
   hparams.set_hparam('basis_file_path', FLAGS.basis_file_path)
   hparams.set_hparam('num_epochs', FLAGS.num_epochs)
   hparams.set_hparam('wavefunction_type', FLAGS.wavefunction_type)
+  hparams.set_hparam('wavefunction_optimizer_type', FLAGS.optimizer)
   hparams.parse(FLAGS.hparams)
   hparams_path = os.path.join(hparams.checkpoint_dir, 'hparams.pbtxt')
 
@@ -132,14 +133,20 @@ def main(argv):
     latest_checkpoint = tf.train.latest_checkpoint(hparams.checkpoint_dir)
     checkpoint_saver.restore(session, latest_checkpoint)
 
+  training_metrics_file = os.path.join(
+      hparams.checkpoint_dir, 'supervised_loss.txt')
   for epoch_number in range(FLAGS.num_epochs):
-    wavefunction_optimizer.run_optimization_epoch(
+    metrics_record = wavefunction_optimizer.run_optimization_epoch(
         train_ops, session, hparams, epoch_number)
     if epoch_number % FLAGS.checkpoint_frequency == 0:
       checkpoint_name = 'model_after_{}_epochs'.format(epoch_number)
       save_path = os.path.join(hparams.checkpoint_dir, checkpoint_name)
       checkpoint_saver.save(session, save_path)
 
+    metrics_file_output = open(training_metrics_file, 'a')
+    metrics_file_output.write('{}\n'.format(metrics_record))
+    metrics_file_output.close()
+
   if FLAGS.generate_vectors:
     vector_generator = evaluation.VectorWavefunctionEvaluator()
     eval_ops = vector_generator.build_eval_ops(

cgs_vmc/run_training.py

@@ -21,6 +21,9 @@
 flags.DEFINE_string(
     'checkpoint_dir', '',
     'Full path to the checkpoint directory.')
+flags.DEFINE_string(
+    'bonds_file', '',
+    'Name of the file that contains bonds of the Hamiltonian.')
 flags.DEFINE_integer(
     'num_sites', 24,
     'Number of sites in the system.')
@@ -74,9 +77,9 @@ def main(argv):
   """Runs wavefunction optimization.
 
   This pipeline optimizes wavefunction specified in flags on a Marshal sign
-  included Heisenberg model. Bonds should be specified in the file J.txt in
-  checkpoint directory, otherwise will default to 1D PBC system. For other
-  tunable parameters see flags description.
+  included Heisenberg model. Bonds should be specified in the bonds_file
+  (e.g.J.txt) in checkpoint directory, otherwise will default to 1D PBC
+  system. For other tunable parameters see flags description.
   """
   del argv  # Not used.
   n_sites = FLAGS.num_sites
@@ -100,7 +103,7 @@ def main(argv):
   with tf.gfile.GFile(hparams_path, 'w') as file:
     file.write(str(hparams.to_proto()))
 
-  bonds_file_path = os.path.join(FLAGS.checkpoint_dir, 'J.txt')
+  bonds_file_path = os.path.join(FLAGS.checkpoint_dir, FLAGS.bonds_file)
   heisenberg_jx = FLAGS.heisenberg_jx
   if os.path.exists(bonds_file_path):
     heisenberg_data = np.genfromtxt(bonds_file_path, dtype=int)

cgs_vmc/training.py

@@ -164,10 +164,10 @@ def build_opt_ops(
         shared_resources, configs, wavefunction)
 
     psi = wavefunction(configs)
-    psi_target = target_wavefunction(configs)
+    psi_target = target_wavefunction(configs) * hparams.scale
 
     loss = tf.reduce_mean(
-        tf.squared_difference(psi, psi_target * np.sqrt(2**n_sites)) /
+        tf.squared_difference(psi, psi_target) /
         (tf.square(tf.stop_gradient(psi)))
     )
     opt_v = wavefunction.get_trainable_variables()
@@ -195,21 +195,27 @@ def run_optimization_epoch(
       session: tf.Session,
       hparams: tf.contrib.training.HParams,
       epoch_number: int,
-  ):
+  )-> np.float32:
     """Runs training epoch by executing `train_ops` in `session`.
 
     Args:
       train_ops: Training operations returned by `build_opt_ops` method.
       session: Active session where to run a training epoch.
       hparams: Hyperparameters of the optimization procedure.
       epoch_number: Number of epoch.
+
+    Returns:
+      Current wavefunctin fidelity loss estimate.
     """
     del epoch_number  # not used by SupervisedWavefunctionOptimizer.
+    loss_values = []
     for _ in range(hparams.num_batches_per_epoch):
       for _ in range(hparams.num_monte_carlo_sweeps * hparams.num_sites):
         session.run(train_ops.mc_step)
-      session.run(train_ops.apply_gradients)
+      _, loss = session.run([train_ops.apply_gradients, train_ops.metrics])
+      loss_values.append(loss)
     session.run(train_ops.epoch_increment)
+    return np.mean(loss_values)
 
 
 class BasisIterationSWO():
@@ -243,18 +249,18 @@ def build_opt_ops(
         hparams.basis_file_path, [tf.float32 for _ in range(n_sites)],
         header=False, field_delim=' ')
     basis_dataset = basis_dataset.map(lambda *x: tf.convert_to_tensor(x))
-    shuffle_batch = scipy.special.binomi(n_sites, n_sites / 2)
+    shuffle_batch = scipy.special.binom(n_sites, n_sites / 2)
     basis_dataset = basis_dataset.shuffle(shuffle_batch)
     basis_dataset = basis_dataset.batch(batch_size)
     basis_dataset = basis_dataset.repeat()
     config_iterator = basis_dataset.make_one_shot_iterator()
     configs = config_iterator.get_next() * 2. - 1.
 
     psi = wavefunction(configs)
-    psi_target = target_wavefunction(configs)
+    psi_target = target_wavefunction(configs) * hparams.scale
 
     loss = tf.reduce_mean(
-        tf.squared_difference(psi, psi_target * np.sqrt(2 ** n_sites)))
+        tf.squared_difference(psi, psi_target))
     opt_v = wavefunction.get_trainable_variables()
     optimizer = create_sgd_optimizer(hparams)
     train_step = optimizer.minimize(loss, var_list=opt_v)
@@ -280,19 +286,25 @@ def run_optimization_epoch(
       session: tf.Session,
       hparams: tf.contrib.training.HParams,
       epoch_number: int,
-  ):
+  )-> np.float32:
     """Runs training epoch by executing `train_ops` in `session`.
 
     Args:
       train_ops: Training operations returned by `build_opt_ops` method.
       session: Active session where to run a training epoch.
       hparams: Hyperparameters of the optimization procedure.
       epoch_number: Number of epoch.
+
+    Returns:
+      Current wavefunctin fidelity loss estimate.
     """
     del epoch_number  # not used by SupervisedWavefunctionOptimizer.
+    loss_values = []
     for _ in range(hparams.num_batches_per_epoch):
-      session.run(train_ops.apply_gradients)
+      _, loss = session.run([train_ops.apply_gradients, train_ops.metrics])
+      loss_values.append(loss)
     session.run(train_ops.epoch_increment)
+    return np.mean(loss_values)
 
 
 class LogOverlapSWO():
@@ -449,7 +461,7 @@ def build_opt_ops(
 
     configs = tf.concat([psi_configs, target_configs], axis=0)
     psi = wavefunction(configs)
-    psi_target = target_wavefunction(configs) * np.sqrt(2 ** n_sites)
+    psi_target = target_wavefunction(configs) * hparams.scale
 
     # # A version of accounting for sampling bias.
     # psi_no_grad = tf.stop_gradient(psi)
@@ -486,21 +498,27 @@ def run_optimization_epoch(
       session: tf.Session,
       hparams: tf.contrib.training.HParams,
       epoch_number: int,
-  ):
+  )-> np.float32:
     """Runs training epoch by executing `train_ops` in `session`.
 
     Args:
       train_ops: Training operations returned by `build_opt_ops` method.
       session: Active session where to run a training epoch.
       hparams: Hyperparameters of the optimization procedure.
       epoch_number: Number of epoch.
+
+    Returns:
+      Current wavefunctin fidelity loss estimate.
     """
     del epoch_number  # not used by SupervisedWavefunctionOptimizer.
+    loss_values = []
     for _ in range(hparams.num_batches_per_epoch):
       for _ in range(hparams.num_monte_carlo_sweeps * hparams.num_sites):
         session.run(train_ops.mc_step)
-      session.run(train_ops.apply_gradients)
+      _, loss = session.run([train_ops.apply_gradients, train_ops.metrics])
+      loss_values.append(loss)
     session.run(train_ops.epoch_increment)
+    return np.mean(loss_values)
 
 
 class EnergyGradientOptimizer(WavefunctionOptimizer):
@@ -809,7 +827,7 @@ def build_opt_ops(
 
     wf_omega = copy.deepcopy(wavefunction)  # building supervisor wavefunction.
     beta = tf.constant(hparams.time_evolution_beta, dtype=tf.float32)
-    beta2 = tf.constant(hparams.time_evolution_befta ** 2, dtype=tf.float32)
+    beta2 = tf.constant(hparams.time_evolution_beta ** 2, dtype=tf.float32)
     psi_omega = wf_omega(configs)
     h_psi_omega = hamiltonian.apply_in_place(wf_omega, configs, psi_omega)
     h_psi_omega_beta = h_psi_omega * beta

cgs_vmc/utils.py

@@ -70,7 +70,8 @@ def create_hparams(**kwargs: Any) -> tf.contrib.training.HParams:
     batch_size: Number of configurations in the batch (num of markov chains).
     num_batches_per_epoch: Number of batches sampled per epoch.
     time_evolution_beta: Imaginary time evolution step for ITSWO.
-    learning_rates: Learning rates used for SGD
+    scale: Scale factor for wavefunction magnitude in supervised learning.
+    learning_rates: Learning rates used for SGD.
     learning_rate_stops: Epoch stops when learning rate is adjusted.
     optimizer: SGD optimizer used to training.
     beta2: Second momenta for optimizer
@@ -137,6 +138,7 @@ def create_hparams(**kwargs: Any) -> tf.contrib.training.HParams:
       batch_size=200,
       num_batches_per_epoch=50,
       time_evolution_beta=0.12,
+      scale=1.0,
 
       learning_rates=[1e-3, 1e-4, 2e-5, 1e-5],
       learning_rate_stops=[300, 600, 1000],

cgs_vmc/wavefunctions.py

@@ -1142,10 +1142,12 @@ def from_hparams(
       name: str = ''
   ) -> 'Wavefunction':
     """Constructs an instance of a class from hparams."""
+    dir_path = hparams.checkpoint_dir
+    adjacency_list_path = os.path.join(dir_path, hparams.adjacency_list_path)
     gnn_params = {
         'num_layers': hparams.num_conv_layers,
         'num_filters': hparams.num_conv_filters,
-        'adj': np.genfromtxt(hparams.adjacency_list_path, dtype=int),
+        'adj': np.genfromtxt(adjacency_list_path, dtype=int),
         'output_activation': layers.NONLINEARITIES[hparams.output_activation],
         'nonlinearity': layers.NONLINEARITIES[hparams.nonlinearity],
     }