Utilities

This module contains utilities for use with various aspects of the Amp calculator.

Module contents

class amp.utilities.Annealer(calc, images, Tmax=None, Tmin=None, steps=None, updates=None, train_forces=True)

Bases: object

Inspired by the simulated annealing implementation of Richard J. Wagner <wagnerr@umich.edu> and Matthew T. Perry <perrygeo@gmail.com> at https://github.com/perrygeo/simanneal.

Performs simulated annealing by calling functions to calculate loss and make moves on a state. The temperature schedule for annealing may be provided manually or estimated automatically.

Can be used by something like:

>>> from amp import Amp
>>> from amp.descriptor.gaussian import Gaussian
>>> from amp.model.neuralnetwork import NeuralNetwork
>>> calc = Amp(descriptor=Gaussian(), model=NeuralNetwork())

which will initialize tha calc object as usual, and then

>>> from amp.utilities import Annealer
>>> Annealer(calc=calc, images=images)

which will perform simulated annealing global search in parameters space, and finally

>>> calc.train(images=images)

for gradient descent optimization.

Parameters:

• calc (object) – Amp calculator.

• images (dict) – Dictionary of images.

• Tmax (float) – Maximum temperature.

• Tmin (float) – Minimum temperature.

• steps (int) – Number of iterations.

• updates (int) – Number of updates.

• train_forces (bool) – Turn off forces.

Tmax = 20.0

Tmin = 2.5

anneal()

Minimizes the loss of a system by simulated annealing.

Parameters:

state – An initial arrangement of the system

Returns:

The best state and loss found.

Return type:

state, loss

auto(minutes, steps=2000)

Minimizes the loss of a system by simulated annealing with automatic selection of the temperature schedule.

Keyword arguments: state – an initial arrangement of the system minutes – time to spend annealing (after exploring temperatures) steps – number of steps to spend on each stage of exploration

Returns the best state and loss found.

copy_state(state)

Returns an exact copy of the provided state Implemented according to self.copy_strategy, one of

• deepcopy : use copy.deepcopy (slow but reliable)

• slice: use list slices (faster but only works if state is list-like)

• method: use the state’s copy() method

copy_strategy = 'copy'

get_loss(state)

Calculate state’s loss

move(state)

Create a state change

static round_figures(x, n)

Returns x rounded to n significant figures.

save_state(fname=None)

Saves state

save_state_on_exit = False

set_schedule(schedule)

Takes the output from auto and sets the attributes

set_user_exit(signum, frame)

Raises the user_exit flag, further iterations are stopped

steps = 10000

static time_string(seconds)

Returns time in seconds as a string formatted HHHH:MM:SS.

update(step, T, L, acceptance, improvement)

Prints the current temperature, loss, acceptance rate, improvement rate, elapsed time, and remaining time.

The acceptance rate indicates the percentage of moves since the last update that were accepted by the Metropolis algorithm. It includes moves that decreased the loss, moves that left the loss unchanged, and moves that increased the loss yet were reached by thermal excitation.

The improvement rate indicates the percentage of moves since the last update that strictly decreased the loss. At high temperatures it will include both moves that improved the overall state and moves that simply undid previously accepted moves that increased the loss by thermal excititation. At low temperatures it will tend toward zero as the moves that can decrease the loss are exhausted and moves that would increase the loss are no longer thermally accessible.

updates = 50.0

user_exit = False

exception amp.utilities.ConvergenceOccurred

Bases: Exception

Kludge to decide when scipy’s optimizers are complete.

class amp.utilities.Data(filename, db=<class 'amp.utilities.FileDatabase'>, calculator=None)

Bases: object

Serves as a container (dictionary-like) for (key, value) pairs that also serves to calculate them.

Works by default with FileDatabase, which writes each entry to disk as a pickle file. Pass db=EphemeralDatabase for inference or MD runs where images are never revisited; that backend holds only the most recent entry in memory and never writes to disk.

Designed to hold things like neighborlists, which have a hash, value format.

This will work like a dictionary in that items can be accessed with data[key], but other advanced dictionary functions should be accessed with through the .d attribute:

>>> data = Data(...)
>>> data.open()
>>> keys = data.d.keys()
>>> values = data.d.values()

calculate_items(images, parallel, log=None)

Calculates the data value with ‘calculator’ for the specified images.

images is a dictionary, and the same keys will be used for the current database.

close()

Safely close the database.

open(mode='r')

Open the database connection with mode specified.

class amp.utilities.EphemeralDatabase(filename)

Bases: object

In-memory drop-in for FileDatabase that never writes to disk.

Intended for inference / MD use, where each image is visited once and will never be revisited. Using FileDatabase in that context writes a growing pile of pickle files that wastes disk I/O and process memory.

All data for a given filename is shared across open()/close() cycles via a class-level store, so fingerprints computed in one phase of a single calculate() call remain accessible in later phases (important for KRR, which needs both test and training fingerprints simultaneously). The store is cleared once per Amp.calculate() call via clear_all(), keeping memory bounded to the working set of one inference step.

The interface mirrors FileDatabase / shelve so it can be passed as the db argument to Data.

classmethod clear_all()

Clear all stored data. Called at the start of each Amp.calculate() to bound memory to one inference step’s working set.

close()

keys()

classmethod open(filename, flag=None)

update(d)

class amp.utilities.FileDatabase(filename)

Bases: object

Using a database file, such as shelve or sqlitedict, that can handle multiple processes writing to the file is hard.

Therefore, we take the stupid approach of having each database entry be a separate file. This class behaves essentially like shelve, but saves each dictionary entry as a plain pickle file within the directory, with the filename corresponding to the dictionary key (which must be a string).

Like shelve, this also keeps an internal (memory dictionary) representation of the variables that have been accessed.

Also includes an archive feature, where files are instead added to a file called ‘archive.tar.gz’ to save disk space. If an entry exists in both the loose and archive formats, the loose is taken to be the new (correct) value.

archive()

Cleans up to save disk space and reduce huge number of files.

That is, puts all files into an archive. Compresses all files in <path>/loose and places them in <path>/archive.tar.gz. If archive exists, appends/modifies.

close()

Only present for compatibility with shelve.

keys()

Return list of keys, both of in-memory and out-of-memory items.

classmethod open(filename, flag=None)

Open present for compatibility with shelve. flag is ignored; this format is always capable of both reading and writing.

update(newitems)

values()

Return list of values, both of in-memory and out-of-memory items. This moves all out-of-memory items into memory.

class amp.utilities.Logger(file, overwrite=False)

Bases: object

Logger that can also deliver timing information.

Parameters:

file (str) – File object or path to the file to write to. Or set to None for a logger that does nothing.

tic(label=None)

Start a timer.

Parameters:

label (str) – Label for managing multiple timers.

class amp.utilities.MessageDictionary(process_id)

Bases: object

Standard container for all messages (typically requests, via zmq.context.socket.send_pyobj) sent from the workers to the master.

This returns a simple dictionary. This is roughly email format. Initialize with process id (e.g., ‘from’). Call with subject and data (body).

class amp.utilities.MetaDict

Bases: dict

Dictionary that can also store metadata. Useful for images dictionary so that images can still be iterated by keys.

metadata = {}

exception amp.utilities.MissingDataError

Bases: Exception

Error to be raised if any images are missing key data, like energy or forces.

exception amp.utilities.TrainingConvergenceError

Bases: Exception

Error to be raised if training does not converge.

amp.utilities.assign_cores(cores, log=None)

Tries to guess cores from environment.

If fed a log object, will write its progress.

amp.utilities.check_images(images, forces, charges)

Checks that all images have energies, and optionally forces and charges, calculated, so that they can be used for training. Raises a MissingDataError if any are missing.

amp.utilities.enforce_magnetic_moments(atoms, supplied_magmom_dict=None)

amp.utilities.extract_an_atomic_chunk(atoms, index, parent_calc=None, cutoff=6.5, vacuum=5.0)

Extract a chunk from atoms centering on the atom with a given index. cutoff defines the range within which atoms are included in the atomic chunk. vacuum represents the thickness of surrounded vacuum layers.

amp.utilities.get_atomic_uncertainties(load, atoms, force=True, label='amp', threshold=None)

Compute atomic uncertainties based on ensemble predictions of either atomic energies or forces. If threshold is specificed, indices of atoms whose atomic uncertainty is larger than the threshold will be identified.

amp.utilities.get_gc_hash(atoms, electrode_potential=None)

Creates a unique signature for a particular ASE atoms object

calculated within grand-canonical framework SJM.

This is used to check whether an image has been seen before. Two identical images are defined as identical both in geometry and electrode potential. This is just an md5 hash of a string representation of the atoms object and target electrode potential.

Parameters:

• atoms (ASE dict) – ASE atoms object.

• electrode_potential (float or None) – If provided, use this value directly (inference path: potential was set on the Amp calculator via calc.set(electrode_potential=…)). If None, read from atoms.calc.results[‘electrode_potential’] (training path: images from SJM/GPAW trajectory).

Return type:

Hash string key of ‘atoms’.

amp.utilities.get_hash(atoms)

Creates a unique signature for a particular ASE atoms object.

This is used to check whether an image has been seen before. This is just an md5 hash of a string representation of the atoms object.

Parameters:

atoms (ASE dict) – ASE atoms object.

Return type:

Hash string key of ‘atoms’.

amp.utilities.get_overfit_mask(nn_model, parametervector)

Compute masked indices for overfit parameters in a raveled parametervector. The basic principle for overfit/regularization is that only weights and slopes should be regularized. In other words, biases and intercepts should be removed. This function is designed for the neural network model in Amp.

amp.utilities.hash_gc_images(images, log=None, ordered=False)

Converts input images calculated within grand-canonical framework SJM including geometry and electrode potential info – which may be a list, a trajectory file, or a database – into a dictionary indexed by their hashes.

Returns this dictionary. If ordered is True, returns an OrderedDict. When duplicate images are encountered (based on encountering an identical hash), a warning is written to the logfile. The number of duplicates of each image can be accessed by examinging dict_images.metadata[‘duplicates’], where dict_images is the returned dictionary.

amp.utilities.hash_images(images, log=None, ordered=False)

Converts input images – which may be a list, a trajectory file, or a database – into a dictionary indexed by their hashes.

Returns this dictionary. If ordered is True, returns an OrderedDict. When duplicate images are encountered (based on encountering an identical hash), a warning is written to the logfile. The number of duplicates of each image can be accessed by examinging dict_images.metadata[‘duplicates’], where dict_images is the returned dictionary.

amp.utilities.importer(name)

Handles strange import cases, like pxssh which might show up in either the package pexpect or pxssh.

amp.utilities.make_filename(label, base_filename)

Creates a filename from the label and the base_filename which should be a string.

Returns None if label is None; that is, it only saves output if a label is specified.

Parameters:

• label (str) – Prefix.

• base_filename (str) – Basic name of the file.

amp.utilities.make_sublists(masterlist, n)

Randomly divides the masterlist into n sublists of roughly equal size.

The intended use is to divide a keylist and assign keys to each task in parallel processing. This also destroys the masterlist (to save some memory).

amp.utilities.now(with_utc=False)

Return type:

String of current time.

amp.utilities.parse_slurm_allocation(env_vars=None, log=None)

If debugging you can pass in a dictionary with custom environment variables.

amp.utilities.randomize_images(images, fraction=0.8)

Randomly assigns ‘fraction’ of the images to a training set and (1 - ‘fraction’) to a test set. Returns two lists of ASE images.

Parameters:

• images (list or str) – List of ASE atoms objects in ASE format. This can also be the path to an ASE trajectory (.traj) or database (.db) file.

• fraction (float) – Portion of train_images to all images.

Returns:

train_images, test_images – Lists of train and test images.

Return type:

list

amp.utilities.read_amp_log_status(logfile)

Reads a partially or fully written amp-log.txt for live job monitoring.

Designed to be safe on logs at any stage of writing: reads only the first 30 lines for the start time and tails the last ~2 KB for the step count, so it does not block or fail on large or incomplete log files.

Uses the ‘Date:’ header line (written at Python process startup, before fingerprinting) as the start-time anchor — giving the true wall-clock age of the job including all setup time. See read_trainlog() in analysis.py for full post-analysis parsing.

Parameters:

logfile (str) – Path to the amp-log.txt file.

Returns:

• start_time (float or None) – UTC timestamp (seconds since epoch) parsed from the ‘Date:’ header, or None if the line has not yet been written.

• last_step (int or None) – Most recent optimizer step number, or None if training has not yet reached the parameter optimization stage.

amp.utilities.setup_parallel(parallel, workercommand, log, setup_publisher=False)

Starts the worker processes and the master to control them.

This makes an SSH connection to each node (including the one the master process runs on), then creates the specified number of processes on each node through its SSH connection. Then sets up ZMQ for efficienty communication between the worker processes and the master process.

Uses the parallel dictionary as defined in amp.Amp. log is an Amp logger. module is the name of the module to be called, which is usually given by self.calc.__module, etc. workercommand is stub of the command used to start the servers, typically like “python -m amp.descriptor.gaussian”. Appended to this will be “ <pid> <serversocket> &” where <pid> is the unique ID assigned to each process and <serversocket> is the address of the server, like ‘node321:34292’.

If setup_publisher is True, also sets up a publisher instead of just a reply socket.

Returns:

server – The ssh connections (pxssh instances; if these objects are destroyed pxssh will close the sessions)

the pid_count, which is the total number of workers started. Each worker can be communicated directly through its PID, an integer between 0 and pid_count

Return type:

(a ZMQ socket)

amp.utilities.start_workers(process_ids, workerhostname, workercommand, log, envcommand)

A function to start a new SSH session and establish processes on that session.

amp.utilities.string2dict(text)

Converts a string into a dictionary.

Basically just calls eval on it, but supplies words like OrderedDict and matrix.