Running germ selection on multiple processors

The code below should be put into a script and run using mpiexec. It's primary function is to pass a MPI Comm object to pygsti.algorithms.germselection.build_up_breadth.

from __future__ import print_function
import time

import pygsti
import pygsti.construction as pc
from pygsti.construction import std2Q_XYICNOT
from pygsti.algorithms import germselection as germsel

from mpi4py import MPI
comm = MPI.COMM_WORLD

def do_greedy_germsel(target_model, forced_germs, candidate_counts,
                      seedStart, outFilename, comm):
    #candidate_counts is a dict of keys = germ lengths, values = # of germs at that length                                                            

    tStart = time.time()

    candidate_germs = []
    for i,(germLength, count) in enumerate(candidate_counts.items()):
        if count == "all upto":
            candidate_germs.extend( pc.list_all_circuits_without_powers_and_cycles(
                    target_model.operations.keys(), maxLength=germLength) )
        else:
            candidate_germs.extend( pc.list_random_circuits_onelen(
                    target_model.operations.keys(), germLength, count, seed=seedStart+i))

    available_germs = pygsti.tools.remove_duplicates( forced_germs + candidate_germs )
    print("%d available germs" % len(available_germs))
    germs = germsel.build_up_breadth(target_model, available_germs,
                     randomizationStrength=1e-3, numCopies=3, seed=1234,
                     opPenalty=10.0, scoreFunc='all', tol=1e-6, threshold=1e5,
                     pretest=False, force=forced_germs, verbosity=5, comm=comm, memLimit=0.5*(1024**3))

    if comm is None or comm.Get_rank() == 0:
        print("Germs (%d) = \n" % len(germs), "\n".join(map(str,germs)))
        print("Total time = %mdl" % (time.time()-tStart))
        pickle.dump(germs,open(outFilename,"wb"))
    return germs
                                                                                                                                         
#2Q case                                                                                                                                              
target_model = std2Q_XYICNOT.target_model()
forced_germs = pygsti.construction.circuit_list([(gl,) for gl in target_model.operations.keys()]) #singletons                                                                                      
candidate_counts = { 3:"all upto", 4:30, 5:20, 6:20, 7:20, 8:20} # germLength:num_candidates                                                          
seedStart = 4
do_greedy_germsel(target_model, forced_germs, candidate_counts,
                  seedStart, "germs_EXAMPLE.pkl", comm)

201 available germs
Starting germ set optimization. Lower score is better.
Memory estimate of 13.6 GB (0.5 GB limit) for all-Jac mode.
Memory estimate of 0.2 GB (0.5 GB limit) for single-Jac mode.
    Initial germ set computation Iter 1 of 6 Gii: 
    Initial germ set computation Iter 2 of 6 Gix: 
    Initial germ set computation Iter 3 of 6 Giy: 
    Initial germ set computation Iter 4 of 6 Gxi: 
    Initial germ set computation Iter 5 of 6 Gyi: 
    Initial germ set computation Iter 6 of 6 Gcnot: 
  Outer iteration: 1 of 1282 amplified, 6 germs
    Inner iter over candidate germs Iter 001 of 195 GiiGix: 
      Score: major=-114.0 minor=275.7550232560748, N: 194
    Inner iter over candidate germs Iter 002 of 195 GiiGiy: 
      Score: major=-114.0 minor=275.75224428270565, N: 194
    Inner iter over candidate germs Iter 003 of 195 GiiGxi: 
      Score: major=-114.0 minor=303.12985712455384, N: 194
    Inner iter over candidate germs Iter 004 of 195 GiiGyi: 
      Score: major=-114.0 minor=359.3300959231123, N: 194
    Inner iter over candidate germs Iter 005 of 195 GiiGcnot: 
      Score: major=-114.0 minor=276.26872683842225, N: 194
    Inner iter over candidate germs Iter 006 of 195 GixGiy: 
      Score: major=-114.0 minor=257.4426897879849, N: 194
    Inner iter over candidate germs Iter 007 of 195 GixGxi: 
      Score: major=-114.0 minor=341.5696468598034, N: 194
    Inner iter over candidate germs Iter 008 of 195 GixGyi: 
      Score: major=-114.0 minor=282.338674492408, N: 194
    Inner iter over candidate germs Iter 009 of 195 GixGcnot: 
      Score: major=-114.0 minor=577.5866443200878, N: 194
    Inner iter over candidate germs Iter 010 of 195 GiyGxi: 
      Score: major=-114.0 minor=294.80873540347477, N: 194
    Inner iter over candidate germs Iter 011 of 195 GiyGyi: 
      Score: major=-114.0 minor=274.01388120453754, N: 194
    Inner iter over candidate germs Iter 012 of 195 GiyGcnot: 
      Score: major=-114.0 minor=292.4187770509611, N: 194
    Inner iter over candidate germs Iter 013 of 195 GxiGyi: 
      Score: major=-114.0 minor=257.14469808616536, N: 194
    Inner iter over candidate germs Iter 014 of 195 GxiGcnot: 
      Score: major=-114.0 minor=295.19014734007106, N: 194
    Inner iter over candidate germs Iter 015 of 195 GyiGcnot: 
      Score: major=-114.0 minor=262.24440122325194, N: 194
    Inner iter over candidate germs Iter 016 of 195 GiiGiiGix: 
      Score: major=-104.0 minor=245.73088449107456, N: 194
    Inner iter over candidate germs Iter 017 of 195 GiiGiiGiy: 
      Score: major=-104.0 minor=236.10434931549185, N: 194
    Inner iter over candidate germs Iter 018 of 195 GiiGiiGxi: 
      Score: major=-104.0 minor=257.4811389850852, N: 194
    Inner iter over candidate germs Iter 019 of 195 GiiGiiGyi: 
      Score: major=-104.0 minor=295.32017448143534, N: 194
    Inner iter over candidate germs Iter 020 of 195 GiiGiiGcnot: 
      Score: major=-104.0 minor=274.607408708754, N: 194
    Inner iter over candidate germs Iter 021 of 195 GiiGixGix: 
      Score: major=-104.0 minor=453.2493918611686, N: 194
    Inner iter over candidate germs Iter 022 of 195 GiiGixGiy: 
      Score: major=-104.0 minor=275.1564063981899, N: 194
    Inner iter over candidate germs Iter 023 of 195 GiiGixGxi: 
      Score: major=-104.0 minor=304.47418773769414, N: 194
    Inner iter over candidate germs Iter 024 of 195 GiiGixGyi: 
      Score: major=-104.0 minor=278.44120228187916, N: 194
    Inner iter over candidate germs Iter 025 of 195 GiiGixGcnot: 
      Score: major=-104.0 minor=352.1721810679982, N: 194
    Inner iter over candidate germs Iter 026 of 195 GiiGiyGix: 
      Score: major=-104.0 minor=266.57020870563014, N: 194
    Inner iter over candidate germs Iter 027 of 195 GiiGiyGiy: 
      Score: major=-104.0 minor=448.32104717482, N: 194
    Inner iter over candidate germs Iter 028 of 195 GiiGiyGxi: 
      Score: major=-104.0 minor=297.3316824804599, N: 194
    Inner iter over candidate germs Iter 029 of 195 GiiGiyGyi: 
      Score: major=-104.0 minor=283.54615357575017, N: 194
    Inner iter over candidate germs Iter 030 of 195 GiiGiyGcnot: 
      Score: major=-104.0 minor=297.12854592193526, N: 194
    Inner iter over candidate germs Iter 031 of 195 GiiGxiGix: 
      Score: major=-104.0 minor=307.71199183761695, N: 194
    Inner iter over candidate germs Iter 032 of 195 GiiGxiGiy: 
      Score: major=-104.0 minor=307.7674841723881, N: 194
    Inner iter over candidate germs Iter 033 of 195 GiiGxiGxi: 
      Score: major=-104.0 minor=410.0988255344368, N: 194
    Inner iter over candidate germs Iter 034 of 195 GiiGxiGyi: 
      Score: major=-104.0 minor=284.3153092741726, N: 194
    Inner iter over candidate germs Iter 035 of 195 GiiGxiGcnot: 
      Score: major=-104.0 minor=303.8102280854164, N: 194
    Inner iter over candidate germs Iter 036 of 195 GiiGyiGix: 
      Score: major=-104.0 minor=283.5892520489328, N: 194
    Inner iter over candidate germs Iter 037 of 195 GiiGyiGiy: 
      Score: major=-104.0 minor=283.3226206509195, N: 194
    Inner iter over candidate germs Iter 038 of 195 GiiGyiGxi: 
      Score: major=-104.0 minor=273.5504913809546, N: 194
    Inner iter over candidate germs Iter 039 of 195 GiiGyiGyi: 
      Score: major=-104.0 minor=563.1049081486988, N: 194
    Inner iter over candidate germs Iter 040 of 195 GiiGyiGcnot: 
      Score: major=-104.0 minor=277.18869907206283, N: 194
    Inner iter over candidate germs Iter 041 of 195 GiiGcnotGix: 
      Score: major=-104.0 minor=352.3317810063291, N: 194
    Inner iter over candidate germs Iter 042 of 195 GiiGcnotGiy: 
      Score: major=-104.0 minor=298.1019089981339, N: 194
    Inner iter over candidate germs Iter 043 of 195 GiiGcnotGxi: 
      Score: major=-104.0 minor=309.6778947901517, N: 194
    Inner iter over candidate germs Iter 044 of 195 GiiGcnotGyi: 
      Score: major=-104.0 minor=279.3478112644729, N: 194
    Inner iter over candidate germs Iter 045 of 195 GiiGcnotGcnot: 
      Score: major=-104.0 minor=375.6548612782157, N: 194
    Inner iter over candidate germs Iter 046 of 195 GixGixGiy: 
      Score: major=-104.0 minor=327.9840714902066, N: 194
    Inner iter over candidate germs Iter 047 of 195 GixGixGxi: 
      Score: major=-104.0 minor=755.1272055918537, N: 194
    Inner iter over candidate germs Iter 048 of 195 GixGixGyi: 
      Score: major=-104.0 minor=362.97934478595647, N: 194
    Inner iter over candidate germs Iter 049 of 195 GixGixGcnot: 
      Score: major=-104.0 minor=532.4198103374531, N: 194
    Inner iter over candidate germs Iter 050 of 195 GixGiyGiy: 
      Score: major=-104.0 minor=293.57122809494956, N: 194
    Inner iter over candidate germs Iter 051 of 195 GixGiyGxi: 
      Score: major=-104.0 minor=313.06071440109775, N: 194
    Inner iter over candidate germs Iter 052 of 195 GixGiyGyi: 
      Score: major=-104.0 minor=291.50473798603565, N: 194
    Inner iter over candidate germs Iter 053 of 195 GixGiyGcnot: 
      Score: major=-104.0 minor=284.2173572664501, N: 194
    Inner iter over candidate germs Iter 054 of 195 GixGxiGiy: 

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
<ipython-input-1-4ee4fa17b7fa> in <module>()
     44 seedStart = 4
     45 do_greedy_germsel(target_model, forced_germs, candidate_counts,
---> 46                   seedStart, "germs_EXAMPLE.pkl", comm)

<ipython-input-1-4ee4fa17b7fa> in do_greedy_germsel(target_model, forced_germs, candidate_counts, seedStart, outFilename, comm)
     30                      randomizationStrength=1e-3, numCopies=3, seed=1234,
     31                      opPenalty=10.0, scoreFunc='all', tol=1e-6, threshold=1e5,
---> 32                      pretest=False, force=forced_germs, verbosity=5, comm=comm, memLimit=0.5*(1024**3))
     33 
     34     if comm is None or comm.Get_rank() == 0:

~/research/pyGSTi/packages/pygsti/algorithms/germselection.py in build_up_breadth(modelList, germsList, randomize, randomizationStrength, numCopies, seed, opPenalty, scoreFunc, tol, threshold, check, force, pretest, memLimit, comm, profiler, verbosity)
   1309                     worstScore = max( worstScore, compute_composite_germ_score(
   1310                         partialDerivDaggerDeriv=testDDD[None,:,:], initN=initN,
-> 1311                         **nonAC_kwargs))
   1312                     testDDDs.append(testDDD) #save in case this is a keeper
   1313 

~/research/pyGSTi/packages/pygsti/algorithms/germselection.py in compute_composite_germ_score(scoreFn, thresholdAC, initN, partialDerivDaggerDeriv, model, partialGermsList, eps, numGaugeParams, opPenalty, germLengths, l1Penalty)
    429 
    430     combinedDDD = _np.sum(partialDerivDaggerDeriv, axis=0)
--> 431     sortedEigenvals = _np.sort(_np.real(_nla.eigvalsh(combinedDDD)))
    432     observableEigenvals = sortedEigenvals[numGaugeParams:]
    433     N_AC = 0

/usr/local/lib/python3.7/site-packages/numpy/linalg/linalg.py in eigvalsh(a, UPLO)
   1126     t, result_t = _commonType(a)
   1127     signature = 'D->d' if isComplexType(t) else 'd->d'
-> 1128     w = gufunc(a, signature=signature, extobj=extobj)
   1129     return w.astype(_realType(result_t), copy=False)
   1130 

KeyboardInterrupt: 

Above is keyboard-interrupted on purpose, as this output was produced with a single processor and it would have taken a very long time.