get_LE_ps.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from scipy.special import spherical_jn as j_l\n",
    "from scipy.special import spherical_in as i_l\n",
    "from spherical_bessel_zeros import Jn_zeros\n",
    "from scipy.integrate import quadrature\n",
    "\n",
    "from dataset_processing import get_dataset_slice\n",
    "from LE_ps import get_LE_expansion, get_LE_ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = 6.0  # Radius of the sphere\n",
    "E_max_2 = 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5 11\n"
     ]
    }
   ],
   "source": [
    "l_big = 26\n",
    "n_big = 26\n",
    "\n",
    "z_ln = Jn_zeros(l_big, n_big)  # Spherical Bessel zeros\n",
    "z_nl = z_ln.T\n",
    "\n",
    "E_nl = z_nl**2\n",
    "E_max = E_max_2 - E_nl[0, 0]\n",
    "n_max = np.where(E_nl[:, 0] <= E_max)[0][-1] + 1\n",
    "l_max = np.where(E_nl[0, :] <= E_max)[0][-1]\n",
    "print(n_max, l_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def R_nl(n, l, r):\n",
    "    return j_l(l, z_nl[n, l]*r/a)\n",
    "\n",
    "def N_nl(n, l):\n",
    "    # Normalization factor for LE basis functions\n",
    "    def function_to_integrate_to_get_normalization_factor(x):\n",
    "        return j_l(l, x)**2 * x**2\n",
    "    integral, _ = sp.integrate.quadrature(function_to_integrate_to_get_normalization_factor, 0.0, z_nl[n, l])\n",
    "    return (1.0/z_nl[n, l]**3 * integral)**(-0.5)\n",
    "\n",
    "def get_LE_function_python(n, l, r):\n",
    "    R = np.zeros_like(r)\n",
    "    for i in range(r.shape[0]):\n",
    "        R[i] = R_nl(n, l, r[i])\n",
    "    return N_nl(n, l)*R*a**(-1.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feed LE (delta) radial spline points to Rust calculator:\n",
    "\n",
    "n_spline_points = 101\n",
    "spline_x = np.linspace(0.0, a, n_spline_points)  # x values\n",
    "\n",
    "spline_f = []\n",
    "for l in range(l_max+1):\n",
    "    for n in range(n_max):\n",
    "        spline_f_single = get_LE_function_python(n, l, spline_x)\n",
    "        spline_f.append(spline_f_single)\n",
    "spline_f = np.array(spline_f).T\n",
    "spline_f = spline_f.reshape(n_spline_points, l_max+1, n_max)  # f(x) values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_LE_function_derivative(n, l, r):\n",
    "    delta = 1e-6\n",
    "    all_derivatives_except_at_zero = (get_LE_function_python(n, l, r[1:]+delta) - get_LE_function_python(n, l, r[1:]-delta)) / (2.0*delta)\n",
    "    derivative_at_zero = (get_LE_function_python(n, l, np.array([delta/10.0])) - get_LE_function_python(n, l, np.array([0.0]))) / (delta/10.0)\n",
    "    return np.concatenate([derivative_at_zero, all_derivatives_except_at_zero])\n",
    "\n",
    "spline_df = []\n",
    "for l in range(l_max+1):\n",
    "    for n in range(n_max):\n",
    "        spline_df_single = get_LE_function_derivative(n, l, spline_x)\n",
    "        spline_df.append(spline_df_single)\n",
    "spline_df = np.array(spline_df).T\n",
    "spline_df = spline_df.reshape(n_spline_points, l_max+1, n_max)  # df/dx values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"splines.txt\", \"w\") as file:\n",
    "    np.savetxt(file, spline_x.flatten(), newline=\" \")\n",
    "    file.write(\"\\n\")\n",
    "\n",
    "with open(\"splines.txt\", \"a\") as file:\n",
    "    np.savetxt(file, (1.0/(4.0*np.pi))*spline_f.flatten(), newline=\" \")\n",
    "    file.write(\"\\n\")\n",
    "    np.savetxt(file, (1.0/(4.0*np.pi))*spline_df.flatten(), newline=\" \")\n",
    "    file.write(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading dataset\n",
      "Shuffling and extracting from dataset (length: 10000)\n",
      "Shuffling and extraction done\n",
      "TensorMap with 1 blocks\n",
      "keys: ['_']\n",
      "        0\n"
     ]
    }
   ],
   "source": [
    "train_slice = \"0:100\"\n",
    "train_structures = get_dataset_slice(\"random-ch4-10k.extxyz\", train_slice)\n",
    "\n",
    "ps = get_LE_ps(train_structures, \"splines.txt\", E_nl, E_max_2, a)\n",
    "print(ps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.7 ('rascaline-branch': venv)",
   "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.7.7"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "df1ed57d087713ef71c7e94baf1db773f5ad4a8102815aa9b5123e587036041f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}