config_wise_reg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting Regularisation for virials and forces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ase, ase.io\n",
    "import numpy as np\n",
    "import os\n",
    "os.chdir('example_data/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read in some data that has atomic forces, virial stresses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "500\n"
     ]
    }
   ],
   "source": [
    "atoms = ase.io.read('hcp_Mg_geomopti_randshear_pm_0.01_product_symm_k0p012.extxyz', index=':') ; print(len(atoms))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set the regularisation inside the Atoms object\n",
    "Regularisation strength for the forces (virials) is controlled by regularise (regularise_v), and if the modulated regularisation is less than a specified tolerence it is set as the tolerence value (this makes sure reg. is not set to 0, as this may cause instabilities in ML fitting). \n",
    "\n",
    "Configwise = True sets the force reg. to be configuration wise (i.e. the reg. is informed by the maximum force experience across all the atoms in a given configuration). False = the reg. is set for each atom individually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "regularise=0.01\n",
    "regularise_v = 0.01\n",
    "configwise=False\n",
    "tol = 1e-3\n",
    "tol_v=1e-3\n",
    "for i, a in enumerate(atoms):\n",
    "    force_arr = a.get_forces()\n",
    "    force_max = max(np.linalg.norm(force_arr, axis=1))\n",
    "    if configwise:\n",
    "        if force_max < tol:\n",
    "            atoms[i].info['force_sigma'] = tol\n",
    "        else:\n",
    "            atoms[i].info['force_sigma'] = regularise*force_max \n",
    "    else:\n",
    "        force_component_sigma = [tol if j < tol else j for j in regularise*abs(force_arr).flatten()]\n",
    "        force_component_sigma = np.array(force_component_sigma)\n",
    "        force_component_sigma = np.reshape(force_component_sigma, [len(a), 3])\n",
    "        atoms[i].arrays[\"force_component_sigma\"] = force_component_sigma\n",
    "\n",
    "        \n",
    "for i, a in enumerate(atoms):\n",
    "    virials = a.info['virial']\n",
    "    virial_component_sigma = regularise_v*abs(virials)\n",
    "    x = [tol_v if j < tol_v else j for j in virial_component_sigma.flatten()]\n",
    "    atoms[i].info['virial_component_sigma'] = np.array(x).flatten()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How to write to file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/ase/io/extxyz.py:997: UserWarning: write_xyz() overwriting array \"forces\" present in atoms.arrays with stored results from calculator\n",
      "  warnings.warn('write_xyz() overwriting array \"{0}\" present '\n"
     ]
    }
   ],
   "source": [
    "Properties='species:S:1:pos:R:3:force:R:3:force_component_sigma:R:3:'\n",
    "ase.io.extxyz.write_extxyz('now_with_reg.extxyz', atoms, columns=['symbols', 'positions', 'forces', 'force_component_sigma'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}