utils module

class edafm.utils.LossLogPlot(log_path, plot_path, loss_labels, loss_weights=None)[source]

Bases: object

Log and plot model training loss history.

Parameters

  • log_path – str. Path where loss log is saved.

  • plot_path – str. Path where plot of loss history is saved.

  • loss_labels – list of str. Labels for different loss components.

  • loss_weights – list of int or str. Weights for different loss components. Empty string for no weight (e.g. Total loss).

add_losses(train_loss, val_loss)[source]

Add losses to log.

Parameters

  • train_loss – list of floats of length len(self.loss_labels). Training losses for the epoch.

  • val_loss – list of floats of length len(self.loss_labels). Validation losses for the epoch.

plot_history(show=False, verbose=1)[source]

Plot and save history of current losses into plot_path.

Parameters

  • show – Bool. Whether to show the plot on screen.

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

edafm.utils.count_parameters(module)[source]

Count pytorch module parameters.

Parameters

module – torch.nn.Module.

edafm.utils.download_molecules(save_path='./Molecules', verbose=1)[source]

Download database of molecules.

Parameters

  • save_path – str. Path where the molecule xyz files will be saved.

  • verbose – int 0 or 1. Whether to print progress information.

edafm.utils.download_weights(weights_type='base', weights_dir='./weights', verbose=1)[source]

Download pretrained weights for EDAFMNet model.

Parameters

  • weights_type – str. Type of weights to download. One of ‘base’, ‘single-channel’, ‘CO-Cl’, ‘Xe-Cl’, ‘constant-noise’, ‘uniform-noise’, ‘no-gradient’, or ‘matched-tips’. See README at https://github.com/SINGROUP/ED-AFM for explanations for the different options.

  • weights_dir – str. Directory where the weight will be downloaded into.

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

edafm.utils.load_checkpoint(model, optimizer=None, file_name='./model.pth', lr_scheduler=None, verbose=1)[source]

Load pytorch checkpoint.

Parameters

  • model – torch.nn.Module. Model where parameters are loaded to.

  • optimizer – torch.optim.Optimizer or None. If not None, load state to this optimizer.

  • file_name – str. Checkpoint file to load from.

  • lr_scheduler – torch.optim.lr_scheduler or None. If not None, try loading state to this scheduler.

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

edafm.utils.read_xyzs(file_paths, return_comment=False)[source]

Read molecule xyz files.

Parameters

  • file_paths – list of str. Paths to xyz files

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

Returns: list of 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].

edafm.utils.save_checkpoint(model, optimizer, epoch, save_dir, lr_scheduler=None, verbose=1)[source]

Save pytorch checkpoint.

Parameters

  • model – torch.nn.Module. Model whose state to save.

  • optimizer – torch.optim.Optimizer. Optimizer whose state to save

  • epoch – int. Training epoch.

  • save_dir – str. Directory to save in.

  • lr_scheduler – torch.optim.lr_scheduler or None. If not None, save state of this scheduler.

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