class, num_fingerprints, arch=None, activation_function=None, weights_initializer=<function random_normal>, bias_initializer=<function zeros>, import_dir=None)[source]

The nueral network is built based on tf.Module.

  • elements (python list) – list of element in the system, e.g. [C,O,H]

  • num_fingerprints (int) – number of fingerprints in dataset

  • arch – network architecture, e.g. [10,10,10], 3 dense layers each with 10 neurons

  • activation_function – any avialiable Tensorflow activation function, e.g. ‘relu’

  • weighs_initializer – the way to initialize weights, default one is tf.random.normal

  • bias_initializer – the way to initialize bias, default one is tf.zeros

  • import_dir – the directory of a saved model to be loaded


Inflate network object from a SavedModel.


imported – a saved tensorflow neural network model

compute_pe(fingerprints, atom_type)[source]

Forward pass of the network to compute potential energy. Parallel networks are used for multiple atom types.

  • fingerprints – 3D array [batch_size,atom_num,fingerprints_num]

  • atom_type – 2D array [batch_size, atom_num]


total potential energy and per-atom potential energy

compute_force_stress(dEdG, input_dict)[source]

Compute force and stress.

  • dEdG – derivatives of per-atom potential energy w.r.t. descriptors

  • input_dict – input dictionary data


force per atom and stress tensor

predict(input_dict, compute_force=True, training=False)[source]

Predict energy, force and stress.

  • input_dict – input dictionary data

  • compute_force (bool) – True to compute force and stress

  • training – True when used during training


potential energy, force and stress

Return type


evaluate(dataset, batch_size=None, return_prediction=True, compute_force=True)[source]

Do evaluation on trained model.

  • dataset – tensorflow dataset used for evaluation

  • batch_size – default is total data size

  • return_prediction (bool) – True for returning prediction resutls

  • compute_force (bool) – True for computing force and stress

loss_function(true, pred)[source]

Build-in tensorflow loss functions can be used. Customized loss function can be also defined here.

  • true (tensorflow tensor) – true output data

  • pred (tensorflow tensor) – predicted results

loss(true_dict, pred_dict, training=False, validation=False)[source]

Compute losses for energy, force and stress.

  • true_dict (dictionary) – dictionary of outputs data

  • pred_dict (dictionary) – dictionary of predicted results

  • training (bool) – True for loss calculation during training

  • validation (bool) – True for loss calculation during validation


if training is true, return total_loss and loss_dict(loss dictionary), otherwise(for evaluation) return loss_dict

train_step(input_dict, output_dict)[source]

A single training step.

  • input_dict – input dictionary data

  • output_dict – output dictionary data


loss dictionary

validation_step(input_dict, output_dict)[source]

A single validation step.

  • input_dict – input dictionary data

  • output_dict – output dictionary data


loss dictionary


Compute scaling factors using training dataset.


Scaling a dataset with the scaling factors calculated with compute_scaling_factors().

train(train_dataset, validation_dataset=None, early_stop=None, scaling=None, batch_size=None, epochs=None, loss_fun=None, optimizer=None, lr=None, loss_weights=None, compute_all_loss=False, shuffle=True, append_loss=False)[source]

Train the neural network.

  • train_dataset (tfdataset) – dataset used for training

  • validation_dataset (tfdataset) – dataset used for validation

  • early_stop (dictionary) – condition for stop training, e.g. {‘train_loss’:[0.01,3]} means stop when train loss is less than 0.01 for 3 times

  • scaling (string) – = None, ‘std’ or ‘norm’

  • batch_size – the training batch size

  • epochs – training epochs

  • loss_fun (string) – loss function, ‘mae’, ‘mse’, ‘rmse’ and others

  • opimizer (string) – any Tensorflow optimizer such as ‘Adam’

  • lr (float) – learning rate

  • loss_weight (dictionary) – weights assigned to loss function, e.g. {‘pe’:1,’force’:1,’stress’:0.1}

  • compute_all_loss (bool) – compute loss for force and stress even when they are not used for training

  • shuffle (bool) – shuffle training dataset during training

  • append_loss (bool) – append loss history

plot_loss(start_epoch=1, figsize=(8, 5))[source]

Plot the losses.

  • start_epoch – plotting starts from start_epoch

  • figsize – set the fingersize, e.g. (8.5)

Functions, model_dir, descriptor=None)[source]

Save a trained model.

  • model_dir – directory for the saved neural network model

  • descriptor (dictionary) – descriptor parameters, used for LAMMPS prediction[source]

Load a saved model.


model_dir – directory of a saved neural network model[source]

Run shell command ‘saved_model_cli show –dir model_dir –tag_set serve –signature_def serving_default’, and get the signature of the network from the shell outputs. SavedModel Command Line Interface (CLI) is a Tensorflow tool to inspect a SavedModel.


model_dir – directory for a saved neural network model[source]

Print the neural network signature.


model_dir – directory for a saved neural network model