dptools package

Subpackages

Submodules

dptools.env module

Functions for controlling different dptools env configurations (deepmd models, sbatch parameters, etc.). i.e. controls things set with CLI set command.

Mostly indented to be used behind the scenes when using CLI commands, but theoretically could be used in your python scripts as well. Slightly confusing usage with global env_file (sorry), so use at your own risk.

dptools.env.clear(keys): Clear specific key-values (or entire env) for loaded global env file.

dptools.env.clear_model(): Clear model related key-values from loaded global env file.

dptools.env.get_dpfaults(key='model')

Like defaults but for dp (haha… ha..). Loads specific key-values from env file.

Parameters:

key (str) –

(model, ensemble, sbatch) Gets env information from env file.

if key is set to model, return env’s deepmd model path and type map
if key is set to ensemble, return all model paths belonging to env’s ensemble
if key is set to sbatch, return env’s Slurm settings

dptools.env.get_env(): Get all key-values from global env file.

dptools.env.load(label): Sets global env file if different from default .env. Used to load custom envs essentially.

dptools.env.set_default_sbatch(warn=True): Sets Slurm parameters from dptools.hpc.hpc_defaults. Mostly used by CLI when no Slurm info has been set to env.

dptools.env.set_env(key, value): Set key-value to global env file.

dptools.env.set_model(model, n_model=''): Set path to deepmd model .pb file to use during simulations evoked by CLI run command.

dptools.env.set_params(params): Save set of simulation parameters to run with CLI run command.

dptools.env.set_sbatch(script)

Reads script and searches for any #SBATCH line to set to env file.

Parameters:: script (str) – Path to .sh script with #SBATCH comments to set Slurm params.

dptools.env.set_training_params(in_json): Save deepmd-kit training parameters to use with CLI train command.

dptools.hpc module

Module for controlling HPC job submissions. Currently only supports Slurm jobs.

class dptools.hpc.SlurmJob(sbatch_comment, commands='', directories='.', file_name='script.sh', zip_commands=False, **kwargs)

Bases: object

Sets up and optionally submits slurm job scripts on HPC systems. Default settings for your HPC system can be added to dptools.hpc.hpc_defaults, or controlled using the CLI set command.

Parameters:

sbatch_comment (str) – Single line containing all Slurm #SBATCH parameters. e.g., '#SBATCH -J job -q regular -N 1 --time=11:00:00'
directories (list[str] or str) – Path(s) to job submission directories.
file_name (str) – Name of submission script to write. Currently only supports .sh scripts, but support for others (e.g., .py) can be easily added upon request.
zip_commands (bool) –
Set to True if each submission dir requires a unique set of commands. If True, then len(commands) must equal len(directories), and each item in commands is used for the respective directory index. i.e., job scripts are written as,
```
for command, dir in zip(commands, directories):
    write_script(command=command, submission_dir=dir)
```
If False, then the same command(s) is used in all submission directories.
**kwargs –
Unpacked dict containing any env variables to set in submission script. e.g. kwargs = dict(TF_INTRA_OP_PARALLELISM_THREADS="1") adds this line to .sh script,
```
export TF_INTRA_OP_PARALLELISM_THREADS=1
```

generate_commands()

get_header()

set_path_stuff(directories, file_name)

set_text(**kwargs)

submit()

property text

write(sub=False)

write_script(path)

dptools.utils module

Assorted utilities for plots, reading things, converting things, etc.

class dptools.utils.Converter(inputs, output, indices=':')

Bases: object

Class to convert between different ASE/VASP/LAMMPS outputs. Mostly useful if wanting to concatenate all MD images from variable-time jobs, but also supports lammps dump conversions and common ASE formats.

Parameters:

inputs (list[str]) – Paths to input structure files to convert. Images from each input are concatenated in output.
output (str) – Name of file with desired conversion extension specified, e.g. 'out.traj'
indices (str) – Index slice, e.g. '::10', :, '1:100', etc.

Example

# concat MD images (ignores 1st image if identical to final image in previous input)
>>> from dptools.utils import Converter
>>> inputs = ['md_000/vasprun.xml', 'md_001/vasprun.xml', 'md_002/vasprun.xml']
>>> output = 'md.traj'
>>> converter = Converter(inputs, outputs)
>>> converter.convert()

convert(**kwargs)

read(**kwargs)

set_reader()

set_types()

dptools.utils.check_type_map(type_map_dict)

dptools.utils.columnize(*data): Takes lists or 1D arrays and concatenates everything into columnized array. Basically just np.column_stack without needing a single tuple arg (i.e. slightly useless).

dptools.utils.convert_dump_cell(lammps_cell)

Converts lammps dump cell format to ase cell.

Parameters:

lammps_cell (numpy.array) – Simulation cell from lammps dump file.

Returns:

cell (numpy.array) – New 3x3 array to use with ase.
shift (numpy.array) – xyz shifts to place atoms at origin aligned with cell.

dptools.utils.get_seed(max_val=999999, n=1)

Returns:: seed (int or list[int]) – n random ints between 0 and max_val.

dptools.utils.graph2typemap(graph)

Determine type_map for a given deepmd model.

Parameters:: graph (str) – Path to deepmd model .pb file.
Returns:: type_map (dict) – Dictionary that maps each atomic symbol to type map index

dptools.utils.next_color()

Convenience function to grab a new color during plotting loops.

Returns:: color (tuple[float]) – New color rgb vals from seaborn.color_palette(‘deep’)

Example

for values in all_values:
    plt.plot(values, color=next_color())

dptools.utils.randomize_seed(in_json)

Take input .json or dict with training parameters and randomize seeds for model ensemble training.

Parameters:: in_json (str or dict) – Path to .json or dict with deepmd training parameters.
Returns:: (dict) – Training parameter dictionary with new randomized seed values.

dptools.utils.read_db(db_name, indices)

Reads ase db and returns entries as list of Atoms objects. Useless, just use ase.io.read(db_name, index=indices) instead.

Parameters:

db_name (str) – Path to ase db file.
indices (str) – Index slice, e.g. ::10, :, 1:100, etc.

Returns:

traj (list[ase.Atoms]) – db entries as list of Atoms objects.

dptools.utils.read_dump(dump, type_map, index=':')

Reads in lammps dump file and returns corresponding ase.Atoms list.

Parameters:

dump (str) – Path to dump file to read.
type_map (dict) – Dictionary with element-index mapping, e.g. {‘Si’: 0, ‘O’: 1}
index (str) – index slice to control which images are returned, e.g. ‘:’, ‘::100’, etc.

Returns:

traj (list[ase.Atoms]) – List of dump images as ase.Atoms objects

dptools.utils.read_type_map(type_map_json)

dptools.utils.str2typemap(tm_str)

dptools.utils.tag_groups(atoms, mult=1.0, retag=True)

Takes Atoms object of disconnected groups (e.g. adsorbates in MOF) and assigns a unique tag to each group. Returns list of nested lists with indices corresponding to each tag

Parameters:

atoms (ase.Atoms) – Atoms object to tag.
mult (float) – Coefficient for adjusting neighbor search distances (ase.neighborlist.natural_cutoffs function).
retag (bool) – If false, a copy of atoms is used for tagging, leaving the original tags unmodified.

Returns:

groups (dict) – dict of all indices and elements corresponding to each group of tagged atoms.

dptools.utils.typemap2str(type_map_dict)