pyGSTi 0.9.12__cp39-cp39-win_amd64.whl → 0.9.13__cp39-cp39-win_amd64.whl

pygsti/forwardsims/mapforwardsim.py

@@ -26,6 +26,8 @@ from pygsti.baseobjs.verbosityprinter import VerbosityPrinter as _VerbosityPrint
 from pygsti.tools import sharedmemtools as _smt
 from pygsti.tools import slicetools as _slct
 from pygsti.tools.matrixtools import _fas
+from pygsti.tools import listtools as _lt
+from pygsti.circuits import CircuitList as _CircuitList
 
 _dummy_profiler = _DummyProfiler()
 
@@ -49,7 +51,7 @@ class SimpleMapForwardSimulator(_ForwardSimulator):
     # If this is done, then MapForwardSimulator wouldn't need to separately subclass DistributableForwardSimulator
 
     def _compute_circuit_outcome_probabilities(self, array_to_fill, circuit, outcomes, resource_alloc, time=None):
-        expanded_circuit_outcomes = circuit.expand_instruments_and_separate_povm(self.model, outcomes)
+        expanded_circuit_outcomes = self.model.expand_instruments_and_separate_povm(circuit, outcomes)
         outcome_to_index = {outc: i for i, outc in enumerate(outcomes)}
         for spc, spc_outcomes in expanded_circuit_outcomes.items():  # spc is a SeparatePOVMCircuit
             # Note: `spc.circuit_without_povm` *always* begins with a prep label.
@@ -157,7 +159,7 @@ class MapForwardSimulator(_DistributableForwardSimulator, SimpleMapForwardSimula
         if method_name == 'bulk_fill_timedep_dchi2': return ('p',)  # just an additional parameter vector
         return super()._array_types_for_method(method_name)
 
-    def __init__(self, model=None, max_cache_size=0, num_atoms=None, processor_grid=None, param_blk_sizes=None,
+    def __init__(self, model=None, max_cache_size=None, num_atoms=None, processor_grid=None, param_blk_sizes=None,
                  derivative_eps=1e-7, hessian_eps=1e-5):
         #super().__init__(model, num_atoms, processor_grid, param_blk_sizes)
         _DistributableForwardSimulator.__init__(self, model, num_atoms, processor_grid, param_blk_sizes)
@@ -193,7 +195,9 @@ class MapForwardSimulator(_DistributableForwardSimulator, SimpleMapForwardSimula
                                    self._processor_grid, self._pblk_sizes)
 
     def create_layout(self, circuits, dataset=None, resource_alloc=None, array_types=('E',),
-                      derivative_dimension=None, verbosity=0):
+                      derivative_dimensions=None, verbosity=0, layout_creation_circuit_cache=None,
+                      circuit_partition_cost_functions=('size', 'propagations'),
+                      load_balancing_parameters=(1.15,.1)):
         """
         Constructs an circuit-outcome-probability-array (COPA) layout for a list of circuits.
 
@@ -214,14 +218,32 @@ class MapForwardSimulator(_DistributableForwardSimulator, SimpleMapForwardSimula
         array_types : tuple, optional
             A tuple of string-valued array types.  See :meth:`ForwardSimulator.create_layout`.
 
-        derivative_dimension : int, optional
+        derivative_dimensions : int or tuple[int], optional
             Optionally, the parameter-space dimension used when taking first
             and second derivatives with respect to the cirucit outcome probabilities.  This must be
             non-None when `array_types` contains `'ep'` or `'epp'` types.
+            If a tuple, then must be length 1.
 
         verbosity : int or VerbosityPrinter
             Determines how much output to send to stdout.  0 means no output, higher
             integers mean more output.
+        
+        layout_creation_circuit_cache : dict, optional (default None)
+            A precomputed dictionary serving as a cache for completed circuits. I.e. circuits 
+            with prep labels and POVM labels appended. Along with other useful pre-computed 
+            circuit structures used in layout creation.
+
+        circuit_partition_cost_functions : tuple of str, optional (default ('size', 'propagations'))
+            A tuple of strings denoting cost function to use in each of the two stages of the algorithm
+            for determining the partitions of the complete circuit set amongst atoms.
+            Allowed options are 'size', which corresponds to balancing the number of circuits, 
+            and 'propagations', which corresponds to balancing the number of state propagations.
+
+        load_balancing_parameters : tuple of floats, optional (default (1.15, .1))
+            A tuple of floats used as load balancing parameters when splitting a layout across atoms,
+            as in the multi-processor setting when using MPI. These parameters correspond to the `imbalance_threshold`
+            and `minimum_improvement_threshold` parameters described in the method `find_splitting_new`
+            of the `PrefixTable` class.
 
         Returns
         -------
@@ -233,7 +255,13 @@ class MapForwardSimulator(_DistributableForwardSimulator, SimpleMapForwardSimula
             if (resource_alloc.mem_limit is not None) else None  # *per-processor* memory limit
         nprocs = resource_alloc.comm_size
         comm = resource_alloc.comm
-        num_params = derivative_dimension if (derivative_dimension is not None) else self.model.num_params
+        if isinstance(derivative_dimensions, int):
+            num_params = derivative_dimensions
+        elif isinstance(derivative_dimensions, tuple):
+            assert len(derivative_dimensions) == 1
+            num_params = derivative_dimensions[0]
+        else:
+            num_params = self.model.num_params
         C = 1.0 / (1024.0**3)
 
         if mem_limit is not None:
@@ -241,15 +269,15 @@ class MapForwardSimulator(_DistributableForwardSimulator, SimpleMapForwardSimula
                 raise MemoryError("Attempted layout creation w/memory limit = %g <= 0!" % mem_limit)
             printer.log("Layout creation w/mem limit = %.2fGB" % (mem_limit * C))
 
-        #Start with how we'd like to split processors up (without regard to memory limit):
-
-        # when there are lots of processors, the from_vector calls dominante over the actual fwdsim,
-        # but we can reduce from_vector calls by having np1, np2 > 0 (each param requires a from_vector
-        # call when using finite diffs) - so we want to choose nc = Ng < nprocs and np1 > 1 (so nc * np1 = nprocs).
-        #work_per_proc = self.model.dim**2
+        #Start with how we'd like to split processors up (without regard to memory limit):        
+        #The current implementation of map (should) benefit more from having a matching between the number of atoms
+        #and the number of processors, at least for up to around two-qubits.
+        default_natoms = nprocs # heuristic
+        #TODO: factor in the mem_limit value to more intelligently set the default number of atoms.
 
         natoms, na, npp, param_dimensions, param_blk_sizes = self._compute_processor_distribution(
-            array_types, nprocs, num_params, len(circuits), default_natoms=2 * self.model.dim)  # heuristic?
+            array_types, nprocs, num_params, len(circuits), default_natoms=default_natoms)  
+        
         printer.log(f'Num Param Processors {npp}')
         
         printer.log("MapLayout: %d processors divided into %s (= %d) grid along circuit and parameter directions." %
@@ -258,7 +286,9 @@ class MapForwardSimulator(_DistributableForwardSimulator, SimpleMapForwardSimula
         assert(_np.prod((na,) + npp) <= nprocs), "Processor grid size exceeds available processors!"
 
         layout = _MapCOPALayout(circuits, self.model, dataset, self._max_cache_size, natoms, na, npp,
-                                param_dimensions, param_blk_sizes, resource_alloc, verbosity)
+                                param_dimensions, param_blk_sizes, resource_alloc,circuit_partition_cost_functions,
+                                verbosity, layout_creation_circuit_cache= layout_creation_circuit_cache,
+                                load_balancing_parameters=load_balancing_parameters)
 
         if mem_limit is not None:
             loc_nparams1 = num_params / npp[0] if len(npp) > 0 else 0
@@ -284,13 +314,6 @@ class MapForwardSimulator(_DistributableForwardSimulator, SimpleMapForwardSimula
                                                   layout._param_dimensions, (loc_nparams1, loc_nparams2),
                                                   (blk1, blk2), max_atom_cachesize, self.model.dim)
 
-            #def approx_mem_estimate(nc, np1, np2):
-            #    approx_cachesize = (num_circuits / nc) * 1.3  # inflate expected # of circuits per atom => cache_size
-            #    return _bytes_for_array_types(array_types, num_elements, num_elements / nc,
-            #                                  num_circuits, num_circuits / nc,
-            #                                  (num_params, num_params), (num_params / np1, num_params / np2),
-            #                                  approx_cachesize, self.model.dim)
-
             GB = 1.0 / 1024.0**3
             if mem_estimate > mem_limit:
                 raise MemoryError("Not enough memory for desired layout! (limit=%.1fGB, required=%.1fGB)" % (
@@ -299,6 +322,41 @@ class MapForwardSimulator(_DistributableForwardSimulator, SimpleMapForwardSimula
                 printer.log("   Esimated memory required = %.1fGB" % (mem_estimate * GB))
 
         return layout
+    
+    @staticmethod
+    def create_copa_layout_circuit_cache(circuits, model, dataset=None):
+        """
+        Helper function for pre-computing/pre-processing circuits structures
+        used in matrix layout creation.
+        """
+        cache = dict()
+        completed_circuits = model.complete_circuits(circuits)
+
+        cache['completed_circuits'] = {ckt: comp_ckt for ckt, comp_ckt in zip(circuits, completed_circuits)}
+
+        split_circuits = model.split_circuits(completed_circuits, split_prep=False)    
+        cache['split_circuits'] = {ckt: split_ckt for ckt, split_ckt in zip(circuits, split_circuits)}
+        
+
+        if dataset is not None:
+            aliases = circuits.op_label_aliases if isinstance(circuits, _CircuitList) else None
+            ds_circuits = _lt.apply_aliases_to_circuits(circuits, aliases)
+            unique_outcomes_list = []
+            for ckt in ds_circuits:
+                ds_row = dataset[ckt]
+                unique_outcomes_list.append(ds_row.unique_outcomes if ds_row is not None else None)
+        else:
+            unique_outcomes_list = [None]*len(circuits)
+
+        expanded_circuit_outcome_list = model.bulk_expand_instruments_and_separate_povm(circuits, 
+                                                                                        observed_outcomes_list = unique_outcomes_list, 
+                                                                                        completed_circuits= completed_circuits)
+        
+        expanded_circuit_cache = {ckt: expanded_ckt for ckt,expanded_ckt in zip(completed_circuits, expanded_circuit_outcome_list)}                
+        cache['expanded_and_separated_circuits'] = expanded_circuit_cache
+
+        return cache
+
 
     def _bulk_fill_probs_atom(self, array_to_fill, layout_atom, resource_alloc):
         # Note: *don't* set dest_indices arg = layout.element_slice, as this is already done by caller

pygsti/forwardsims/mapforwardsim_calc_densitymx.pyx

@@ -1,6 +1,4 @@
 # encoding: utf-8
-# cython: profile=False
-# cython: linetrace=False
 
 #***************************************************************************************************
 # Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
@@ -28,12 +26,6 @@ from ..tools import slicetools as _slct
 #from ..tools import optools as _ot
 from ..tools.matrixtools import _fas
 
-#DEBUG REMOVE MEMORY PROFILING
-#import os, psutil
-#process = psutil.Process(os.getpid())
-#def print_mem_usage(prefix):
-#    print("%s: mem usage = %.3f GB" % (prefix, process.memory_info().rss / (1024.0**3)))
-
 #Use 64-bit integers
 ctypedef long long INT
 ctypedef unsigned long long UINT
@@ -59,15 +51,16 @@ def propagate_staterep(staterep, operationreps):
 # Python -> C Conversion functions
 # -----------------------------------------
 
-cdef vector[vector[INT]] convert_maplayout(layout_atom, operation_lookup, rho_lookup):
+@cython.wraparound(False)
+cdef vector[vector[INT]] convert_maplayout(prefix_table_contents, operation_lookup, rho_lookup):
     # c_layout :
     # an array of INT-arrays; each INT-array is [iDest,iStart,iCache,<remainder gate indices>]
     cdef vector[INT] intarray
-    cdef vector[vector[INT]] c_layout_atom = vector[vector[INT]](len(layout_atom.table))
-    for kk, (iDest, iStart, remainder, iCache) in enumerate(layout_atom.table.contents):
+    cdef vector[vector[INT]] c_layout_atom = vector[vector[INT]](len(prefix_table_contents))
+    for kk, (iDest, iStart, remainder, iCache) in enumerate(prefix_table_contents):
         if iStart is None: iStart = -1 # so always an int
         if iCache is None: iCache = -1 # so always an int
-        remainder = remainder.circuit_without_povm.layertup
+        remainder = remainder
         intarray = vector[INT](3 + len(remainder))
         intarray[0] = iDest
         intarray[1] = iStart
@@ -107,6 +100,8 @@ cdef vector[vector[INT]] convert_and_wrap_dict_of_intlists(d, wrapper):
             ret[i][j] = wrapper[intlist[j]]
     return ret
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
 cdef vector[StateCRep*] create_rhocache(INT cacheSize, INT state_dim):
     cdef INT i
     cdef vector[StateCRep*] rho_cache = vector[StateCRep_ptr](cacheSize)
@@ -114,11 +109,14 @@ cdef vector[StateCRep*] create_rhocache(INT cacheSize, INT state_dim):
         rho_cache[i] = new StateCRep(state_dim)
     return rho_cache
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
 cdef void free_rhocache(vector[StateCRep*] rho_cache):
     cdef UINT i
     for i in range(rho_cache.size()): # fill cache with empty but alloc'd states
         del rho_cache[i]
 
+@cython.wraparound(False)
 cdef vector[OpCRep*] convert_opreps(operationreps):
     # c_opreps : an array of OpCReps
     cdef vector[OpCRep*] c_opreps = vector[OpCRep_ptr](len(operationreps))
@@ -130,19 +128,20 @@ cdef StateCRep* convert_rhorep(rhorep):
     # extract c-reps from rhorep and ereps => c_rho and c_ereps
     return (<StateRep?>rhorep).c_state
 
+@cython.wraparound(False)
 cdef vector[StateCRep*] convert_rhoreps(rhoreps):
     cdef vector[StateCRep*] c_rhoreps = vector[StateCRep_ptr](len(rhoreps))
     for ii,rrep in rhoreps.items(): # (ii = python variable)
         c_rhoreps[ii] = (<StateRep?>rrep).c_state
     return c_rhoreps
 
+@cython.wraparound(False)
 cdef vector[EffectCRep*] convert_ereps(ereps):
     cdef vector[EffectCRep*] c_ereps = vector[EffectCRep_ptr](len(ereps))
     for i in range(len(ereps)):
         c_ereps[i] = (<EffectRep>ereps[i]).c_effect
     return c_ereps
 
-
 # -----------------------------------------
 # Mapfill functions
 # -----------------------------------------
@@ -159,7 +158,6 @@ def mapfill_probs_atom(fwdsim, np.ndarray[double, mode="c", ndim=1] array_to_fil
     shared_mem_leader = resource_alloc.is_host_leader if (resource_alloc is not None) else True
 
     dest_indices = _slct.to_array(dest_indices)  # make sure this is an array and not a slice
-    #dest_indices = np.ascontiguousarray(dest_indices) #unneeded
 
     #Get (extension-type) representation objects
     rho_lookup = { lbl:i for i,lbl in enumerate(layout_atom.rho_labels) } # rho labels -> ints for faster lookup
@@ -169,7 +167,7 @@ def mapfill_probs_atom(fwdsim, np.ndarray[double, mode="c", ndim=1] array_to_fil
     ereps = [fwdsim.model._circuit_layer_operator(elbl, 'povm')._rep for elbl in layout_atom.full_effect_labels]  # cache these in future
 
     # convert to C-mode:  evaltree, operation_lookup, operationreps
-    cdef vector[vector[INT]] c_layout_atom = convert_maplayout(layout_atom, operation_lookup, rho_lookup)
+    cdef vector[vector[INT]] c_layout_atom = convert_maplayout(layout_atom.table.contents, operation_lookup, rho_lookup)
     cdef vector[StateCRep*] c_rhos = convert_rhoreps(rhoreps)
     cdef vector[EffectCRep*] c_ereps = convert_ereps(ereps)
     cdef vector[OpCRep*] c_opreps = convert_opreps(operationreps)
@@ -182,17 +180,6 @@ def mapfill_probs_atom(fwdsim, np.ndarray[double, mode="c", ndim=1] array_to_fil
     cdef vector[vector[INT]] final_indices_per_circuit = convert_and_wrap_dict_of_intlists(
         layout_atom.elindices_by_expcircuit, dest_indices)
 
-    #DEBUG REMOVE
-    #print_mem_usage("MAPFILL PROBS begin")
-    #for i in [1808, 419509, 691738, 497424]:
-    #    from ..evotypes.densitymx.opreps import OpRepComposed
-    #    op = operationreps[i]
-    #    if isinstance(op.embedded_rep, OpRepComposed):
-    #        extra = " factors = " + ', '.join([str(type(opp)) for opp in op.embedded_rep.factor_reps])
-    #    else:
-    #        extra = ""
-    #    print("ID ",i,str(type(op)),str(type(op.embedded_rep)), extra)
-
     if shared_mem_leader:
         #Note: dm_mapfill_probs could have taken a resource_alloc to employ multiple cpus to do computation.
         # Since array_fo_fill is assumed to be shared mem it would need to only update `array_to_fill` *if*
@@ -203,6 +190,7 @@ def mapfill_probs_atom(fwdsim, np.ndarray[double, mode="c", ndim=1] array_to_fil
     free_rhocache(rho_cache)  #delete cache entries
 
 
+@cython.wraparound(False)
 cdef dm_mapfill_probs(double[:] array_to_fill,
                       vector[vector[INT]] c_layout_atom,
                       vector[OpCRep*] c_opreps,
@@ -234,7 +222,7 @@ cdef dm_mapfill_probs(double[:] array_to_fill,
     # - all rho_cache entries have been allocated via "new"
     #REMOVE print("MAPFILL PROBS begin cfn")
     for k in range(<INT>c_layout_atom.size()):
-        t0 = pytime.time() # DEBUG
+        #t0 = pytime.time() # DEBUG
         intarray = c_layout_atom[k]
         i = intarray[0]
         istart = intarray[1]
@@ -298,22 +286,27 @@ cdef dm_mapfill_probs(double[:] array_to_fill,
     del prop2
     del shelved
 
-
+@cython.wraparound(False)
 def mapfill_dprobs_atom(fwdsim,
                         np.ndarray[double, ndim=2] array_to_fill,
                         dest_indices,
                         dest_param_indices,
                         layout_atom, param_indices, resource_alloc, double eps):
 
-    #cdef double eps = 1e-7
+    cdef int num_params = fwdsim.model.num_params
+    cdef int model_dim = fwdsim.model.dim
 
     if param_indices is None:
-        param_indices = list(range(fwdsim.model.num_params))
+        param_indices = list(range(num_params))
     if dest_param_indices is None:
         dest_param_indices = list(range(_slct.length(param_indices)))
 
-    param_indices = _slct.to_array(param_indices)
-    dest_param_indices = _slct.to_array(dest_param_indices)
+    cdef np.ndarray[np.int64_t, ndim=1, mode='c'] param_indices_array = _slct.to_array(param_indices)
+    cdef np.ndarray[np.int64_t, ndim=1, mode='c'] dest_param_indices_array = _slct.to_array(dest_param_indices)
+
+    cdef np.int64_t[::1] param_indices_view = param_indices_array
+    cdef np.int64_t[::1] dest_param_indices_view = dest_param_indices_array
+    
 
     #Get (extension-type) representation objects
     # NOTE: the circuit_layer_operator(lbl) functions cache the returned operation
@@ -327,7 +320,7 @@ def mapfill_dprobs_atom(fwdsim,
     ereps = [fwdsim.model._circuit_layer_operator(elbl, 'povm')._rep for elbl in layout_atom.full_effect_labels]  # cache these in future
 
     # convert to C-mode:  evaltree, operation_lookup, operationreps
-    cdef vector[vector[INT]] c_layout_atom = convert_maplayout(layout_atom, operation_lookup, rho_lookup)
+    cdef vector[vector[INT]] c_layout_atom = convert_maplayout(layout_atom.table.contents, operation_lookup, rho_lookup)
     cdef vector[StateCRep*] c_rhos = convert_rhoreps(rhoreps)
     cdef vector[EffectCRep*] c_ereps = convert_ereps(ereps)
     cdef vector[OpCRep*] c_opreps = convert_opreps(operationreps)
@@ -335,53 +328,60 @@ def mapfill_dprobs_atom(fwdsim,
     # create rho_cache = vector of StateCReps
     #print "DB: creating rho_cache of size %d * %g GB => %g GB" % \
     #   (layout_atom.cache_size, 8.0 * fwdsim.model.dim / 1024.0**3, layout_atom.cache_size * 8.0 * fwdsim.model.dim / 1024.0**3)
-    cdef vector[StateCRep*] rho_cache = create_rhocache(layout_atom.cache_size, fwdsim.model.dim)
+    cdef vector[StateCRep*] rho_cache = create_rhocache(layout_atom.cache_size, model_dim)
 
     cdef vector[vector[INT]] elabel_indices_per_circuit = convert_dict_of_intlists(layout_atom.elbl_indices_by_expcircuit)
     cdef vector[vector[INT]] final_indices_per_circuit = convert_dict_of_intlists(layout_atom.elindices_by_expcircuit)
 
-    orig_vec = fwdsim.model.to_vector().copy()
+    cdef np.ndarray[np.float64_t, ndim=1, mode='c'] orig_vec = fwdsim.model.to_vector().copy()
+    cdef double[::1] orig_vec_view = orig_vec
     fwdsim.model.from_vector(orig_vec, close=False)  # ensure we call with close=False first
 
     nEls = layout_atom.num_elements
-    probs = np.empty(nEls, 'd') #must be contiguous!
-    probs2 = np.empty(nEls, 'd') #must be contiguous!
+    cdef np.ndarray[np.float64_t, ndim=1, mode='c'] probs = np.empty(nEls, dtype=np.float64) #must be contiguous!
+    cdef np.ndarray[np.float64_t, ndim=1, mode='c'] probs2 = np.empty(nEls, dtype=np.float64) #must be contiguous!
+
+    cdef double[::1] probs_view = probs
+    cdef double[::1] probs2_view = probs2
 
     #if resource_alloc.comm_rank == 0:
     #    print("MAPFILL DPROBS ATOM 1"); t=pytime.time(); t0=pytime.time()
-    dm_mapfill_probs(probs, c_layout_atom, c_opreps, c_rhos, c_ereps, &rho_cache,
-                     elabel_indices_per_circuit, final_indices_per_circuit, fwdsim.model.dim)
+    dm_mapfill_probs(probs_view, c_layout_atom, c_opreps, c_rhos, c_ereps, &rho_cache, elabel_indices_per_circuit, final_indices_per_circuit, model_dim)
     #if resource_alloc.comm_rank == 0:
     #    print("MAPFILL DPROBS ATOM 2 %.3fs" % (pytime.time() - t)); t=pytime.time()
 
-    shared_mem_leader = resource_alloc.is_host_leader
-
-    #Get a map from global parameter indices to the desired
-    # final index within array_to_fill
-    iParamToFinal = {i: dest_index for i, dest_index in zip(param_indices, dest_param_indices)}
+    cdef bint shared_mem_leader = resource_alloc.is_host_leader
+
+    #add typing to indices we'll be using below:
+    cdef int i
+    cdef int first_param_idx
+    cdef int iFinal
+
+    #Split off the first finite difference step, as the pattern I want in the loop with each step
+    #is to simultaneously undo the previous update and apply the new one.
+    if len(param_indices_view)>0:
+        #probs2_view[:] = probs_view[:]
+        first_param_idx = param_indices_view[0]
+        iFinal = dest_param_indices_view[0]
+        fwdsim.model.set_parameter_value(first_param_idx, orig_vec_view[first_param_idx]+eps)
+        if shared_mem_leader:  # don't fill assumed-shared array-to_fill on non-mem-leaders
+            dm_mapfill_probs(probs2_view, c_layout_atom, c_opreps, c_rhos, c_ereps, &rho_cache, elabel_indices_per_circuit, final_indices_per_circuit, model_dim)
+            array_to_fill[dest_indices, iFinal] = (probs2 - probs) / eps
+
+    for i in range(1, len(param_indices_view)):
+        #probs2_view[:] = probs_view[:]
+        iFinal = dest_param_indices_view[i]
+        fwdsim.model.set_parameter_values([param_indices_view[i-1], param_indices_view[i]], [orig_vec_view[param_indices_view[i-1]], orig_vec_view[param_indices_view[i]]+eps])
+
+        if shared_mem_leader:  # don't fill assumed-shared array-to_fill on non-mem-leaders
+            dm_mapfill_probs(probs2_view, c_layout_atom, c_opreps, c_rhos, c_ereps,  &rho_cache, elabel_indices_per_circuit, final_indices_per_circuit, model_dim)
+            array_to_fill[dest_indices, iFinal] = (probs2 - probs) / eps
+        
+    #reset the final model parameter we changed to it's original value.
+    fwdsim.model.set_parameter_value(param_indices_view[len(param_indices_view)-1], orig_vec_view[param_indices_view[len(param_indices_view)-1]])
 
-    for i in range(fwdsim.model.num_params):
-        #print("dprobs cache %d of %d" % (i,self.Np))
-        if i in iParamToFinal:
-            #if resource_alloc.comm_rank == 0:
-            #    print("MAPFILL DPROBS ATOM 3 (i=%d) %.3fs elapssed=%.1fs" % (i, pytime.time() - t, pytime.time() - t0)); t=pytime.time()
-            iFinal = iParamToFinal[i]
-            vec = orig_vec.copy(); vec[i] += eps
-            fwdsim.model.from_vector(vec, close=True)
-            #Note: dm_mapfill_probs could have taken a resource_alloc to employ multiple cpus to do computation.
-            # If probs2 were shared mem (seems not benefit to this?) it would need to only update `probs2` *if*
-            # it were the host leader.
-            if shared_mem_leader:  # don't fill assumed-shared array-to_fill on non-mem-leaders
-                dm_mapfill_probs(probs2, c_layout_atom, c_opreps, c_rhos, c_ereps, &rho_cache,
-                                 elabel_indices_per_circuit, final_indices_per_circuit, fwdsim.model.dim)
-                #_fas(array_to_fill, [dest_indices, iFinal], (probs2 - probs) / eps)  # I don't think this is needed
-                array_to_fill[dest_indices, iFinal] = (probs2 - probs) / eps
-
-    #if resource_alloc.comm_rank == 0:
-    #    print("MAPFILL DPROBS ATOM 4 elapsed=%.1fs" % (pytime.time() - t0))
-    fwdsim.model.from_vector(orig_vec, close=True)
     free_rhocache(rho_cache)  #delete cache entries
-
+    
 
 cdef double TDchi2_obj_fn(double p, double f, double n_i, double n, double omitted_p, double min_prob_clip_for_weighting, double extra):
     cdef double cp, v, omitted_cp
@@ -475,7 +475,6 @@ def mapfill_TDterms(fwdsim, objective, array_to_fill, dest_indices, num_outcomes
     #comm is currently ignored
     #TODO: if layout_atom is split, distribute among processors
     for iDest, iStart, remainder, iCache in layout_atom.table.contents:
-        remainder = remainder.circuit_without_povm.layertup
         rholabel = remainder[0]; remainder = remainder[1:]
         rhoVec = fwdsim.model._circuit_layer_operator(rholabel, 'prep')
 

pygsti/forwardsims/mapforwardsim_calc_generic.py

@@ -51,7 +51,6 @@ def mapfill_probs_atom(fwdsim, mx_to_fill, dest_indices, layout_atom, resource_a
 
     #TODO: if layout_atom is split, distribute somehow among processors(?) instead of punting for all but rank-0 above
     for iDest, iStart, remainder, iCache in layout_atom.table.contents:
-        remainder = remainder.circuit_without_povm.layertup
 
         if iStart is None:  # then first element of remainder is a state prep label
             rholabel = remainder[0]
@@ -73,7 +72,63 @@ def mapfill_probs_atom(fwdsim, mx_to_fill, dest_indices, layout_atom, resource_a
                 mx_to_fill[final_indices] = povmreps[povm_lbl].probabilities(final_state, None, effect_labels)
         else:
             ereps = [effectreps[j] for j in layout_atom.elbl_indices_by_expcircuit[iDest]]
+            #print(ereps)
+            if shared_mem_leader:
+                for j, erep in zip(final_indices, ereps):
+                    mx_to_fill[j] = erep.probability(final_state)  # outcome probability
+    #raise Exception
+#Version of the probability calculation that updates circuit probabilities conditionally based on
+#Whether the circuit is sensitive to the parameter. If not we leave that circuit alone.
+def cond_update_probs_atom(fwdsim, mx_to_fill, dest_indices, layout_atom, param_index, resource_alloc):
+
+    # The required ending condition is that array_to_fill on each processor has been filled.  But if
+    # memory is being shared and resource_alloc contains multiple processors on a single host, we only
+    # want *one* (the rank=0) processor to perform the computation, since array_to_fill will be
+    # shared memory that we don't want to have muliple procs using simultaneously to compute the
+    # same thing.  Thus, we carefully guard any shared mem updates/usage
+    # using "if shared_mem_leader" (and barriers, if needed) below.
+    shared_mem_leader = resource_alloc.is_host_leader if (resource_alloc is not None) else True
+
+    dest_indices = _slct.to_array(dest_indices)  # make sure this is an array and not a slice
+    cacheSize = layout_atom.jac_table.cache_size_by_parameter[param_index]
+
+    #Create rhoCache
+    rho_cache = [None] * cacheSize  # so we can store (s,p) tuples in cache
+
+    #Get operationreps and ereps now so we don't make unnecessary ._rep references
+    rhoreps = {rholbl: fwdsim.model._circuit_layer_operator(rholbl, 'prep')._rep for rholbl in layout_atom.rho_labels}
+    operationreps = {gl: fwdsim.model._circuit_layer_operator(gl, 'op')._rep for gl in layout_atom.op_labels}
+    povmreps = {plbl: fwdsim.model._circuit_layer_operator(plbl, 'povm')._rep for plbl in layout_atom.povm_labels}
+    if any([(povmrep is None) for povmrep in povmreps.values()]):
+        effectreps = {i: fwdsim.model._circuit_layer_operator(Elbl, 'povm')._rep
+                      for i, Elbl in enumerate(layout_atom.full_effect_labels)}  # cache these in future
+    else:
+        effectreps = None  # not needed, as we use povm reps directly
+
+
+    #TODO: if layout_atom is split, distribute somehow among processors(?) instead of punting for all but rank-0 above
+
+    for iDest, iStart, remainder, iCache in layout_atom.jac_table.contents_by_parameter[param_index]:
+ 
+        if iStart is None:  # then first element of remainder is a state prep label
+            rholabel = remainder[0]
+            init_state = rhoreps[rholabel]
+            remainder = remainder[1:]
+        else:
+            init_state = rho_cache[iStart]  # [:,None]
+
+        final_state = propagate_staterep(init_state, [operationreps[gl] for gl in remainder])
+        if iCache is not None: rho_cache[iCache] = final_state  # [:,0] #store this state in the cache
+
+        final_indices = [dest_indices[j] for j in layout_atom.elindices_by_expcircuit[iDest]]
+
+        if effectreps is None:
+            povm_lbl, *effect_labels = layout_atom.povm_and_elbls_by_expcircuit[iDest]
 
+            if shared_mem_leader:
+                mx_to_fill[final_indices] = povmreps[povm_lbl].probabilities(final_state, None, effect_labels)
+        else:
+            ereps = [effectreps[j] for j in layout_atom.elbl_indices_by_expcircuit[iDest]]
             if shared_mem_leader:
                 for j, erep in zip(final_indices, ereps):
                     mx_to_fill[j] = erep.probability(final_state)  # outcome probability
@@ -82,11 +137,10 @@ def mapfill_probs_atom(fwdsim, mx_to_fill, dest_indices, layout_atom, resource_a
 def mapfill_dprobs_atom(fwdsim, mx_to_fill, dest_indices, dest_param_indices, layout_atom, param_indices,
                         resource_alloc, eps):
 
-    #eps = 1e-7
-    #shared_mem_leader = resource_alloc.is_host_leader if (resource_alloc is not None) else True
+    num_params = fwdsim.model.num_params
 
     if param_indices is None:
-        param_indices = list(range(fwdsim.model.num_params))
+        param_indices = list(range(num_params))
     if dest_param_indices is None:
         dest_param_indices = list(range(_slct.length(param_indices)))
 
@@ -105,19 +159,43 @@ def mapfill_dprobs_atom(fwdsim, mx_to_fill, dest_indices, dest_param_indices, la
     nEls = layout_atom.num_elements
     probs, shm = _smt.create_shared_ndarray(resource_alloc, (nEls,), 'd', memory_tracker=None)
     probs2, shm2 = _smt.create_shared_ndarray(resource_alloc, (nEls,), 'd', memory_tracker=None)
+    #probs2_test, shm2_test = _smt.create_shared_ndarray(resource_alloc, (nEls,), 'd', memory_tracker=None)
+    
+    #mx_to_fill_test = mx_to_fill.copy()
+
     mapfill_probs_atom(fwdsim, probs, slice(0, nEls), layout_atom, resource_alloc)  # probs != shared
 
-    for i in range(fwdsim.model.num_params):
-        #print("dprobs cache %d of %d" % (i,self.Np))
-        if i in iParamToFinal:
-            iFinal = iParamToFinal[i]
-            vec = orig_vec.copy(); vec[i] += eps
-            fwdsim.model.from_vector(vec, close=True)
-            mapfill_probs_atom(fwdsim, probs2, slice(0, nEls), layout_atom, resource_alloc)
-            _fas(mx_to_fill, [dest_indices, iFinal], (probs2 - probs) / eps)
-    fwdsim.model.from_vector(orig_vec, close=True)
+    #Split off the first finite difference step, as the pattern I want in the loop with each step
+    #is to simultaneously undo the previous update and apply the new one.
+    if len(param_indices)>0:
+        probs2[:] = probs[:]
+        first_param_idx = param_indices[0]
+        iFinal = iParamToFinal[first_param_idx]
+        fwdsim.model.set_parameter_value(first_param_idx, orig_vec[first_param_idx]+eps)
+        #mapfill_probs_atom(fwdsim, probs2, slice(0, nEls), layout_atom, resource_alloc)
+        cond_update_probs_atom(fwdsim, probs2, slice(0, nEls), layout_atom, first_param_idx, resource_alloc)
+        #assert _np.linalg.norm(probs2_test-probs2) < 1e-10
+        #print(f'{_np.linalg.norm(probs2_test-probs2)=}')
+        _fas(mx_to_fill, [dest_indices, iFinal], (probs2 - probs) / eps)
+
+
+    for i in range(1, len(param_indices)):
+        probs2[:] = probs[:]
+        iFinal = iParamToFinal[param_indices[i]]
+        fwdsim.model.set_parameter_values([param_indices[i-1], param_indices[i]], 
+                                          [orig_vec[param_indices[i-1]], orig_vec[param_indices[i]]+eps])
+        #mapfill_probs_atom(fwdsim, probs2, slice(0, nEls), layout_atom, resource_alloc)
+        cond_update_probs_atom(fwdsim, probs2, slice(0, nEls), layout_atom, param_indices[i], resource_alloc)
+        #assert _np.linalg.norm(probs2_test-probs2) < 1e-10
+        #print(f'{_np.linalg.norm(probs2_test-probs2)=}')
+        _fas(mx_to_fill, [dest_indices, iFinal], (probs2 - probs) / eps)
+
+    #reset the final model parameter we changed to it's original value.
+    fwdsim.model.set_parameter_value(param_indices[-1], orig_vec[param_indices[-1]])
+
     _smt.cleanup_shared_ndarray(shm)
     _smt.cleanup_shared_ndarray(shm2)
+    #_smt.cleanup_shared_ndarray(shm2_test)
 
 
 def mapfill_TDchi2_terms(fwdsim, array_to_fill, dest_indices, num_outcomes, layout_atom, dataset_rows,