mlspm.utils#

class mlspm.utils.Checkpointer(model: Module, optimizer: Optimizer, additional_data: dict = {}, checkpoint_dir: str = './Checkpoints', keep_last_epoch: bool = True)[source]#

Bases: object

Keep checkpoints of a Pytorch model and optimizer over epochs, keeping only the one with best loss. Also load the latest checkpoint in the beginning if any exist.

Parameters:

  • model – Pytorch model.

  • optimizer – Pytorch optimizer.

  • additional_module – Additional modules whose state will be saved to the checkpoints.

  • checkpoint_dir – Path to directory where checkpoints are saved.

  • keep_last_epoch – Also keep the last epoch even if it does not have the best loss.

property best_epoch#
property best_loss#
next_epoch(loss: float)[source]#

Advance epoch and save state if loss improved.

Parameters:

loss – Loss value for the current epoch.

revert_to_best_epoch()[source]#

Revert model state to the best epoch.

mlspm.utils.batch_write_xyzs(xyzs: list[ndarray], outdir: str = './', start_ind: int = 0, verbose: int = 1)[source]#

Write a batch of xyz files 0_mol.xyz, 1_mol.xyz, …

Parameters:

  • xyzs – Molecules to write.

  • outdir – Directory where files are saved.

  • start_ind – Index where file numbering starts.

  • verbose – 0 or 1. Whether to print output information.

mlspm.utils.count_parameters(module: Module) → int[source]#

Count trainable parameters in a Pytorch module.

Parameters:

module – Pytorch module.

mlspm.utils.load_checkpoint(model: Module, optimizer: Optimizer = None, file_name: str = './model.pth', additional_data: dict | None = None, rank: int | None = None)[source]#

Load a Pytorch model/optimizer checkpoint.

Parameters:

  • model – Model whose state to load.

  • optimizer – Optimizer whose state to load.

  • file_name – Checkpoint file to load from.

  • additional_data – If not None, a dictionary with additional modules or other data that will be updated with additional data found in the checkpoint.

  • rank – Process rank for distributed training.

mlspm.utils.read_xyzs(file_paths: list[str], return_comment: bool = False) → list[ndarray][source]#

Read molecule xyz files.

Parameters:

  • file_paths – Paths to xyz files

  • return_comment – If True, also return the comment string on second line of file.

Returns:

Arrays of shape (num_atoms, 4) or (num_atoms, 5). Each row in the arrays corresponds to one atom with [x, y, z, element] or [x, y, z, charge, element].

mlspm.utils.save_checkpoint(model: Module, optimizer: Optimizer, epoch: int, save_dir: str, additional_data: dict = {})[source]#

Save a Pytorch model/optimizer checkpoint.

Parameters:

  • model – Model whose state to save.

  • optimizer – Optimizer whose state to save.

  • epoch – Training epoch.

  • save_dir – Directory to save in.

  • additional_data – A dictionary of additional modules or data to save to the checkpoint.

mlspm.utils.write_to_xyz(molecule: ndarray, outfile: str = './pos.xyz', comment_str: str = '', verbose: int = 1)[source]#

Write molecule into xyz file.

Parameters:

  • molecule – Molecule to write. np.array of shape (num_atoms, 4) or (num_atoms, 5). Each row corresponds to one atom with [x, y, z, element] or [x, y, z, charge, element].

  • outfile – Path where xyz file will be saved.

  • comment_str – Comment written to the second line of the file.

  • verbose – 0 or 1. Whether to print output information.