Merge pull request #7 from mila-iqia/ovito_visualisation

PR: Ovito visualisation

crystal_diffusion/analysis/README_ovito.md

@@ -0,0 +1,15 @@
+# OVITO visualisation
+
+## MaxVol Gamma (nbh grades)
+
+The .xyz file resulting from *crystal_diffusion/analysis/ovito_visualisation.py* can be loaded directly in OVITO.
+
+See *examples/local/mtp_to_ovito.sh* for a bash script example.
+
+Once loaded, all atoms appear white. You can change the radius by clicking on **Particles** on the right hand side.
+The menu **Particle display** can customized the view, in particular, the **Standard radius**. It should defaults at 0.5
+
+To get the gamma value, click on **Particles**. On the right hand side near the top, click on the scrolldown menu
+**Add modification** -> **Color coding**. A reasonable start value (around 1) and end value (around 10) should be set.
+**Adjust range (all frames)** will match the colorbar to be the same across the frames. Low gamme value should look
+blue, and high values are red.

crystal_diffusion/analysis/ovito_visualisation.py

@@ -0,0 +1,83 @@
+"""Convert MTP predictions to OVITO readable file.
+
+Draw the MD simulation with the MaxVol values. Some tweaking is required in the OVITO UI. See README_OVITO.md for more
+information.
+"""
+import argparse
+import os
+
+import numpy as np
+import pandas as pd
+import yaml
+
+
+def main():
+    """Read MTP output files and convert to xyz format readable by OVITO."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prediction_file", help="MTP prediction files. Should contain position and MaxVol gamma.")
+    parser.add_argument("--lammps_output", help="LAMMPS output file. Should contain the bounding box information.")
+    parser.add_argument("--output_name", help="Name of the output file that can be loaded by OVITO. ")
+    args = parser.parse_args()
+
+    assert os.path.exists(args.lammps_output), f"LAMMPS out file {args.lammps_output} does not exist."
+
+    lattice = get_lattice_from_lammps(args.lammps_output)
+
+    assert os.path.exists(args.prediction_file), f"Provided prediction file {args.prediction_file} does not exist."
+
+    mtp_predictions_to_ovito(args.prediction_file, lattice, args.output_name)
+
+
+def get_lattice_from_lammps(lammps_output_file: str) -> np.ndarray:
+    """Read periodic bounding box from LAMMPS output file.
+
+    Args:
+        lammps_output_file: lammps output file (dump)
+
+    Returns:
+        lattice: 3x3 array with lattice coordinates
+    """
+    with (open(lammps_output_file, 'r') as f):
+        l_yaml = yaml.safe_load_all(f)
+        for d in l_yaml:  # loop over LAMMPS outputs to get the MD box - we only need the first step
+            # lattice in yaml is 3 x 2 [0, x_lim]
+            # we assume a rectangular lattice for now with the 2nd coordinates as the lattice vectors
+            lattice = np.zeros((3, 3))
+            for i, x in enumerate(d['box']):
+                lattice[i, i] = x[1]
+            break
+    return lattice
+
+
+def mtp_predictions_to_ovito(pred_file: str, lattice: np.ndarray, output_name: str):
+    """Convert output csv to ovito readable file.
+
+    Args:
+        pred_file: output csv to convert
+        lattice: lattice parameters in a 3x3 numpy array
+        output_name: name of resulting file. An .xyz extension is added if not already in the name.
+    """
+    lattice = list(map(str, lattice.flatten()))  # flatten and convert to string for formatting
+    lattice_str = 'Lattice=\"' + " ".join(lattice) + '\" Origin=\"0 0 0\" pbc=\"T T T\"'
+    df = pd.read_csv(pred_file)  # read the predictions
+    xyz_file = ""  # output will be written in file
+    for struct in sorted(df['structure_index'].unique()):  # iterate over LAMMPS steps
+        xyz_values = df.loc[df['structure_index'] == struct, ['x', 'y', 'z']].to_numpy()
+        gamma_values = df.loc[df['structure_index'] == struct, 'nbh_grades'].to_numpy()  # nbh_grade
+        n_atom = xyz_values.shape[0]
+        frame_txt = f"{n_atom}\n"
+        frame_txt += (lattice_str + ' Properties=pos:R:3:MaxVolGamma:R:1\n')
+        # here, we can add properties to filter out atoms or identify some of them
+        for i in range(n_atom):
+            frame_txt += f"{' '.join(map(str, xyz_values[i, :]))} {gamma_values[i]}\n"
+        xyz_file += frame_txt
+
+    if not output_name.endswith(".xyz"):
+        output_name += ".xyz"
+
+    with open(output_name, 'w') as f:
+        f.write(xyz_file)
+
+
+if __name__ == '__main__':
+    main()

crystal_diffusion/train_mtp.py

@@ -2,7 +2,7 @@
 
 Running the main() runs a debugging example. Entry points are train_mtp and evaluate_mtp.
 """
-import os
+import argparse
 from typing import Any, Dict, List, NamedTuple, Tuple
 
 import numpy as np
@@ -13,17 +13,6 @@
 
 from crystal_diffusion.models.mtp import MTPWithMLIP3
 
-MLIP_PATH = os.path.join(os.getcwd(), "mlip-3")
-SAVE_DIR = os.path.join(os.getcwd(), "debug_mlip3")  # for demo only
-
-# TODO list of yaml files should come from an external call
-# yaml dump file
-lammps_yaml = ['examples/local/mtp_example/dump.si-300-1.yaml']
-# yaml thermodynamic variables
-lammps_thermo_yaml = ['examples/local/mtp_example/thermo_log.yaml']
-# note that the YAML output does not contain the map from index to atomic species
-# this will have to be taken from elsewhere
-# use a manual map for now
 atom_dict = {1: 'Si'}
 
 
@@ -55,7 +44,8 @@ def extract_structure_and_forces_from_file(filename: str, atom_dict: Dict[int, A
             coords = [[x[i] for i in coords_idx] for x in d['data']]
             pm_structure = Structure(lattice=lattice,
                                      species=species,
-                                     coords=coords)
+                                     coords=coords,
+                                     coords_are_cartesian=True)
             structures.append(pm_structure)
             force_idx = [d['keywords'].index(x) for x in ['fx', 'fy', 'fz']]
             structure_forces = [[x[i] for i in force_idx] for x in d['data']]
@@ -194,8 +184,19 @@ def get_metrics_from_pred(df_orig: pd.DataFrame, df_predict: pd.DataFrame) -> Tu
 
 def main():
     """Train and evaluate an example for MTP."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--lammps_yaml', help='path to LAMMPS yaml file', required=True, nargs='+')
+    parser.add_argument('--lammps_thermo', help='path to LAMMPS thermo output', required=True, nargs='+')
+    parser.add_argument('--mlip_dir', help='directory to MLIP compilation folder', required=True)
+    parser.add_argument('--output_dir', help='path to folder where outputs will be saved', required=True)
+    args = parser.parse_args()
+
+    lammps_yaml = args.lammps_yaml
+    lammps_thermo_yaml = args.lammps_thermo
+    assert len(lammps_yaml) == len(lammps_thermo_yaml), "LAMMPS outputs yaml should match thermodynamics output."
+
     mtp_inputs = prepare_mtp_inputs_from_lammps(lammps_yaml, lammps_thermo_yaml, atom_dict)
-    mtp = train_mtp(mtp_inputs, MLIP_PATH, SAVE_DIR)
+    mtp = train_mtp(mtp_inputs, args.mlip_dir, args.output_dir)
     print("Training is done")
     df_orig, df_predict = evaluate_mtp(mtp_inputs, mtp)
     energy_mae, force_mae = get_metrics_from_pred(df_orig, df_predict)

examples/local/mtp_example/train_mtp.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+# Change this path to the root folder of the repo
+ROOT_DIR="${HOME}/ic-collab/courtois_collab/crystal_diffusion"
+
+MLIP_PATH="${ROOT_DIR}/mlip-3"
+SAVE_DIR="${ROOT_DIR}/debug_mlip3"
+LAMMPS_YAML="${ROOT_DIR}/examples/local/mtp_example/dump.si-300-1.yaml"
+LAMMPS_THERMO="${ROOT_DIR}/examples/local/mtp_example/thermo_log.yaml"
+
+python crystal_diffusion/train_mtp.py \
+    --lammps_yaml $LAMMPS_YAML \
+    --lammps_thermo $LAMMPS_THERMO \
+    --mlip_dir $MLIP_PATH \
+    --output_dir $SAVE_DIR

examples/local/mtp_to_ovito.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+MTP_PREDICTION=./experiments/mtp_find_region/predictions.csv
+LAMMPS_OUTPUT=./experiments/mtp_find_region/dump.si-10000-4.yaml
+OVITO_OUTPUT=./test_si_structure_ovito.xyz
+
+python crystal_diffusion/analysis/ovito_visualisation.py \
+    --prediction_file $MTP_PREDICTION \
+    --lammps_output $LAMMPS_OUTPUT \
+    --output_name $OVITO_OUTPUT