test_functional.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import numpy as np\n",
    "import scipy\n",
    "from python.ADMPForce import read_mpid_inputs, setup_ewald_parameters\n",
    "from scipy.stats import special_ortho_group\n",
    "from python.utils import convert_cart2harm\n",
    "import mpidplugin\n",
    "\n",
    "mScales = np.array([0.0, 0.0, 0.0, 1.0])\n",
    "pScales = np.array([0.0, 0.0, 0.0, 1.0])\n",
    "dScales = np.array([0.0, 0.0, 0.0, 1.0])\n",
    "rc = 8 # in Angstrom\n",
    "ethresh = 1e-4\n",
    "\n",
    "pdb = 'tests/samples/waterbox_31ang.pdb'\n",
    "xml = 'tests/samples/mpidwater.xml'\n",
    "# get a random geometry for testing\n",
    "\n",
    "positions, box, list_elems, params = read_mpid_inputs(pdb, xml)\n",
    "\n",
    "scipy.random.seed(1000)\n",
    "R1 = special_ortho_group.rvs(3)\n",
    "R2 = special_ortho_group.rvs(3)\n",
    "\n",
    "positions[0:3] = positions[0:3].dot(R1)\n",
    "positions[3:6] = positions[3:6].dot(R2)\n",
    "positions[3:] += np.array([3.0, 0.0, 0.0])\n",
    "\n",
    "\n",
    "kappa, K1, K2, K3 = setup_ewald_parameters(rc, ethresh, box)\n",
    "kappa = 0.328532611\n",
    "print(kappa, K1, K2, K3)\n",
    "\n",
    "lmax = 2\n",
    "Q_lh = params['multipoles_lh']\n",
    "axis_types = params['axis_types']\n",
    "axis_indices = params['axis_indices']\n"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "0.328532611 49 49 49\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "from jax import jit, grad, value_and_grad, partial\n",
    "from python.pme import gen_pme_reciprocal\n",
    "\n",
    "pme_reciprocal_energy = value_and_grad(gen_pme_reciprocal(axis_types, axis_indices))\n",
    "\n",
    "pme_reci = jit(pme_reciprocal_energy, static_argnums=(4,5,6,7))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "ene0, f0 = pme_reci(positions, box,  Q_lh, kappa, lmax, K1, K2, K3)\n",
    "print(ene0, f0)\n",
    "\n",
    "%timeit ene0, f0 = pme_reci(positions, box,  Q_lh, kappa, lmax, K1, K2, K3)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "169.62799040544323 [[-14.17624607  -3.38865009   0.06555066]\n",
      " [  6.14597744   0.14120547  -0.13661119]\n",
      " [  7.52349044   2.85472991   0.65742572]\n",
      " ...\n",
      " [ -6.32465615   1.30861852   2.51086882]\n",
      " [  3.58574134  -0.4424106    0.53681311]\n",
      " [  2.34454043  -1.10928516  -2.76049837]]\n",
      "411 ms ± 19.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "from python.pme import pme_self\n",
    "\n",
    "pme_self_eandf = jit(value_and_grad(pme_self), static_argnums=2)\n",
    "\n",
    "eneself, fself = pme_self_eandf(Q_lh, kappa)\n",
    "\n",
    "print(eneself, fself)\n",
    "%timeit eneself, fself = pme_self_eandf(Q_lh, kappa)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "-434474.77274520777 [[ 5.46670600e+02  8.77286610e+00 -0.00000000e+00 ... -0.00000000e+00\n",
      "  -1.77069278e-02 -0.00000000e+00]\n",
      " [-2.73335300e+02 -0.00000000e+00 -0.00000000e+00 ... -0.00000000e+00\n",
      "  -0.00000000e+00 -0.00000000e+00]\n",
      " [-2.73335300e+02 -0.00000000e+00 -0.00000000e+00 ... -0.00000000e+00\n",
      "  -0.00000000e+00 -0.00000000e+00]\n",
      " ...\n",
      " [ 5.46670600e+02  8.77286610e+00 -0.00000000e+00 ... -0.00000000e+00\n",
      "  -1.77069278e-02 -0.00000000e+00]\n",
      " [-2.73335300e+02 -0.00000000e+00 -0.00000000e+00 ... -0.00000000e+00\n",
      "  -0.00000000e+00 -0.00000000e+00]\n",
      " [-2.73335300e+02 -0.00000000e+00 -0.00000000e+00 ... -0.00000000e+00\n",
      "  -0.00000000e+00 -0.00000000e+00]]\n",
      "26.6 µs ± 1.78 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "N_a = 2988\n",
    "findiff = np.empty(N_a, 3))\n",
    "delta = np.zeros((N_a,3))\n",
    "for i in range(N_a):\n",
    "    for j in range(3):\n",
    "        delta[i][j] = 1\n",
    "        findiff[i][j] = pme_reci(positions +delta, box,  Q_lh, kappa, lmax, K1, K2, K3)\n",
    "        delta[i][j] = 0\n",
    "print(findiff)"
   ],
   "outputs": [
    {
     "output_type": "error",
     "ename": "SyntaxError",
     "evalue": "unmatched ')' (76799231.py, line 2)",
     "traceback": [
      "\u001b[0;36m  File \u001b[0;32m\"/tmp/ipykernel_13802/76799231.py\"\u001b[0;36m, line \u001b[0;32m2\u001b[0m\n\u001b[0;31m    findiff = np.empty(N_a, 3))\u001b[0m\n\u001b[0m                              ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m unmatched ')'\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "print(\"Testing behaviour of the same function, under different jax wrappers\")\n",
    "print(\"====================================================================\")\n",
    "ene1, f = value_and_grad(pme_reciprocal_energy)(positions, box,  Q_lh, kappa, lmax, K1, K2, K3)\n",
    "print(\"value_and_grad:      \", ene1)\n",
    "ene2 = pme_reciprocal_energy(positions, box,  Q_lh, kappa, lmax, K1, K2, K3)\n",
    "print(\"only value:          \", ene2)\n",
    "ene3, f = jit(value_and_grad(pme_reciprocal_energy), static_argnums=(4,5,6,7))(positions, box,  Q_lh, kappa, lmax, K1, K2, K3)\n",
    "print(\"jit, value_and_grad: \", ene3)\n",
    "ene4 = jit(pme_reciprocal_energy, static_argnums=(4,5,6,7))(positions, box,  Q_lh, kappa, lmax, K1, K2, K3)\n",
    "print(\"jit, only value:     \", ene4)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Testing behaviour of the same function, under different jax wrappers\n",
      "====================================================================\n",
      "value_and_grad:       3.8504010162466096\n",
      "only value:           3.850401016246607\n",
      "jit, value_and_grad:  3.8504010162464084\n",
      "jit, only value:      3.8504010162402884\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "def f(x):\n",
    "    tmp = []\n",
    "    for i in range(x[3].astype(int)):\n",
    "        xx = x[i]\n",
    "        tmp.append(xx**2)\n",
    "    return jnp.sum(jnp.array(tmp))\n",
    "\n",
    "x = jnp.array([1.,2.,3.,4.])\n",
    "\n",
    "jitf = jit(f)\n",
    "\n",
    "print(jitf(x))"
   ],
   "outputs": [
    {
     "output_type": "error",
     "ename": "TracerIntegerConversionError",
     "evalue": "The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=0/1)>\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTracerIntegerConversionError\u001b[0m              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_6076/446721950.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0mjitf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjitf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "    \u001b[0;31m[... skipping hidden 12 frame]\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_6076/446721950.py\u001b[0m in \u001b[0;36mf\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0mtmp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mastype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m         \u001b[0mxx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m         \u001b[0mtmp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mxx\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/mpid/lib/python3.9/site-packages/jax/core.py\u001b[0m in \u001b[0;36m__index__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    479\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    480\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m__index__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 481\u001b[0;31m     \u001b[0;32mraise\u001b[0m \u001b[0mTracerIntegerConversionError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    482\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    483\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTrace\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTracerIntegerConversionError\u001b[0m: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=0/1)>\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError"
     ]
    }
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.9.7",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.9.7 64-bit ('mpid': conda)"
  },
  "interpreter": {
   "hash": "68dd20fd376d19659cce1299928e7dec6f1628c8b11d51a3a8e92f78598a03f1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}