Merge 465698d into e6126cf

docs/make.jl

@@ -25,6 +25,7 @@ examples = [
     "Subsample a-HfO2 dataset with DPP and fit with ACE" => "DPP-ACE-aHfO2-1/fit-dpp-ace-ahfo2.jl",
     "Subsample Na dataset with DPP and fit with ACE" => "DPP-ACE-Na/fit-dpp-ace-na.jl",
     "Subsample Si dataset with DPP, fit with ACE, and cross validate" => "DPP-ACE-Si/fit-dpp-ace-si.jl",
+    "Reduce ACE descriptors with PCA and fit a-HfO2 dataset" => "PCA-ACE-aHfO2/fit-pca-ace-ahfo2.jl",
     "Load Ar+Lennard-Jones dataset and postprocess" => "LJ-Ar/lennard-jones-ar.jl"
 ]
 

examples/ACE-aHfO2/fit-ace-aHfO2.jl

@@ -25,14 +25,14 @@ ds = load_data(ds_path, uparse("eV"), uparse("Å"))
 
 # 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)
+conf_train, conf_test = split(ds[1:1000], n_train, n_test)
 
 # ## Create ACE basis, compute descriptors and add them to the dataset.
 
 # Create ACE basis
 basis = ACE(species           = [:Hf, :O],
             body_order        = 3,
-            polynomial_degree = 3,
+            polynomial_degree = 4,
             rcutoff           = 5.0,
             wL                = 1.0,
             csp               = 1.0,

examples/PCA-ACE-aHfO2/Project.toml

@@ -0,0 +1,14 @@
+[deps]
+AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+DisplayAs = "0b91fe84-8a4c-11e9-3e1d-67c38462b6d6"
+InteratomicPotentials = "a9efe35a-c65d-452d-b8a8-82646cd5cb04"
+IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+PotentialLearning = "82b0a93c-c2e3-44bc-a418-f0f89b0ae5c2"
+ProgressBars = "49802e3a-d2f1-5c88-81d8-b72133a6f568"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
+UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"

examples/PCA-ACE-aHfO2/fit-pca-ace-ahfo2.jl

@@ -0,0 +1,158 @@
+# # Reduce ACE descriptors with PCA and fit a-HfO2 dataset
+
+# ## Load packages, define paths, and create experiment folder.
+
+# Load packages
+using AtomsBase, InteratomicPotentials, PotentialLearning
+using Unitful, UnitfulAtomic
+using LinearAlgebra, Random, DisplayAs
+
+# Define paths.
+path = joinpath(dirname(pathof(PotentialLearning)), "../examples/PCA-ACE-aHfO2")
+ds_path =  "$path/../data/a-HfO2/a-HfO2-300K-NVT-6000.extxyz"
+res_path = "$path/results/"
+
+# Load utility functions.
+include("$path/../utils/utils.jl")
+include("$path/pca.jl")
+
+# Create experiment folder.
+run(`mkdir -p $res_path`)
+
+# ## Load atomistic dataset and split it into training and test.
+
+# Load atomistic dataset: atomistic configurations (atom positions, geometry, etc.) + DFT data (energies, forces, etc.)
+ds = load_data(ds_path, uparse("eV"), uparse("Å"))
+
+# 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[1:1000], n_train, n_test)
+
+# ## Create ACE basis, compute descriptors and add them to the dataset.
+
+# Create 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)
+
+# Update training dataset by adding energy and force descriptors.
+ds_train = DataSet(conf_train .+ e_descr_train .+ f_descr_train)
+
+# ## Dimension reduction of energy and force descriptors of training dataset.
+n_desc = 20
+pca = PCAState(tol = n_desc)
+fit!(ds_train, pca)
+transform!(ds_train, pca)
+
+# ## 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
+
+# ## Post-process output: calculate metrics, create plots, and save 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)
+
+# 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 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 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
+
+# Compute 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
+
+# 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)
+