Hierarchical Inference with N2JNet

Once our N2JNet is trained, we want to generate convergence predictions on individual sightlines and combine those predictions into a hierarchical inference of the population’s convergence statistics. We proceed similarly as with training, e.g.

$python infer.py nersc_config.yml

The “inference” section of the nersc_config.yml config file in the repo provides an example of how to configure inference. We take a look at it here.

First, we configure the properties of InferenceManager, a class that manages inference, with the device type, output directory (where all the inference results will be stored), and sampling seed.

inference_manager:
  device_type: 'cpu'
  out_dir: '/global/cscratch1/sd/jwp/n2j/apj_v4/seed1/inference_E10_N1000'
  seed: 1025

Then we specify the test healpixes as well as the subsampling distribution for the test sets. For instance, if we had n_subsample_test of 1000 and dist_name of ‘norm’ with dist_kwargs such that loc and scale were 0.04 and 0.005, respectively, we subsample 1,000 sightlines with a Gaussian distribution between with mean 0.04 and standard deviation 0.005. You can use any distribution supported by scipy.stats. The distributional parameters are the true hyperparameters governing the test population, which our hierarchical inference scheme will attempt to retrieve.

test_data:
  seed: 1
  batch_size: 1000
  test_hp: [10326, 9686]
  n_test: [50000, 50000]
  n_subsample_test: 1000
  dist_name: 'norm'
  dist_kwargs:
    loc: 0.04
    scale: 0.005

The summary statistics matching serves as a useful comparison. We provide a grid of closeness thresholds and a minimum number of matches for a threshold to be considered valid. In the case below, we choose the smallest threshold that resulted in more than 200 matches.

summary_stats:
  thresholds:
    N: [0, 1, 2, 4, 8, 16, 32, 64, 128]
    N_inv_dist: [0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
  min_matches: 200
# Replace with your own

The checkpoint path of the trained N2JNet must be passed.

checkpoint_path: '/global/cscratch1/sd/jwp/n2j/apj_v4/seed1/N2JNet_epoch=118_10-25-2021_08:10.mdl'

If we want to run hierarchical inference, we set run_mcmc as True. If we want to stop with generating individual predictions, this can be set as False. We can configure the MCMC, such as the number of “run” iterations, number of “burn” iterations, number of walkers, and the number of CPU cores.

run_mcmc: True
extra_mcmc_kwargs:
  n_run: 50
  n_burn: 20
  n_walkers: 10
  plot_chain: True
  clear: True
  n_cores: 1