Skip to content

Commit

Permalink
Moredocs nico (#240)
Browse files Browse the repository at this point in the history 
* Model training example

* added eval to 01 docs


---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
@Tetracarbonylnickel @pre-commit-ci
Tetracarbonylnickel and pre-commit-ci[bot] authored Mar 6, 2024
1 parent f407970 commit 70423bd
Show file tree
Hide file tree
Showing 3 changed files with 606 additions and 70 deletions.
2 changes: 1 addition & 1 deletion apax/config/md_config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ class MDConfig(BaseModel, frozen=True, extra="forbid"):
JaxMD allocates a maximal number of neighbors.
This argument lets you add additional capacity to avoid recompilation.
The default is usually fine.
intitial_structure:
initial_structure:
Path to the starting structure of the simulation.
sim_dir:
Directory where simulation file will be stored.
Expand Down
399 changes: 345 additions & 54 deletions examples/01_Model_Training.ipynb
View file

Large diffs are not rendered by default.

275 changes: 260 additions & 15 deletions examples/02_Molecular_Dynamics.ipynb
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,23 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import warnings\n",
"warnings.simplefilter('ignore')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Molecular Dynamics\n",
"\n",
"In this tutorial we will cover how to use trained models to drive MD simulations.\n",
"For this purpose, apax offers two options: ASE and JaxMD.\n",
"For this purpose, apax offers two options: ASE and JaxMD. Keep in mind that JaxMD can be GPU/TPU accelerated and is therefore much faster.\n",
"Both will be covered below."
]
},
Expand All @@ -18,7 +28,72 @@
"## Basic Model Training\n",
"\n",
"First we need to train a model.\n",
"If you have the parameters from tutorial 01, you can point the paths to those models and skip the current section."
"If you have the parameters from tutorial 01, you can point the paths to those models and skip the current section to the [ASE MD](##-The-ASE-calculator) or the [JaxMD](##-JaxMD) section."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"!apax template train # generating the config file in the cwd"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Precomputing NL: 100%|███████████████████████████████████████| 1000/1000 [00:00<00:00, 12924.12it/s]\n",
"Precomputing NL: 100%|█████████████████████████████████████████| 100/100 [00:00<00:00, 11632.43it/s]\n",
"Epochs: 100%|██████████████████████████████████████| 100/100 [03:36<00:00, 2.17s/it, val_loss=0.31]\n"
]
}
],
"source": [
"from pathlib import Path\n",
"from apax.utils.datasets import download_benzene_DFT, mod_md_datasets\n",
"from apax.train.run import run\n",
"from apax.utils.helpers import mod_config\n",
"import yaml\n",
"\n",
"\n",
"# Download and modify the dataset\n",
"data_path = Path(\"project\")\n",
"experiment = \"benzene_md\"\n",
"\n",
"file_path = download_benzene_DFT(data_path)\n",
"file_path = mod_md_datasets(file_path)\n",
"\n",
"\n",
"# Modify the config file (can be done manually)\n",
"config_path = Path(\"config.yaml\")\n",
"\n",
"config_updates = {\n",
" \"n_epochs\": 100,\n",
" \"data\": {\n",
" \"experiment\": experiment,\n",
" \"directory\": str(data_path / \"models\"),\n",
" \"data_path\": str(file_path),\n",
" \"energy_unit\": \"kcal/mol\",\n",
" \"pos_unit\": \"Ang\",\n",
" }\n",
"}\n",
"config_dict = mod_config(config_path, config_updates)\n",
"\n",
"\n",
"# dump config for cli showcase\n",
"with open(\"config.yaml\", \"w\") as conf:\n",
" yaml.dump(config_dict, conf, default_flow_style=False)\n",
"\n",
"\n",
"# Train model\n",
"run(config_dict)\n"
]
},
{
Expand All @@ -29,11 +104,47 @@
"\n",
"If you require some ASE features during your simulation, we provide an alternative to the JaxMD interface.\n",
"\n",
"An ASE calculator of a trained model can be instantiated as follows\n",
"Please refer to the [ASE documentation](https://wiki.fysik.dtu.dk/ase/ase/calculators/calculators.html) to see how to use ASE calculators.\n",
"\n",
"An ASE calculator of a trained model can be instantiated as follows."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"from ase.io import read\n",
"from apax.md import ASECalculator\n",
"from ase.md.langevin import Langevin\n",
"from ase import units\n",
"import numpy as np\n",
"from ase.io.trajectory import Trajectory\n",
"\n",
"\n",
"# read starting structure and define modelpath\n",
"atoms = read(file_path, index=0)\n",
"model_dir = data_path / f\"models/{experiment}\"\n",
"\n",
"\n",
"# initiolize the apax ase calculator and assign it to the starting structure\n",
"calc = ASECalculator(model_dir=model_dir)\n",
"atoms.calc = calc\n",
"\n",
"\n",
"CODE\n",
"# perform MD simulation\n",
"dyn = Langevin(\n",
" atoms=atoms,\n",
" timestep=0.5 * units.fs,\n",
" temperature_K=300,\n",
" friction=0.01 / units.fs,\n",
")\n",
"\n",
"Please refer to the ASE documentation LINK to see how to use ASE calculators."
"traj = Trajectory('example.traj', 'w', atoms)\n",
"dyn.attach(traj.write, interval=100)\n",
"dyn.run(1000)\n",
"traj.close()"
]
},
{
Expand All @@ -43,8 +154,8 @@
"## JaxMD\n",
"\n",
"While the ASE interface is convenient and flexible, it is not meant for high performance applications.\n",
"For these purposes, apax comes with an interface to JaxMD.\n",
"JaxMD LINK is a high performance molecular dynamics engine built on top of Jax LINK.\n",
"For these purposes, apax comes with an interface to [JaxMD](https://jax-md.readthedocs.io/en/main/#).\n",
"JaxMD is a high performance molecular dynamics engine built on top of [Jax](https://jax.readthedocs.io/en/latest/index.html).\n",
"The CLI provides easy access to standard NVT and NPT simulations.\n",
"More complex simulation loops are relatively easy to build yourself in JaxMD (see their colab notebooks for examples). \n",
"Trained apax models can of course be used as `energy_fn` in such custom simulations.\n",
Expand All @@ -56,31 +167,128 @@
"metadata": {},
"source": [
"### Configuration\n",
"We can once again use the template command to give ourselves a quickstart.\n",
"\n",
"`apax template md --minimal`\n",
"We can once again use the template command to give ourselves a quickstart.\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"!apax template md"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"Open the config and specify the starting structure and simulation parameters.\n",
"If you specify the data set file itself, the first structure of the data set is going to be used as the initial structure.\n",
"Your `md_config_minimal.yaml` should look similar to this:\n",
"\n",
"```yaml\n",
"ensemble:\n",
" temperature: 300 # K\n",
" \n",
"duration: 20_000 # fs\n",
"initial_structure: md17.extxyz\n",
"```\n",
"initial_structure: project/benzene_mod.xyz\n",
"```\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"from apax.utils.helpers import mod_config\n",
"import yaml\n",
"\n",
"\n",
"config_path = Path(\"md_config.yaml\")\n",
"\n",
"config_updates = {\n",
" \"initial_structure\": str(file_path), # if the model from example 01 is used change this\n",
" \"duration\": 1000, #fs\n",
" \"ensemble\": {\n",
" \"temperature\": 300,\n",
" }\n",
"}\n",
"config_dict = mod_config(config_path, config_updates)\n",
"\n",
"with open(\"md_config.yaml\", \"w\") as conf:\n",
" yaml.dump(config_dict, conf, default_flow_style=False)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"As with training configurations, we can use the `validate` command to ensure our input is valid before we submit the calculation.\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[32mSuccess!\u001b[0m\n",
"md_config.yaml is a valid MD config.\n"
]
}
],
"source": [
"!apax validate md md_config.yaml"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Running the simulation\n",
"\n",
"The simulation can be started by running\n",
"The simulation can be started by running"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO | 12:39:38 | reading structure\n",
"INFO | 12:39:39 | Unable to initialize backend 'cuda': \n",
"INFO | 12:39:39 | Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
"INFO | 12:39:39 | Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory\n",
"INFO | 12:39:39 | initializing model\n",
"INFO | 12:39:39 | loading checkpoint from /home/linux3_i1/segreto/uni/dev/apax/examples/project/models/benzene_md/best\n",
"INFO | 12:39:39 | Initializing new trajectory file at md/md.h5\n",
"INFO | 12:39:39 | initializing simulation\n",
"INFO | 12:39:41 | running simulation for 1.0 ps\n",
"Simulation: 100%|███████████████████████████████████| 2000/2000 [00:10<00:00, 183.72it/s, T=196.3 K]\n",
"INFO | 12:39:52 | simulation finished after elapsed time: 10.93 s\n"
]
}
],
"source": [
"!apax md config.yaml md_config.yaml"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"`apax md config.yaml md_config_minimal.yaml`\n",
"\n",
"where `config.yaml` is the configuration file that was used to train the model.\n",
"\n",
Expand All @@ -98,6 +306,29 @@
"TODO"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To remove all the created files and clean up yor working directory run"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!rm -r project md config.yaml example.traj md_config.yaml"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand All @@ -107,8 +338,22 @@
}
],
"metadata": {
"kernelspec": {
"display_name": "apax",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python"
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.8"
}
},
"nbformat": 4,
Expand Down

0 comments on commit 70423bd

Please sign in to comment.