diff --git a/docs/batching.ipynb b/docs/batching.ipynb index a013dbf..64a4089 100644 --- a/docs/batching.ipynb +++ b/docs/batching.ipynb @@ -35,7 +35,28 @@ "the batch we build a `Hamiltonian` for each bond length and stack the built modules to \n", "create a batched-Hamiltonian. This example uses the `sto-3g` basis set and the simple\n", "local density approximation of density functional theory but the this formulation isn't\n", - "unique to these choices for how the Hamiltonian is represented." + "unique to these choices for how the Hamiltonian is represented.\n", + "\n", + ":::{note}\n", + "The following code cell will install MESS into the Google Colab runtime.\n", + "Select the 🚀 in the toolbar above to try this out!\n", + ":::" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-cell" + ] + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if 'google.colab' in sys.modules:\n", + " !pip install git+https://github.com/graphcore-research/mess.git\n" ] }, { diff --git a/docs/optim.ipynb b/docs/optim.ipynb index 2c29636..bcd4086 100644 --- a/docs/optim.ipynb +++ b/docs/optim.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": 19, "metadata": { "tags": [ "remove-input" @@ -29,35 +29,50 @@ "On the second point, there are many possible approaches to solving this constrained \n", "optimisation problem. In the following we setup minimising the total energy with the \n", "[Adam optimiser](https://optax.readthedocs.io/en/latest/api/optimizers.html#adam) \n", - "from the [optax library](https://optax.readthedocs.io/en/latest/index.html)." + "from the [optax library](https://optax.readthedocs.io/en/latest/index.html).\n", + "\n", + "\n", + ":::{note}\n", + "The following code cell will install MESS into the Google Colab runtime.\n", + "Select the 🚀 in the toolbar above to try this out!\n", + ":::" ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 1, "metadata": { "tags": [ "hide-cell" ] }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/home/ubuntu/miniforge3/envs/jax/lib/python3.10/site-packages/pyscf/dft/libxc.py:771: UserWarning: Since PySCF-2.3, B3LYP (and B3P86) are changed to the VWN-RPA variant, corresponding to the original definition by Stephens et al. (issue 1480) and the same as the B3LYP functional in Gaussian. To restore the VWN5 definition, you can put the setting \"B3LYP_WITH_VWN5 = True\" in pyscf_conf.py\n", - " warnings.warn('Since PySCF-2.3, B3LYP (and B3P86) are changed to the VWN-RPA variant, '\n" - ] - } - ], + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + " !pip install git+https://github.com/graphcore-research/mess.git" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": { + "tags": [ + "hide-cell" + ] + }, + "outputs": [], "source": [ + "\n", + "\n", "import jax\n", "import jax.numpy as jnp\n", "import optax\n", "import seaborn as sns\n", "from tqdm.notebook import tqdm\n", "\n", - "from mess import Hamiltonian, basisset, molecule\n", + "from mess import Hamiltonian, basisset\n", "from mess.structure import nuclear_energy\n", "\n", "sns.set_theme(style=\"whitegrid\")" @@ -65,13 +80,15 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 21, "metadata": {}, "outputs": [], "source": [ - "mol = molecule(\"water\")\n", + "from mess.interop import from_pyquante\n", + "\n", + "mol = from_pyquante(\"c6h6\")\n", "basis = basisset(mol, \"6-31g\")\n", - "H = Hamiltonian(basis, xc_method=\"lda\")\n", + "H = Hamiltonian(basis, xc_method=\"pbe\")\n", "optimiser = optax.adam(learning_rate=0.1)" ] }, @@ -86,7 +103,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 22, "metadata": {}, "outputs": [], "source": [ @@ -94,6 +111,7 @@ "\n", "\n", "@jax.jit\n", + "@jax.value_and_grad\n", "def total_energy(Z):\n", " C = H.orthonormalise(Z)\n", " P = basis.density_matrix(C)\n", @@ -113,18 +131,18 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 23, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "8a43134ddb704605bd039d816026672b", + "model_id": "e144922bde6449c0bccd36030c6d8e08", "version_major": 2, "version_minor": 0 }, "text/plain": [ - " 0%| | 0/128 [00:00" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "history = jnp.stack(history)\n", + "history = history - jnp.min(history)\n", + "ax = sns.lineplot(history)\n", + "ax.set_yscale(\"log\")\n", + "ax.set_xlabel(\"Iteration\")\n", + "ax.set_ylabel(\"Total energy (Hartree)\");" + ] + }, { "cell_type": "code", "execution_count": 6, @@ -173,6 +256,180 @@ "ax.set_xlabel(\"Iteration\")\n", "ax.set_ylabel(\"Total energy (Hartree)\");" ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'cart'" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "backend = \"pyscf_cart\"\n", + "backend.startswith(\"pyscf_\")\n", + "backend.split(\"_\")[1]" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/ubuntu/miniforge3/envs/jax/lib/python3.10/site-packages/pyscf/dft/libxc.py:771: UserWarning: Since PySCF-2.3, B3LYP (and B3P86) are changed to the VWN-RPA variant, corresponding to the original definition by Stephens et al. (issue 1480) and the same as the B3LYP functional in Gaussian. To restore the VWN5 definition, you can put the setting \"B3LYP_WITH_VWN5 = True\" in pyscf_conf.py\n", + " warnings.warn('Since PySCF-2.3, B3LYP (and B3P86) are changed to the VWN-RPA variant, '\n" + ] + } + ], + "source": [ + "from mess import molecule\n", + "from mess.interop import to_pyscf\n", + "import numpy as np\n", + "import numpy.linalg as npl\n", + "\n", + "np.set_printoptions(linewidth=120, precision=6)\n", + "\n", + "mol = molecule(\"water\")\n", + "scfmol = to_pyscf(mol, basis_name=\"def2-TZVPPD\")\n", + "# scfmol.cart = False\n", + "# scfmol = scfmol.build()\n", + "S = scfmol.intor(\"int1e_ovlp_cart\")\n", + "s, U = npl.eigh(S)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([2.227773e-04, 4.519461e-04, 1.138834e-03, 5.175996e-04, 3.437612e-03, 8.825716e-04, 1.000282e-03, 4.020117e-03,\n", + " 3.256600e-03, 5.913186e-03, 5.233487e-03, 6.822549e-03, 6.770662e-03, 1.884973e-03, 1.011013e-02, 8.236594e-04,\n", + " 1.085325e-02, 8.021005e-04, 1.065478e-02, 7.274325e-03, 7.035361e-04, 2.655205e-02, 1.293354e-02, 1.585945e-02,\n", + " 6.655871e-03, 1.770198e-02, 3.971056e-02, 9.050910e-04, 5.530916e-03, 6.185895e-03, 1.594431e-02, 1.441710e-02,\n", + " 1.686322e-02, 2.529782e-02, 7.736134e-04, 1.433575e-02, 4.884687e-03, 5.009570e-02, 2.044539e-02, 2.933490e-02,\n", + " 4.542135e-02, 6.228932e-02, 1.673826e-02, 1.258803e-02, 1.978192e-02, 6.286715e-02, 2.073157e-02, 6.370569e-02,\n", + " 1.735051e-02, 5.812182e-03, 1.039793e-01, 3.430022e-02, 5.516339e-02, 3.182214e-02, 3.353706e-03, 1.008233e-01,\n", + " 4.393568e-02, 2.208200e-02, 4.042325e-02, 1.120331e-01, 2.074160e-02, 4.688071e-02, 2.036013e-01, 7.769408e-02,\n", + " 5.559280e-02, 5.492349e-02, 8.038952e-02, 4.374833e-01, 1.038853e-01, 2.356155e-02, 1.074352e-01, 5.213506e-02,\n", + " 4.636286e-01, 1.819283e-01, 2.744310e-01, 6.971762e-01, 1.009265e+00, 5.941142e-01, 5.952234e-01, 1.692097e+00,\n", + " 7.889425e+00])" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "np.diff(s)" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "70" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "S = scfmol.intor(\"int1e_ovlp_sph\")\n", + "s, U = npl.eigh(S)\n", + "np.sum(np.diff(s) > 1e-3)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(82, 82)" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "U.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "True" + ] + }, + "execution_count": 21, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sc = s[s > 1e-3]\n", + "Uc = U[:, s > 1e-3]\n", + "\n", + "X = Uc @ np.diag(np.power(sc, -0.5)) @ Uc.T\n", + "\n", + "sc = np.where(s < 1e-3, 1.0, s)\n", + "Uc = np.where(s < 1e-3, 0.0, U)\n", + "\n", + "X2 = Uc @ np.diag(np.power(sc, -0.5)) @ Uc.T\n", + "\n", + "np.allclose(X, X2)" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(82, 82)" + ] + }, + "execution_count": 23, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "(Uc @ np.diag(np.power(sc, -0.5))).shape" + ] } ], "metadata": { diff --git a/docs/tour.ipynb b/docs/tour.ipynb index 0ee3d09..2590b2f 100644 --- a/docs/tour.ipynb +++ b/docs/tour.ipynb @@ -30,7 +30,7 @@ "object will display a 3D visualisation\n", "\n", ":::{note}\n", - "The following code cell will install mess into the Google Colab runtime.\n", + "The following code cell will install MESS into the Google Colab runtime.\n", "Select the 🚀 in the toolbar above to try this out!\n", ":::" ]