pygsti.protocols.stability.StabilityAnalysis

class StabilityAnalysis(significance=0.05, transform='auto', marginalize='auto', mergeoutcomes=None, constnumtimes='auto', ids=False, frequencies='auto', freqpointers=None, freqstest=None, tests='auto', inclass_correction=None, betweenclass_weighting='auto', estimator='auto', modelselector=None, verbosity=1, name=None)

Bases: Protocol

Stability Analysis protocol

Parameters:

• ds (DataSet or MultiDataSet) – A DataSet containing time-series data to be analyzed for signs of instability.

• significance (float, optional) – The global significance level. With defaults for all other inputs (a wide range of non-default options), the family-wise error rate of the set of all hypothesis tests performed is controlled to this value.

• transform (str, optional) –

  The type of transform to use in the spectral analysis. Options are:

  • ’auto’: An attempt is made to choose the best transform given the “meta-data” of the data,

    e.g., the variability in the time-step between data points. For beginners, ‘auto’ is the best option. If you are familiar with the underlying methods, the meta-data of the input, and the relative merits of the different transform, then it is probably better to choose this yourself – as the auto-selection is not hugely sophisticated.

  • ’dct’The Type-II Discrete Cosine Transform (with an orthogonal normalization). This is

    the only tested option, and it is our recommended option when the data is approximately equally-spaced, i.e., the time-step between each “click” for each circuit is almost a constant. (the DCT transform implicitly assumes that this time-step is exactly constant)

  • ’dft’The discrete Fourier transform (with an orthogonal normalization).

    This is an experimental feature, and the results are unreliable with this transform

  • ’lsp’The Lomb-Scargle periodogram.

    This is an experimental feature, and the code is untested with this transform

• marginalize (str or bool, optional) – True, False or ‘auto’. Whether or not to marginalize multi-qubit data, to look for instability in the marginalized probability distribution over the two outcomes for each qubit. Cannot be set to True if mergeoutcomes is not None.

• mergeoutcomes (None or Dict, optional) – If not None, a dictionary of outcome-merging dictionaries. Each dictionary contained as a value of mergeoutcomes is used to create a new DataSet, where the values have been merged according to that dictionary (see the aggregate_dataset_outcomes() function inside datasetconstructions.py). The corresponding key is used as the key for that DataSet, when it is stored in a MultiDataSet, and the instability analysis is implemented on each DataSet. This is a more general data coarse-grainin option than marginalize.

• constnumtimes (str or bool, optional) – True, False or ‘auto’. If True then data is discarded from the end of the “clickstream” for each circuit until all circuits have the same length clickstream, i.e., the same number of data acquisition times. If ‘auto’ then it is set to True or False depending on the meta-data of the data and the type of transform being used.

• ids (True or False, optional) – Whether the multiple DataSets should be treat as generated from independent random variables. If the input is a DataSet and marginalize is False and mergeoutcomes is None then this input is irrelevant: there is only ever one DataSet being analyzed. But in general multiple DataSets are concurrently analyzed. This is irrelevant for independent analyses of the DataSets, but the analysis is capable of also implementing a joint analysis of the DataSets. This joint analysis is only valid on the assumption of independent DataSets, and so this analysis will not be permitted unless ids is set to True. Note that the set of N marginalized data from N-qubit circuits are generally not independent – even if the circuits contain no 2-qubit gates then crosstalk can causes dependencies. However, as long as the dependencies are weak then settings this to True is likely ok.

• frequencies ('auto' or list, optional) –

  The frequencies that the power spectra are calculated for. If ‘auto’ these are automatically determined from the meta-data of the time-series data (e.g., using the mean time between data points) and the transform being used. If not ‘auto’, then a list of lists, where each list is a set of frequencies that are the frequencies corresponding to one or more power spectra. The frequencies that should be paired to a given power spectrum are specified by freqpointers.

  These frequencies (whether automatically calculated or explicitly input) have a fundmentally different meaning depending on whether the transform is time-stamp aware (here, the LSP) or not (here, the DCT and DFT).

  Time-stamp aware transforms take the frequencies to calculate powers at as an input, so the specified frequencies are, explicitly, the frequencies associated with the powers. The task of choosing the frequencies amounts to picking the best set of frequencies at which to interrogate the true probability trajectory for components. As there are complex factors involved in this choice that the code has no way of knowing, sometimes it is best to choose them yourself. E.g., if different frequencies are used for different circuits it isn’t possible to (meaningfully) averaging power spectra across circuits, but this might be preferable if the time-step is sufficiently different between different circuits – it depends on your aims.

  For time-stamp unaware transforms, these frequencies should be the frequencies that, given that we’re implementing the, e.g., DCT, the generated power spectrum is implicitly with respect to. In the case of data on a fixed time-grid, i.e., equally spaced data, then there is a precise set of frequencies implicit in the transform (which will be accurately extracted with frequencies set to auto). Otherwise, these frequencies are explicitly at least slightly ad hoc, and choosing these frequencies amounts to choosing those frequencies that “best” approximate the properties being interogatted with fitting each, e.g., DCT basis function to the (timestamp-free) data. The ‘auto’ option bases there frequencies solely on the mean time step and the number of times, and is a decent option when the time stamps are roughly equally spaced for each circuit.

  These frequencies should be in units of 1/t where ‘t’ is the unit of the time stamps.

• freqpointers (dict, optional) – Specifies which frequencies correspond to which power spectra. The keys are power spectra labels, and the values are integers that point to the index of frequencies (a list of lists) that the relevant frequencies are found at. Whenever a power spectra is not included in freqpointers then this defaults to 0. So if frequencies is specified and is a list containing a single list (of frequencies) then freqpointers can be left as the empty dictionary.

• freqstest (None or list, optional) – If not not None, a list of the frequency indices at which to test the powers. Leave as None to perform comprehensive testing of the power spectra.

• tests ('auto' or tuple, optional) –

  Specifies the set of hypothesis tests to perform. If ‘auto’ then an set of tests is automatically chosen. This set of tests will be suitable for most purposes, but sometimes it is useful to override this. If a tuple, the elements are “test classes”, that specifies a set of hypothesis tests to run, and each test class is itself specified by a tuple. The tests specified by each test class in this tuple are all implemented. A test class is a tuple containing some subset of ‘dataset’, ‘circuit’ and ‘outcome’, which specifies a set of power spectra. Specifically, a power spectra has been calculated for the clickstream for every combination of eachinput DataSet (e.g., there are multiple DataSets if there has been marginalization of multi-qubit data), each Circuit in the DataSet, and each possible outcome in the DataSet. For each of “dataset”, “circuit” and “outcome” not included in a tuple defining a test class, the corresponding “axis” of the 3-dimensional array of spectra is averaged over, and these spectra are then tested. So the tuple () specifies the “test class” whereby we test the power spectrum obtained by averaging all power spectra; the tuple (‘dataset’,’circuit’) specifies the “test class” whereby we average only over outcomes, obtaining a single power spectrum for each DataSet and Circuit combination, which we test.

  The default option for “tests” is appropriate for most circumstances, and it consists of (), (‘dataset’) and (‘dataset’, ‘circuit’) with duplicates removed (e.g., if there is a single DataSet then () is equivalent to (‘dataset’)).

• inclass_correction (dict, optional) – A dictionary with keys ‘dataset’, ‘circuit’, ‘outcome’ and ‘spectrum’, and values that specify the type of multi-test correction used to account for the multiple tests being implemented. This specifies how the statistically significance is maintained within the tests implemented in a single “test class”.

• betweenclass_weighting ('auto' or dict, optional) – The weighting to use to maintain statistical significance between the different classes of test being implemented. If ‘auto’ then a standard Bonferroni correction is used.

• estimator (str, optional) –

  The name of the estimator to use. This is the method used to estimate the parameters of a parameterized model for each probability trajectory, after that parameterized model has been selected with the model selection methods. Allowed values are:

  • ’auto’. The estimation method is chosen automatically, default to the fast method that is also

    reasonably reliable.

  • ’filter’. Performs a type of signal filtering: implements the transform used for generating power

    spectra (e.g., the DCT), sets the amplitudes to zero for all freuquencies that the model selection has not included in the model, inverts the transform, and then performs some minor post-processing to guarantee probabilities within [0, 1]. This method is less statically well-founded than ‘mle’, but it is faster and typically gives similar results. This method is not an option for non-invertable transforms, such as the Lomb-Scargle periodogram.

  • ’mle’. Implements maximum likelihood estimation, on the parameterized model chosen by the model

    selection. The most statistically well-founded option, but can be slower than ‘filter’ and relies on numerical optimization.

• modelselector (tuple, optional) – The model selection method. If not None, a “test class” tuple, specifying which test results to use to decide which frequencies are significant for each circuit, to then construct a parameterized model for each probability trajectory. This can be typically set to None, and it will be chosen automatically. But if you wish to use specific test results for the model selection then this should be set.

• verbosity (int, optional) – The amount of print-to-screen

• name (str, optional) – The name of this protocol, also used to (by default) name the results produced by this protocol. If None, the class name will be used.

Implements instability (“drift”) detection and characterization on timeseries data from any set of quantum circuits on any number of qubits. This uses the StabilityAnalyzer object, and directly accessing that object allows for some more complex analyzes to be performed. That object also offers a more step-by-step analysis procedure, which may be helpful for exploring the optional arguments of this analysis.

Parameters:

• ds (DataSet or MultiDataSet) – A DataSet containing time-series data to be analyzed for signs of instability.

• significance (float, optional) – The global significance level. With defaults for all other inputs (a wide range of non-default options), the family-wise error rate of the set of all hypothesis tests performed is controlled to this value.

• transform (str, optional) –

  The type of transform to use in the spectral analysis. Options are:

  • ’auto’: An attempt is made to choose the best transform given the “meta-data” of the data,

    e.g., the variability in the time-step between data points. For beginners, ‘auto’ is the best option. If you are familiar with the underlying methods, the meta-data of the input, and the relative merits of the different transform, then it is probably better to choose this yourself – as the auto-selection is not hugely sophisticated.

  • ’dct’The Type-II Discrete Cosine Transform (with an orthogonal normalization). This is

    the only tested option, and it is our recommended option when the data is approximately equally-spaced, i.e., the time-step between each “click” for each circuit is almost a constant. (the DCT transform implicitly assumes that this time-step is exactly constant)

  • ’dft’The discrete Fourier transform (with an orthogonal normalization).

    This is an experimental feature, and the results are unreliable with this transform

  • ’lsp’The Lomb-Scargle periodogram.

    This is an experimental feature, and the code is untested with this transform

• marginalize (str or bool, optional) – True, False or ‘auto’. Whether or not to marginalize multi-qubit data, to look for instability in the marginalized probability distribution over the two outcomes for each qubit. Cannot be set to True if mergeoutcomes is not None.

• mergeoutcomes (None or Dict, optional) – If not None, a dictionary of outcome-merging dictionaries. Each dictionary contained as a value of mergeoutcomes is used to create a new DataSet, where the values have been merged according to that dictionary (see the aggregate_dataset_outcomes() function inside datasetconstructions.py). The corresponding key is used as the key for that DataSet, when it is stored in a MultiDataSet, and the instability analysis is implemented on each DataSet. This is a more general data coarse-grainin option than marginalize.

• constnumtimes (str or bool, optional) – True, False or ‘auto’. If True then data is discarded from the end of the “clickstream” for each circuit until all circuits have the same length clickstream, i.e., the same number of data acquisition times. If ‘auto’ then it is set to True or False depending on the meta-data of the data and the type of transform being used.

• ids (True or False, optional) – Whether the multiple DataSets should be treat as generated from independent random variables. If the input is a DataSet and marginalize is False and mergeoutcomes is None then this input is irrelevant: there is only ever one DataSet being analyzed. But in general multiple DataSets are concurrently analyzed. This is irrelevant for independent analyses of the DataSets, but the analysis is capable of also implementing a joint analysis of the DataSets. This joint analysis is only valid on the assumption of independent DataSets, and so this analysis will not be permitted unless ids is set to True. Note that the set of N marginalized data from N-qubit circuits are generally not independent – even if the circuits contain no 2-qubit gates then crosstalk can causes dependencies. However, as long as the dependencies are weak then settings this to True is likely ok.

• frequencies ('auto' or list, optional) –

  The frequencies that the power spectra are calculated for. If ‘auto’ these are automatically determined from the meta-data of the time-series data (e.g., using the mean time between data points) and the transform being used. If not ‘auto’, then a list of lists, where each list is a set of frequencies that are the frequencies corresponding to one or more power spectra. The frequencies that should be paired to a given power spectrum are specified by freqpointers.

  These frequencies (whether automatically calculated or explicitly input) have a fundmentally different meaning depending on whether the transform is time-stamp aware (here, the LSP) or not (here, the DCT and DFT).

  Time-stamp aware transforms take the frequencies to calculate powers at as an input, so the specified frequencies are, explicitly, the frequencies associated with the powers. The task of choosing the frequencies amounts to picking the best set of frequencies at which to interrogate the true probability trajectory for components. As there are complex factors involved in this choice that the code has no way of knowing, sometimes it is best to choose them yourself. E.g., if different frequencies are used for different circuits it isn’t possible to (meaningfully) averaging power spectra across circuits, but this might be preferable if the time-step is sufficiently different between different circuits – it depends on your aims.

  For time-stamp unaware transforms, these frequencies should be the frequencies that, given that we’re implementing the, e.g., DCT, the generated power spectrum is implicitly with respect to. In the case of data on a fixed time-grid, i.e., equally spaced data, then there is a precise set of frequencies implicit in the transform (which will be accurately extracted with frequencies set to auto). Otherwise, these frequencies are explicitly at least slightly ad hoc, and choosing these frequencies amounts to choosing those frequencies that “best” approximate the properties being interogatted with fitting each, e.g., DCT basis function to the (timestamp-free) data. The ‘auto’ option bases there frequencies solely on the mean time step and the number of times, and is a decent option when the time stamps are roughly equally spaced for each circuit.

  These frequencies should be in units of 1/t where ‘t’ is the unit of the time stamps.

• freqpointers (dict, optional) – Specifies which frequencies correspond to which power spectra. The keys are power spectra labels, and the values are integers that point to the index of frequencies (a list of lists) that the relevant frequencies are found at. Whenever a power spectra is not included in freqpointers then this defaults to 0. So if frequencies is specified and is a list containing a single list (of frequencies) then freqpointers can be left as the empty dictionary.

• freqstest (None or list, optional) – If not not None, a list of the frequency indices at which to test the powers. Leave as None to perform comprehensive testing of the power spectra.

• tests ('auto' or tuple, optional) –

  Specifies the set of hypothesis tests to perform. If ‘auto’ then an set of tests is automatically chosen. This set of tests will be suitable for most purposes, but sometimes it is useful to override this. If a tuple, the elements are “test classes”, that specifies a set of hypothesis tests to run, and each test class is itself specified by a tuple. The tests specified by each test class in this tuple are all implemented. A test class is a tuple containing some subset of ‘dataset’, ‘circuit’ and ‘outcome’, which specifies a set of power spectra. Specifically, a power spectra has been calculated for the clickstream for every combination of eachinput DataSet (e.g., there are multiple DataSets if there has been marginalization of multi-qubit data), each Circuit in the DataSet, and each possible outcome in the DataSet. For each of “dataset”, “circuit” and “outcome” not included in a tuple defining a test class, the corresponding “axis” of the 3-dimensional array of spectra is averaged over, and these spectra are then tested. So the tuple () specifies the “test class” whereby we test the power spectrum obtained by averaging all power spectra; the tuple (‘dataset’,’circuit’) specifies the “test class” whereby we average only over outcomes, obtaining a single power spectrum for each DataSet and Circuit combination, which we test.

  The default option for “tests” is appropriate for most circumstances, and it consists of (), (‘dataset’) and (‘dataset’, ‘circuit’) with duplicates removed (e.g., if there is a single DataSet then () is equivalent to (‘dataset’)).

• inclass_correction (dict, optional) – A dictionary with keys ‘dataset’, ‘circuit’, ‘outcome’ and ‘spectrum’, and values that specify the type of multi-test correction used to account for the multiple tests being implemented. This specifies how the statistically significance is maintained within the tests implemented in a single “test class”.

• betweenclass_weighting ('auto' or dict, optional) – The weighting to use to maintain statistical significance between the different classes of test being implemented. If ‘auto’ then a standard Bonferroni correction is used.

• estimator (str, optional) –

  The name of the estimator to use. This is the method used to estimate the parameters of a parameterized model for each probability trajectory, after that parameterized model has been selected with the model selection methods. Allowed values are:

  • ’auto’. The estimation method is chosen automatically, default to the fast method that is also

    reasonably reliable.

  • ’filter’. Performs a type of signal filtering: implements the transform used for generating power

    spectra (e.g., the DCT), sets the amplitudes to zero for all freuquencies that the model selection has not included in the model, inverts the transform, and then performs some minor post-processing to guarantee probabilities within [0, 1]. This method is less statically well-founded than ‘mle’, but it is faster and typically gives similar results. This method is not an option for non-invertable transforms, such as the Lomb-Scargle periodogram.

  • ’mle’. Implements maximum likelihood estimation, on the parameterized model chosen by the model

    selection. The most statistically well-founded option, but can be slower than ‘filter’ and relies on numerical optimization.

• modelselector (tuple, optional) – The model selection method. If not None, a “test class” tuple, specifying which test results to use to decide which frequencies are significant for each circuit, to then construct a parameterized model for each probability trajectory. This can be typically set to None, and it will be chosen automatically. But if you wish to use specific test results for the model selection then this should be set.

• verbosity (int, optional) – The amount of print-to-screen

Return type:

StabilityAnalysis

Methods

__init__([significance, transform, ...])	Implements instability ("drift") detection and characterization on timeseries data from any set of quantum circuits on any number of qubits.
add_mongodb_write_ops(write_ops, mongodb[, ...])	Accumulate write and update operations for writing this object to a MongoDB database.
from_dir(dirname[, quick_load])	Initialize a new Protocol object from dirname.
from_mongodb(mongodb, doc_id, **kwargs)	Create and initialize an object from a MongoDB instance.
from_mongodb_doc(mongodb, collection_name, ...)	Create and initialize an object from a MongoDB instance and pre-loaded primary document.
remove_from_mongodb(mongodb, doc_id[, ...])	Remove the documents corresponding to an instance of this class from a MongoDB database.
remove_me_from_mongodb(mongodb[, session, ...])
run(data[, memlimit, comm])	Run this protocol on data.
run_mpi(data, num_ranks, *[, mpiexec, ...])	Run this protocol in parallel using MPI workers launched as a subprocess.
setup_nameddict(final_dict)	Initializes a set of nested NamedDict dictionaries describing this protocol.
stage_slurm(data, num_ranks, slurm, work_dir, *)	Write all working files to work_dir and generate a SLURM batch script ready for sbatch submission.
write(dirname)	Write this protocol to a directory.
write_to_mongodb(mongodb[, session, ...])	Write this object to a MongoDB database.

Attributes

collection_name

add_mongodb_write_ops(write_ops, mongodb, overwrite_existing=False, **kwargs)

Accumulate write and update operations for writing this object to a MongoDB database.

Similar to write_to_mongodb() but collects write operations instead of actually executing any write operations on the database. This function may be preferred to write_to_mongodb() when this object is being written as a part of a larger entity and executing write operations is saved until the end.

As in write_to_mongodb(), self.collection_name is the collection name and _id is either: 1) the ID used by a previous write or initial read-in, if one exists, OR 2) a new random bson.objectid.ObjectId

Parameters:

• write_ops (WriteOpsByCollection) – An object that keeps track of pymongo write operations on a per-collection basis. This object accumulates write operations to be performed at some point in the future.

• mongodb (pymongo.database.Database) – The MongoDB instance to write data to.

• overwrite_existing (bool, optional) – Whether existing documents should be overwritten. The default of False causes a ValueError to be raised if a document with the given _id already exists and is different from what is being written.

• **kwargs (dict) – Additional keyword arguments potentially used by subclass implementations. Any arguments allowed by a subclass’s _add_auxiliary_write_ops_and_update_doc method is allowed here.

Returns:

The identifier (_id value) of the main document that was written.

Return type:

bson.objectid.ObjectId

classmethod from_dir(dirname, quick_load=False)

Initialize a new Protocol object from dirname.

quick_loadbool, optional

Setting this to True skips the loading of components that may take a long time to load.

Parameters:

• dirname (str) – The directory name.

• quick_load (bool, optional) – Setting this to True skips the loading of components that may take a long time to load.

Return type:

Protocol

classmethod from_mongodb(mongodb, doc_id, **kwargs)

Create and initialize an object from a MongoDB instance.

Parameters:

• mongodb (pymongo.database.Database) – The MongoDB instance to load from.

• doc_id (bson.objecctid.ObjectId or dict) – The object ID or filter used to find a single object ID within the database. This document is loaded from the collection given by the collection_name attribute of this class.

• **kwargs (dict) – Additional keyword arguments potentially used by subclass implementations. Any arguments allowed by a subclass’s _create_obj_from_doc_and_mongodb method is allowed here.

Return type:

object

classmethod from_mongodb_doc(mongodb, collection_name, doc, **kwargs)

Create and initialize an object from a MongoDB instance and pre-loaded primary document.

Parameters:

• mongodb (pymongo.database.Database) – The MongoDB instance to load from.

• collection_name (str) – The collection name within mongodb that doc was loaded from. This is needed for the sole purpose of setting the created (returned) object’s database “coordinates”.

• doc (dict) – The already-retrieved main document for the object being loaded. This takes the place of giving an identifier for this object.

• **kwargs (dict) – Additional keyword arguments potentially used by subclass implementations. Any arguments allowed by a subclass’s _create_obj_from_doc_and_mongodb method is allowed here.

Return type:

object

classmethod remove_from_mongodb(mongodb, doc_id, collection_name=None, session=None, recursive='default')

Remove the documents corresponding to an instance of this class from a MongoDB database.

Parameters:

• mongodb (pymongo.database.Database) – The MongoDB instance to remove documents from.

• doc_id (bson.objectid.ObjectId) – The identifier of the root document stored in the database.

• collection_name (str, optional) – the MongoDB collection within mongodb where the main document resides. If None, then <this_class>.collection_name is used (which is usually what you want).

• session (pymongo.client_session.ClientSession, optional) – MongoDB session object to use when interacting with the MongoDB database. This can be used to implement transactions among other things.

• recursive (RecursiveRemovalSpecification, optional) – An object that filters the type of documents that are removed. Used when working with inter-related experiment designs, data, and results objects to only remove the types of documents you know aren’t being shared with other documents.

Return type:

None

run(data, memlimit=None, comm=None)

Run this protocol on data.

Parameters:

• data (ProtocolData) – The input data.

• memlimit (int, optional) – A rough per-processor memory limit in bytes.

• comm (mpi4py.MPI.Comm, optional) – When not None, an MPI communicator used to run this protocol in parallel.

Return type:

StabilityAnalysisResults

run_mpi(data, num_ranks, *, mpiexec='auto', extra_mpi_args=None, ranks_per_host=None, env=None, persistent_dir=None, dry_run=False, blas_threads_per_rank=0, **run_kwargs)

Run this protocol in parallel using MPI workers launched as a subprocess. The subprocess environment variables will be set as

worker_env = {**_os.environ, **_blas_env, **(env or {})}

where _blas_env is inferred from blas_threads_per_rank (see Notes).

This method can be called from anywhere (e.g., Jupyter notebooks or scripts) without requiring the caller to manage MPI communicators or write launcher scripts manually.

Parameters:

• data (ProtocolData) – The input data.

• num_ranks (int) – Number of MPI worker processes to launch. When 1, falls back to a plain run() call with no MPI overhead.

• mpiexec (str, keyword-only) – MPI launcher executable name or path. 'auto' (default) searches PATH for mpiexec, mpirun, or mpiexec.hydra in that order.

• extra_mpi_args (list of str, keyword-only) – Extra arguments inserted between the launcher and the Python executable, e.g. ['--oversubscribe'] or ['--hostfile', 'hosts.txt'].

• ranks_per_host (int, keyword-only) – Number of ranks per virtual host. Sets PYGSTI_MAX_HOST_PROCS in the worker environment. None (default) uses actual hostnames.

• env (dict, keyword-only) – Extra environment variables forwarded to worker processes, merged on top of the current os.environ. ranks_per_host takes precedence over PYGSTI_MAX_HOST_PROCS supplied here.

• persistent_dir (str or Path, keyword-only) – Permanent directory for data and results. Created if it does not already exist. When None (default), the existing on-disk data path is reused if available; otherwise a temporary directory is used and deleted on return. Required when dry_run=True.

• dry_run (bool, keyword-only) –

  When True, we write working files to persistent_dir and return None. The recommended launch command is printed to console if self.verbosity is unset or positive.

  Requires persistent_dir to be set.

• blas_threads_per_rank (int, keyword-only) – Number of threads each worker rank allows BLAS libraries to use. 0 (default) auto-detects an appropriate value. See Notes.

• **run_kwargs – Keyword arguments forwarded to Protocol.run for each worker. Examples of arguments you might want to forward include simulator, optimizers, or disable_checkpointing. These arguments will be serialized via pickle.

Returns:

None when dry_run=True.

Return type:

ProtocolResults or None

Notes

Shared memory and host topology. pyGSTi uses a two-level communicator hierarchy: a host communicator (host_comm) groups ranks that share a physical host and can exchange data via shared memory, while an inter-host communicator (interhost_comm) connects one rank per host for cross-host MPI. Collective operations (gather, allreduce, broadcast) use a two-phase pattern — first within the host, then across hosts — when shared memory is enabled.

Setting PYGSTI_USE_SHARED_MEMORY=0 (via env) disables shared memory entirely. In this case host_comm is None and all communication goes directly through the full communicator; the host grouping has no effect. ranks_per_host is therefore only meaningful when shared memory is enabled.

Choosing ranks_per_host. On a single compute node all ranks share one hostname, so by default pyGSTi places them all in one host group and uses shared memory across all of them. Setting ranks_per_host=k splits the ranks into num_ranks // k virtual hosts of k ranks each, producing k-rank shared-memory groups connected by MPI across groups. This is useful for testing multi-node behavior on a single machine, or for matching a known NUMA or hardware partition boundary.

Thread oversubscription. When blas_threads_per_rank > 0, that value is set for OMP_NUM_THREADS, OPENBLAS_NUM_THREADS, MKL_NUM_THREADS, NUMEXPR_NUM_THREADS, and BLIS_NUM_THREADS in the worker environment. When blas_threads_per_rank == 0 (the default), the value is computed as max(1, num_cpus // num_ranks) where num_cpus is the physical CPU core count of the current machine (falling back to the logical count if the physical count is unavailable). These variables are applied before env, so explicit entries in env override them.

Job schedulers. To generate a SLURM batch script, use stage_slurm() instead. For other schedulers (PBS, LSF, etc.), use dry_run=True together with persistent_dir to write the data and runner script to a permanent location. The recommended MPI launch command will be written to console if self.verbosity is unset or positive.

setup_nameddict(final_dict)

Initializes a set of nested NamedDict dictionaries describing this protocol.

This function is used by ProtocolResults objects when they’re creating nested dictionaries of their contents. This function returns a set of nested, single (key,val)-pair named-dictionaries which describe the particular attributes of this Protocol object named within its self._nameddict_attributes tuple. The final nested dictionary is set to be final_dict, which allows additional result quantities to easily be added.

Parameters:

final_dict (NamedDict) – the final-level (innermost-nested) NamedDict in the returned nested dictionary.

Return type:

NamedDict

stage_slurm(data, num_ranks, slurm, work_dir, *, ranks_per_host=None, blas_threads_per_rank=0, **run_kwargs)

Write all working files to work_dir and generate a SLURM batch script ready for sbatch submission.

This is the SLURM-specific counterpart to run_mpi(). It does not launch any processes; call sbatch <slurm.script_path> from a terminal or batch system to submit the job.

Parameters:

• data (ProtocolData) – The input data.

• num_ranks (int) – Total number of MPI worker processes.

• slurm (SlurmSettings) – SLURM script options. See SlurmSettings for details.

• work_dir (str or Path) – Permanent directory where data, the protocol, the pickled kwargs, and the runner script are written. All paths embedded in the generated batch script point here.

• ranks_per_host (int, keyword-only) – Number of ranks per shared-memory group. When set, a commented export PYGSTI_MAX_HOST_PROCS=<value> line is emitted in the batch script (uncomment to activate). None (default) uses a suggested value of num_ranks // slurm.nodes.

• blas_threads_per_rank (int, keyword-only) –

  Number of BLAS threads per rank. 0 (default) auto-detects based on the current machine’s CPU count. This value sets both --cpus-per-task in the batch script and the BLAS environment variable exports.

  Note

  Auto-detection measures the machine on which you call stage_slurm (typically a login node), which may differ from the HPC compute nodes. Pass an explicit value if your compute nodes have a different core count.

• **run_kwargs – Forwarded to run() in each worker. Serialized via pickle.

Returns:

script_path – The file path of the generated SLURM script.

Return type:

str

Notes

The generated script uses srun as the MPI launcher. srun is SLURM-native and inherits the job allocation automatically, so no -n flag is needed.

The script also contains a commented block of optional #SBATCH directives (--account, --qos, --constraint, etc.) that are commonly needed but too site-specific to set automatically. Uncomment and edit the relevant lines before submitting.

write(dirname)

Write this protocol to a directory.

Parameters:

dirname (str) – The directory name to write. This directory will be created if needed, and the files in an existing directory will be overwritten.

Return type:

None

write_to_mongodb(mongodb, session=None, overwrite_existing=False, **kwargs)

Write this object to a MongoDB database.

The collection name used is self.collection_name, and the _id is either: 1) the ID used by a previous write or initial read-in, if one exists, OR 2) a new random bson.objectid.ObjectId

Parameters:

• mongodb (pymongo.database.Database) – The MongoDB instance to write data to.

• session (pymongo.client_session.ClientSession, optional) – MongoDB session object to use when interacting with the MongoDB database. This can be used to implement transactions among other things.

• overwrite_existing (bool, optional) – Whether existing documents should be overwritten. The default of False causes a ValueError to be raised if a document with the given _id already exists and is different from what is being written.

• **kwargs (dict) – Additional keyword arguments potentially used by subclass implementations. Any arguments allowed by a subclass’s _add_auxiliary_write_ops_and_update_doc method is allowed here.

Returns:

The identifier (_id value) of the main document that was written.

Return type:

bson.objectid.ObjectId