Adding optional mpirun command to pass to python interface.

examples/pythoninterface/example_spinchain.py

@@ -0,0 +1,168 @@
+#  Quandary's python interface functions are defined in /path/to/quandary/quandary.py. Import them here. 
+from quandary import * 
+np.random.seed(9001)
+
+### Define some case's coefficients
+def get_deterministic_params(N, h_amp, U_amp, J_amp):
+    h = h_amp * np.ones(N) 
+    U = U_amp * np.ones(N-1)
+    J = J_amp * np.ones(N-1)
+    return h, U, J
+
+def get_xx_params(N, J_amp):
+    h = np.zeros(N)
+    U = np.zeros(N-1)
+    J = J_amp * np.ones(N-1)
+    return h, U, J
+
+def get_disordered_xx_params(N, h_amp, J_amp):
+    h = np.random.uniform(-h_amp, h_amp, N)
+    U = np.zeros(N)
+    J = J_amp * np.ones(N)
+    return h, U, J
+
+def get_xxz_params(N, U_amp, J_amp):
+    h = np.zeros(N)
+    U = U_amp * np.ones(N)
+    J = J_amp * np.ones(N)
+    return h, U, J
+
+def get_disordered_xxz_params(N, h_amp, U_amp, J_amp):
+    h = np.random.uniform(-h_amp, h_amp, N)
+    U = U_amp * np.ones(N)
+    J = J_amp * np.ones(N)
+    return h, U, J
+
+def mapCoeffs_SpinChainToQuandary(N:int, h:list, U:list, J:list):
+	"""
+	Map spin chain coefficents J, U and h onto Quandary coefficients
+
+	Parameters:
+	------------
+ 	N	: int			:  number of spin sites
+ 	J	: float			:  J-coefficient per linear chain coupling: J[i] couples i and i+1
+ 	U	: list(float)	:  U-coefficient per linear chain coupling: J[i] couples i and i+1
+ 	h	: list(float)	:  h-coefficient per spin site
+
+	Output:
+	--------
+ 	freq01	  :  01-transition frequency per qubit (omega_j)
+ 	crossker  :  crossker coupling (linear chain) (xi_ij)
+ 	Jkl       :  dipole-dipole coupling (linear chain) (J_ij)
+	"""
+
+	# 01 transition frequencies [GHz] per site (omega_q) 
+	freq01 = np.zeros(N)
+	for i in range(1,N-1):
+		freq01[i] = (-2*h[i] -2*U[i] -2*U[i-1]) / (2*np.pi)
+	freq01[0]   = (-2*h[0]   - 2*U[0] ) / (2*np.pi)
+	freq01[N-1] = (-2*h[N-1] - 2*U[N-2]) / (2*np.pi)
+	# Jkl and Xi term [GHz] (Format [0<->1, 0<->2, ..., 1<->2, ... ])
+	Jkl=[]
+	crosskerr=[]
+	couplingID = 0
+	for i in range(N):
+		for j in range(i+1, N):
+			if j == i+1:  # linear chain coupling
+				valJ = - 2*J[couplingID] / (2*np.pi)
+				valC = - 4*U[couplingID] / (2*np.pi)
+			else:
+				valJ = 0.0
+				valC = 0.0
+			Jkl.append(valJ)
+			crosskerr.append(valC)
+		couplingID += 1
+	return freq01, crosskerr, Jkl
+#####
+
+# Specify the number of spin sites
+N = 8
+
+# Get spin chain case coefficients
+U_amp = 1.0 
+J_amp = 1.0 
+h_amp = 1.0
+h, U, J = get_disordered_xx_params(N, U_amp, J_amp)
+
+# Specify the initial state (list of 0 or 1 per site). Here: domain wall |111...000>
+initstate= np.zeros(N, dtype=int)
+for i in range(int(N/2)):
+   initstate[i] = 1 
+
+# Set the simulation duration and step size
+T = 10.0
+# T = 4.0
+dT = 0.01 	# Double check if this is small enough (e.g. by re-running with smaller dT and compare results)
+
+# Quandary run options
+verbose = True
+ncores=8   				# Numbers of cores for Quandary 
+mpi_exec="mpirun -np" 	# MPI executable, e.g. "srun -n" for LC 
+
+# Prepare Quandary
+initcondstr = "pure, "
+for e in initstate:
+	initcondstr += str(e) +", "
+freq01, crosskerr, Jkl = mapCoeffs_SpinChainToQuandary(N, h, U, J)
+quandary = Quandary(Ne=[2 for _ in range(N)], Ng=[0 for _ in range(N)], freq01=freq01, rotfreq=np.zeros(N), crosskerr=crosskerr, Jkl=Jkl, initialcondition=initcondstr, T=T, dT=dT, initctrl_MHz=0.0, carrier_frequency=[[0.0] for _ in range(N)], verbose=verbose)
+
+# Storage for averaged magnetization. Matrix rows = sites, cols = time
+magnet = np.zeros((N,quandary.nsteps+1))
+# Storage for averaged domain wall spread
+nhalf = np.zeros(quandary.nsteps+1)
+
+# Number of samples for h
+# nsamples = 10
+nsamples = 1
+
+# Iterate over randomized h and run quandary
+for isample in range(nsamples):
+
+	# Sample new set of coefficients and map to Quandary's coeffs
+	h, U, J = get_disordered_xx_params(N, U_amp, J_amp)
+	quandary.freq01, quandary.crosskerr, quandary.Jkl = mapCoeffs_SpinChainToQuandary(N, h, U, J)
+	quandary.update() 
+
+	# Run forward simulation
+	datadir = "./N"+str(N)+"_sample" + str(isample)+"_run_dir_parallel"  
+	t, pt, qt, infidelity, expectedEnergy, population = quandary.simulate(datadir=datadir, maxcores=ncores, mpi_exec=mpi_exec)
+
+	# Compute magnetization from expected Energy (-2*expEnergy + 1) and add to average
+	for isite in range(N):
+		for nt in range(len(t)):
+			exp = expectedEnergy[isite][0][nt] ## expected energy for this site and time
+			magnet_this = -2.0*exp + 1.0
+			magnet[isite,nt] += magnet_this / nsamples
+
+	# Compute domain wall spread (N_1/2) from expected Energy and add to average
+	for nt in range(len(t)): # iterate over time
+		nhalf_i = 0.0
+		for isite in range(int(N/2)): # iterate over first half of the sites
+			nhalf_i -= expectedEnergy[isite][0][nt] 
+		nhalf_i += int(N/2)
+		nhalf[nt] += nhalf_i / nsamples
+
+# # Plot sigmaz dynamics of one spin site
+# isite = 0
+# myexp = [-2*e+1 for e in expectedEnergy[isite][0]]
+# plt.plot(t, myexp)
+# plt.show()
+
+# Plot heatmap of magnetization
+fig, ax = plt.subplots(figsize=(6,4))
+mycmap = plt.get_cmap('coolwarm')
+plt.imshow(magnet, interpolation='none', aspect='auto', cmap=mycmap)
+plt.colorbar()
+plt.title(r"Heat Map of Spin Chain $\langle \sigma^z_j \rangle$")
+plt.xlabel("Time-step index")
+plt.ylabel("Spin Site Index $j$")
+plt.show()
+
+# Plot domain wall spread:
+plt.plot(t, nhalf)
+plt.xlabel("Time-step index")
+plt.ylabel("N_1/2")
+plt.grid(True)
+plt.legend()
+plt.show()
+

quandary.py

@@ -199,9 +199,7 @@ def __post_init__(self):
             self.initctrl_MHz = [10.0 for _ in range(len(self.Ne))]
         if len(self.Hsys) > 0 and not self.standardmodel: # User-provided Hamiltonian operators 
             self.standardmodel=False   
-        else: # Using standard Hamiltonian model
-            Ntot = [sum(x) for x in zip(self.Ne, self.Ng)]
-            self.Hsys, self.Hc_re, self.Hc_im = hamiltonians(N=Ntot, freq01=self.freq01, selfkerr=self.selfkerr, crosskerr=self.crosskerr, Jkl=self.Jkl, rotfreq=self.rotfreq, verbose=self.verbose)
+        else: # Using standard Hamiltonian model. Set it up in python only if needed, i.e. only if dT or the carrier wave frequencies need to be computed.
             self.standardmodel=True
         if len(self.targetstate) > 0:
             self.optim_target = "file"
@@ -226,6 +224,9 @@ def __post_init__(self):
 
         # Estimate the number of required time steps
         if self.dT < 0:
+            if self.standardmodel==True: # set up the standard Hamiltonian first
+                Ntot = [sum(x) for x in zip(self.Ne, self.Ng)]
+                self.Hsys, self.Hc_re, self.Hc_im = hamiltonians(N=Ntot, freq01=self.freq01, selfkerr=self.selfkerr, crosskerr=self.crosskerr, Jkl=self.Jkl, rotfreq=self.rotfreq, verbose=self.verbose)
             self.nsteps = estimate_timesteps(T=self.T, Hsys=self.Hsys, Hc_re=self.Hc_re, Hc_im=self.Hc_im, maxctrl_MHz=self.maxctrl_MHz, Pmin=self.Pmin)
             self.dT = self.T/self.nsteps
         else:
@@ -250,6 +251,10 @@ def __post_init__(self):
         if self.spline_order == 0 and len(self.carrier_frequency) == 0:
             self.carrier_frequency = [[0.0] for _ in range(len(self.freq01))]
         if len(self.carrier_frequency) == 0: 
+            # set up the standard Hamiltonian first, if needed
+            if self.standardmodel==True and len(self.Hsys<=0):
+                Ntot = [sum(x) for x in zip(self.Ne, self.Ng)]
+                self.Hsys, self.Hc_re, self.Hc_im = hamiltonians(N=Ntot, freq01=self.freq01, selfkerr=self.selfkerr, crosskerr=self.crosskerr, Jkl=self.Jkl, rotfreq=self.rotfreq, verbose=self.verbose)
             self.carrier_frequency, _ = get_resonances(Ne=self.Ne, Ng=self.Ng, Hsys=self.Hsys, Hc_re=self.Hc_re, Hc_im=self.Hc_im, rotfreq=self.rotfreq, verbose=self.verbose, cw_amp_thres=self.cw_amp_thres, cw_prox_thres=self.cw_prox_thres, stdmodel=self.standardmodel)
 
         if self.verbose: 
@@ -280,7 +285,7 @@ def update(self):
         self.uT         = uT_org.copy()
 
 
-    def simulate(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir", quandary_exec="", cygwinbash="", batchargs=[]):
+    def simulate(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir", quandary_exec="", cygwinbash="", mpi_exec="mpirun -np ", batchargs=[]):
         """ 
         Simulate the quantm dynamics using the current settings. 
 
@@ -293,6 +298,7 @@ def simulate(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir"
         datadir       : Data directory for storing output files. Default: "./run_dir"
         quandary_exec : Location of Quandary's C++ executable, if not in $PATH
         cygwinbash    : To run on Windows through Cygwin, set the path to Cygwin/bash.exe. Default: None.
+        mpi_exec      : String for MPI launcher prefix, e.g. "mpirun -np" or "srun -n". The string should include the flag for core counts, but not the number of cores itself which will be appended automatically
         batchargs     : [str(time), str(accountname), int(nodes)] If given, submits a batch job rather than local execution. Specify the max. runtime (string), the account name (string) and the number of requested nodes (int). Note, the number of executing *cores* is defined through 'maxcores'. 
 
         Returns:
@@ -308,10 +314,10 @@ def simulate(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir"
         if len(pt0) > 0 and len(qt0) > 0:
             pcof0 = self.downsample_pulses(pt0=pt0, qt0=qt0)
 
-        return self.__run(pcof0=pcof0, runtype="simulation", overwrite_popt=False, maxcores=maxcores, datadir=datadir, quandary_exec=quandary_exec, cygwinbash=cygwinbash, batchargs=batchargs)
+        return self.__run(pcof0=pcof0, runtype="simulation", overwrite_popt=False, maxcores=maxcores, datadir=datadir, quandary_exec=quandary_exec, cygwinbash=cygwinbash,mpi_exec=mpi_exec, batchargs=batchargs)
 
 
-    def optimize(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir", quandary_exec="", cygwinbash="", batchargs=[]):
+    def optimize(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir", quandary_exec="", cygwinbash="", mpi_exec="mpirun -np ", batchargs=[]):
         """ 
         Optimize the quantm dynamics using the current settings. 
 
@@ -322,6 +328,7 @@ def optimize(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir"
         pt, qt          :  p,q-control pulses [MHz] at each time point for each oscillator (List of list)
         datadir       : Data directory for storing output files. Default: "./run_dir"
         quandary_exec : Location of Quandary's C++ executable, if not in $PATH
+        mpi_exec      : String for MPI launcher prefix, e.g. "mpirun -np" or "srun -n". The string should include the flag for core counts, but not the number of cores itself which will be appended automatically
         cygwinbash    : To run on Windows through Cygwin, set the path to Cygwin/bash.exe. Default: None.
         batchargs     : [str(time), str(accountname), int(nodes)] If given, submits a batch job rather than local execution. Specify the max. runtime (string), the account name (string) and the number of requested nodes (int). Note, the number of executing *cores* is defined through 'maxcores'. 
 
@@ -337,10 +344,10 @@ def optimize(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir"
         if len(pt0) > 0 and len(qt0) > 0:
             pcof0 = self.downsample_pulses(pt0=pt0, qt0=qt0)
 
-        return self.__run(pcof0=pcof0, runtype="optimization", overwrite_popt=True, maxcores=maxcores, datadir=datadir, quandary_exec=quandary_exec, cygwinbash=cygwinbash, batchargs=batchargs)
+        return self.__run(pcof0=pcof0, runtype="optimization", overwrite_popt=True, maxcores=maxcores, datadir=datadir, quandary_exec=quandary_exec, cygwinbash=cygwinbash, mpi_exec=mpi_exec, batchargs=batchargs)
 
 
-    def evalControls(self, *, pcof0=[], points_per_ns=1,datadir="./run_dir", quandary_exec="", cygwinbash=""):
+    def evalControls(self, *, pcof0=[], points_per_ns=1,datadir="./run_dir", quandary_exec="", mpi_exec="mpirun -np ", cygwinbash=""):
         """
         Evaluate control pulses on a specific sample rate.       
 
@@ -350,6 +357,7 @@ def evalControls(self, *, pcof0=[], points_per_ns=1,datadir="./run_dir", quandar
         points_per_ns :  sample rate of the resulting controls. Default: 1ns 
         datadir       :  Directory for output files. Default: "./run_dir"
         quandary_exec :  Path to Quandary's C++ executable if not in $PATH
+        mpi_exec      : String for MPI launcher prefix, e.g. "mpirun -np" or "srun -n". The string should include the flag for core counts, but not the number of cores itself which will be appended automatically
         cygwinbash    : To run on Windows through Cygwin, set the path to Cygwin/bash.exe. Default: None.
 
         Returns:
@@ -367,7 +375,7 @@ def evalControls(self, *, pcof0=[], points_per_ns=1,datadir="./run_dir", quandar
         os.makedirs(datadir_controls, exist_ok=True)
         runtype = 'evalcontrols'
         configfile_eval= self.__dump(pcof0=pcof0, runtype=runtype, datadir=datadir_controls)
-        err = execute(runtype=runtype, ncores=1, config_filename=configfile_eval, datadir=datadir_controls, quandary_exec=quandary_exec, verbose=False, cygwinbash=cygwinbash)
+        err = execute(runtype=runtype, ncores=1, config_filename=configfile_eval, datadir=datadir_controls, quandary_exec=quandary_exec, verbose=False, mpi_exec=mpi_exec, cygwinbash=cygwinbash)
         time, pt, qt, _, _, _, pcof, _, _ = self.get_results(datadir=datadir_controls, ignore_failure=True)
 
         # Save pcof to config.popt
@@ -428,7 +436,7 @@ def downsample_pulses(self, *, pt0=[], qt0=[]):
         return pcof0
 
 
-    def __run(self, *, pcof0=[], runtype="optimization", overwrite_popt=False, maxcores=-1, datadir="./run_dir", quandary_exec="", cygwinbash="", batchargs=[]):
+    def __run(self, *, pcof0=[], runtype="optimization", overwrite_popt=False, maxcores=-1, datadir="./run_dir", quandary_exec="", cygwinbash="", mpi_exec="mpirun -np ", batchargs=[]):
         """
         Internal helper function to launch processes to execute the C++ Quandary code:
           1. Writes quandary config files to file system
@@ -442,17 +450,22 @@ def __run(self, *, pcof0=[], runtype="optimization", overwrite_popt=False, maxco
         config_filename = self.__dump(pcof0=pcof0, runtype=runtype, datadir=datadir)
 
         # Set default number of cores to the number of initial conditions, unless otherwise specified. Make sure ncores is an integer divisible of ninit.
-        ncores = self._ninit
+        ncores_init = self._ninit
         if maxcores > -1:
-            ncores = min(self._ninit, maxcores)
+            ncores_init = min(self._ninit, maxcores)
         for i in range(self._ninit, 0, -1):
             if self._ninit % i == 0:  # i is a factor of ninit
-                if i <= ncores:
-                    ncores = i
+                if i <= ncores_init:
+                    ncores_init = i
                     break
+        # Set remaining number of cores for petsc
+        ncores_petsc = 1
+        if maxcores > ncores_init and maxcores % ncores_init == 0:
+            ncores_petsc = int(maxcores / ncores_init)
+        ncores = ncores_init * ncores_petsc
 
         # Execute subprocess to run Quandary
-        err = execute(runtype=runtype, ncores=ncores, config_filename=config_filename, datadir=datadir, quandary_exec=quandary_exec, verbose=self.verbose, cygwinbash=cygwinbash, batchargs=batchargs)
+        err = execute(runtype=runtype, ncores=ncores, config_filename=config_filename, datadir=datadir, quandary_exec=quandary_exec, verbose=self.verbose, cygwinbash=cygwinbash, mpi_exec=mpi_exec, batchargs=batchargs)
         if self.verbose:
             print("Quandary data dir: ", datadir, "\n")
 
@@ -1332,7 +1345,7 @@ def timestep_richardson_est(quandary, tol=1e-8, order=2, quandary_exec=""):
     return errs_J, errs_u, dts
 
 
-def execute(*, runtype="simulation", ncores=1, config_filename="config.cfg", datadir=".", quandary_exec="", verbose=False, cygwinbash="", batchargs=[]):
+def execute(*, runtype="simulation", ncores=1, config_filename="config.cfg", datadir=".", quandary_exec="", verbose=False, cygwinbash="", mpi_exec="mpirun -np ", batchargs=[]):
     """ 
     Helper function to evoke a subprocess that executes Quandary.
 
@@ -1344,6 +1357,7 @@ def execute(*, runtype="simulation", ncores=1, config_filename="config.cfg", dat
     quandary_exec       (string)    : Absolute path to quandary's executable. Default: "" (expecting quandary to be in the $PATH)
     verbose             (Bool)      : Flag to print more output. Default: False
     cygwinbash          (string)    : Path to Cygwin bash.exe, if running on Windows machine. Default: None
+    mpi_exec            (string)    : MPI launcher prefix, e.g. "mpirun -np" or "srun -n". The string should include the flag for core counts, but not the number of cores itself which will be appended automatically
     batchargs           (List)      : Submit to batch system by setting batchargs= [maxime, accountname, nodes]. Default: []. Compare end of this file. Specify the max. runtime (string), the account name (string) and the number of requested nodes (int). Note, the number of executing *cores* is defined through 'ncores'. 
 
     Returns:
@@ -1367,7 +1381,7 @@ def execute(*, runtype="simulation", ncores=1, config_filename="config.cfg", dat
         if len(batchargs)>0:
             myrun = batch_run  # currently set to "srun -n"
         else:
-            myrun = "mpirun -np "
+            myrun = mpi_exec
         runcommand = f"{myrun} {ncores} " + runcommand
     if verbose:
         print("Running Quandary ... ")