Reduce ACE descriptors with PCA and fit a-HfO2 dataset

Setup experiment

Load packages.

using AtomsBase, InteratomicPotentials, PotentialLearning
using Unitful, UnitfulAtomic
using LinearAlgebra, Random, DisplayAs

Define paths.

base_path = haskey(ENV, "BASE_PATH") ? ENV["BASE_PATH"] : "../../"
ds_path   = "$base_path/examples/data/a-HfO2/a-HfO2-300K-NVT-6000.extxyz"
res_path  = "$base_path/examples/PCA-ACE-aHfO2/results/";

Load utility functions.

include("$base_path/examples/utils/utils.jl");

Create experiment folder.

run(`mkdir -p $res_path`);

Load datasets

Load atomistic dataset: atomistic configurations (atom positions, geometry, etc.) + DFT data (energies, forces, etc.)

ds = load_data(ds_path, uparse("eV"), uparse("Å"))[1:1000]; # Load first 1K samples.

Split atomistic dataset into training and test.

n_train, n_test = 50, 50 # Only 50 samples per dataset are used in this example.
conf_train, conf_test = split(ds, n_train, n_test)

(DataSet{num_configs = 50} 
	 Configuration{S, AtomsBase.FlexibleSystem{3, AtomsBase.Atom, Unitful.Quantity{Float64, 𝐋, Unitful.FreeUnits{(Å,), 𝐋, nothing}}}, Energy, Forces}
	 Configuration{S, AtomsBase.FlexibleSystem{3, AtomsBase.Atom, Unitful.Quantity{Float64, 𝐋, Unitful.FreeUnits{(Å,), 𝐋, nothing}}}, Energy, Forces}
	 ⋮
	 Configuration{S, AtomsBase.FlexibleSystem{3, AtomsBase.Atom, Unitful.Quantity{Float64, 𝐋, Unitful.FreeUnits{(Å,), 𝐋, nothing}}}, Energy, Forces}, DataSet{num_configs = 50} 
	 Configuration{S, AtomsBase.FlexibleSystem{3, AtomsBase.Atom, Unitful.Quantity{Float64, 𝐋, Unitful.FreeUnits{(Å,), 𝐋, nothing}}}, Energy, Forces}
	 Configuration{S, AtomsBase.FlexibleSystem{3, AtomsBase.Atom, Unitful.Quantity{Float64, 𝐋, Unitful.FreeUnits{(Å,), 𝐋, nothing}}}, Energy, Forces}
	 ⋮
	 Configuration{S, AtomsBase.FlexibleSystem{3, AtomsBase.Atom, Unitful.Quantity{Float64, 𝐋, Unitful.FreeUnits{(Å,), 𝐋, nothing}}}, Energy, Forces})

Compute descriptors

Create and save ACE basis.

basis = ACE(species           = [:Hf, :O],
            body_order        = 3,
            polynomial_degree = 4,
            rcutoff           = 5.0,
            wL                = 1.0,
            csp               = 1.0,
            r0                = 1.0)
@save_var res_path basis;

Compute ACE descriptors for energy and forces based on the atomistic training configurations.

println("Computing energy descriptors of training dataset...")
e_descr_train = compute_local_descriptors(conf_train, basis;
                                          pbar=false)
println("Computing force descriptors of training dataset...")
f_descr_train = compute_force_descriptors(conf_train, basis;
                                          pbar=false)

50-element Vector{ForceDescriptors}:
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ⋮
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}
 ForceDescriptors{n = 96, d = 3}

Update training dataset by adding energy and force descriptors.

ds_train = DataSet(conf_train .+ e_descr_train .+ f_descr_train);

Reduce descriptor dimension

Reduce dimension of energy and force descriptors in the training dataset.

n_desc = 20
pca = PCAState(tol = n_desc)
fit!(ds_train, pca)
transform!(ds_train, pca);

Learn coefficients

Learn ACE coefficients based on ACE descriptors and DFT data.

println("Learning energies and forces...")
lb = LBasisPotential(basis)
ws, int = [1.0, 1.0], true
learn!(lb, ds_train, ws, int)
@save_var res_path lb.β
@save_var res_path lb.β0
lb.β, lb.β0

([-0.4324525770446828, 0.5058727101898061, -0.751977259533252, -1.1955913360018775, -1.431348593448467, 1.2183093663704279, -4.161132879769713, -16.64069638932541, -39.53727523347586, -7.335922481308213, 45.2517822296508, -3.3544373514336465, 46.73942374497546, 75.700248071378, -20.7920586715346, -73.28983595977158, 50.57463180261907, 163.36894632983143, 24.382147202789657, -107.10445358073386], [-952.5579941858003])

Post-process results

Compute ACE descriptors for energy and forces based on the atomistic test configurations.

println("Computing energy descriptors of test dataset...")
e_descr_test = compute_local_descriptors(conf_test, basis;
                                         pbar = false)
println("Computing force descriptors of test dataset...")
f_descr_test = compute_force_descriptors(conf_test, basis;
                                         pbar = false);

Computing energy descriptors of test dataset...
Computing force descriptors of test dataset...

Update test dataset by adding energy and force descriptors.

ds_test = DataSet(conf_test .+ e_descr_test .+ f_descr_test);

Dimension reduction of energy and force descriptors of test dataset.

transform!(ds_test, pca);

Get and save true and predicted values for energies and forces.

n_atoms_train = length.(get_system.(ds_train))
n_atoms_test = length.(get_system.(ds_test))

e_train, e_train_pred = get_all_energies(ds_train) ./ n_atoms_train,
                        get_all_energies(ds_train, lb) ./ n_atoms_train
f_train, f_train_pred = get_all_forces(ds_train),
                        get_all_forces(ds_train, lb)
@save_var res_path e_train
@save_var res_path e_train_pred
@save_var res_path f_train
@save_var res_path f_train_pred

e_test, e_test_pred = get_all_energies(ds_test) ./ n_atoms_test,
                      get_all_energies(ds_test, lb) ./ n_atoms_test
f_test, f_test_pred = get_all_forces(ds_test),
                      get_all_forces(ds_test, lb)
@save_var res_path e_test
@save_var res_path e_test_pred
@save_var res_path f_test
@save_var res_path f_test_pred;

Compute and save training metrics.

e_train_metrics = get_metrics(e_train, e_train_pred,
                              metrics = [mae, rmse, rsq],
                              label = "e_train")
f_train_metrics = get_metrics(f_train, f_train_pred,
                              metrics = [mae, rmse, rsq, mean_cos],
                              label = "f_train")
train_metrics = merge(e_train_metrics, f_train_metrics)
@save_dict res_path train_metrics
train_metrics

OrderedCollections.OrderedDict{String, Float64} with 7 entries:
  "e_train_mae"      => 0.00230555
  "e_train_rmse"     => 0.00280742
  "e_train_rsq"      => 0.416908
  "f_train_mae"      => 0.245075
  "f_train_rmse"     => 0.31095
  "f_train_rsq"      => 0.636352
  "f_train_mean_cos" => 0.79433

Compute and save test metrics.

e_test_metrics = get_metrics(e_test, e_test_pred,
                             metrics = [mae, rmse, rsq],
                             label = "e_test")
f_test_metrics = get_metrics(f_test, f_test_pred,
                             metrics = [mae, rmse, rsq, mean_cos],
                             label = "f_test")
test_metrics = merge(e_test_metrics, f_test_metrics)
@save_dict res_path test_metrics
test_metrics

OrderedCollections.OrderedDict{String, Float64} with 7 entries:
  "e_test_mae"      => 0.0022566
  "e_test_rmse"     => 0.00265939
  "e_test_rsq"      => 0.355292
  "f_test_mae"      => 0.24496
  "f_test_rmse"     => 0.311153
  "f_test_rsq"      => 0.624556
  "f_test_mean_cos" => 0.781521

Plot and save energy results.

e_plot = plot_energy(e_train, e_train_pred,
                     e_test, e_test_pred)
@save_fig res_path e_plot
DisplayAs.PNG(e_plot)

Plot and save force results.

f_plot = plot_forces(f_train, f_train_pred,
                     f_test, f_test_pred)
@save_fig res_path f_plot
DisplayAs.PNG(f_plot)

Plot and save training force cosine.

e_train_plot = plot_energy(e_train, e_train_pred)
f_train_plot = plot_forces(f_train, f_train_pred)
f_train_cos  = plot_cos(f_train, f_train_pred)
@save_fig res_path e_train_plot
@save_fig res_path f_train_plot
@save_fig res_path f_train_cos
DisplayAs.PNG(f_train_cos)

Plot and save test force cosine.

e_test_plot = plot_energy(e_test, e_test_pred)
f_test_plot = plot_forces(f_test, f_test_pred)
f_test_cos  = plot_cos(f_test, f_test_pred)
@save_fig res_path e_test_plot
@save_fig res_path f_test_plot
@save_fig res_path f_test_cos
DisplayAs.PNG(f_test_cos)

This page was generated using Literate.jl.