adding tools for plotting with py3dmol

notebooks/nanoDFT-demo.ipynb

@@ -189,7 +189,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "id": "efa3bc23",
    "metadata": {},
    "outputs": [],
@@ -219,7 +219,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 4,
    "id": "25acb1b7",
    "metadata": {},
    "outputs": [
@@ -254,7 +254,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 16,
    "id": "fc79c21d",
    "metadata": {},
    "outputs": [
@@ -282,67 +282,14 @@
     "molecular_orbitals.shape"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "85db5266",
-   "metadata": {},
-   "source": [
-    "Below we define functions for visualising the molecular orbitals of benzene using the [py3Dmol](https://3dmol.org/) package"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "id": "63fbdbd2",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def cube_data(axes, value) -> str:\n",
-    "    fmt = \"cube format\\n\\n\"\n",
-    "    x, y, z = axes\n",
-    "    nx, ny, nz = [ax.shape[0] for ax in axes]\n",
-    "    fmt += \"0 \" + \" \".join([f\"{v:12.6f}\" for v in [x[0], y[0], z[0]]]) + \"\\n\"\n",
-    "    fmt += f\"{nx} \" + \" \".join([f\"{v:12.6f}\" for v in [x[1] - x[0], 0.0, 0.0]]) + \"\\n\"\n",
-    "    fmt += f\"{ny} \" + \" \".join([f\"{v:12.6f}\" for v in [0.0, y[1] - y[0], 0.0]]) + \"\\n\"\n",
-    "    fmt += f\"{nz} \" + \" \".join([f\"{v:12.6f}\" for v in [0.0, 0.0, z[1] - z[0]]]) + \"\\n\"\n",
-    "\n",
-    "    line = \"\"\n",
-    "    for i in range(len(value)):\n",
-    "        line += f\"{value[i]:12.6f}\"\n",
-    "\n",
-    "        if i % 6 == 0:\n",
-    "            fmt += line + \"\\n\"\n",
-    "            line = \"\"\n",
-    "\n",
-    "    return fmt\n",
-    "\n",
-    "\n",
-    "def build_transferfn(value) -> dict:\n",
-    "    v = np.percentile(value, [99.9, 75])\n",
-    "    a = [0.02, 0.0005]\n",
-    "    return {\n",
-    "        \"transferfn\": [\n",
-    "            {\"color\": \"blue\", \"opacity\": a[0], \"value\": -v[0]},\n",
-    "            {\"color\": \"blue\", \"opacity\": a[1], \"value\": -v[1]},\n",
-    "            {\"color\": \"white\", \"opacity\": 0.0, \"value\": 0.0},\n",
-    "            {\"color\": \"red\", \"opacity\": a[1], \"value\": v[1]},\n",
-    "            {\"color\": \"red\", \"opacity\": a[0], \"value\": v[0]},\n",
-    "        ]\n",
-    "    }\n",
-    "\n",
-    "\n",
-    "def plot_orbital(orbital, mol):\n",
-    "    xyzfmt = f\"{len(mol.atom)}\\n\\n\" + mol.tostring()\n",
-    "    v = py3Dmol.view(data=xyzfmt, style={\"stick\": {\"radius\": 0.06}, \"sphere\": {\"radius\": 0.2}})\n",
-    "    v.addVolumetricData(cube_data(axes, orbital), \"cube\", build_transferfn(orbital))\n",
-    "    return v"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "8efd393a",
    "metadata": {},
    "source": [
+    "Below we use `plot_volume` from the `pyscf_ipu.experimental.plot` module to \n",
+    "visualise the molecular orbitals of benzene using the [py3Dmol](https://3dmol.org/) package.\n",
+    "\n",
     "Try changing the `mo_index` variable to select the different molecular orbitals benzene."
    ]
   },
@@ -353,10 +300,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from pyscf_ipu.experimental.plot import plot_volume\n",
+    "from pyscf_ipu.experimental.interop import from_pyscf\n",
     "mo_index = 5\n",
     "\n",
     "orbital = molecular_orbitals[:, mo_index]\n",
-    "mol_view = plot_orbital(orbital, mol)\n",
+    "structure, _ = from_pyscf(mol)\n",
+    "mol_view = plot_volume(structure, orbital, axes)\n",
     "mol_view.spin()"
    ]
   }
@@ -377,7 +327,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.0"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

pyscf_ipu/experimental/interop.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+from typing import Tuple
+
+import numpy as np
+from periodictable import elements
+from pyscf import gto
+
+from .basis import Basis, basisset
+from .structure import Structure
+
+
+def to_pyscf(
+    structure: Structure, basis_name: str = "sto-3g", unit: str = "Bohr"
+) -> "gto.Mole":
+    mol = gto.Mole(unit=unit, spin=structure.num_electrons % 2, cart=True)
+    mol.atom = [
+        (symbol, pos)
+        for symbol, pos in zip(structure.atomic_symbol, structure.position)
+    ]
+    mol.basis = basis_name
+    mol.build(unit=unit)
+    return mol
+
+
+def from_pyscf(mol: "gto.Mole") -> Tuple[Structure, Basis]:
+    atomic_number = []
+    position = []
+
+    for i in range(mol.natm):
+        sym, pos = mol.atom[i]
+        atomic_number.append(elements.symbol(sym).number)
+        position.append(pos)
+
+    structure = Structure(
+        atomic_number=np.array(atomic_number),
+        position=np.array(position),
+        is_bohr=mol.unit != "Angstom",
+    )
+
+    basis = basisset(structure, basis_name=mol.basis)
+
+    return structure, basis

pyscf_ipu/experimental/mesh.py

@@ -28,13 +28,20 @@ def uniform_mesh(
     axes = [jnp.linspace(-bi, bi, ni) for bi, ni in zip(b, n)]
     mesh = jnp.stack(jnp.meshgrid(*axes, indexing="ij"), axis=-1)
     mesh = mesh.reshape(-1, ndim)
-    return mesh
+    return mesh, axes
 
 
 def electron_density(
     basis: Basis, mesh: FloatNx3, C: Optional[FloatNxN] = None
 ) -> FloatN:
-    C = jnp.eye(basis.num_orbitals) if C is None else C
-    orbitals = basis(mesh) @ C
+    orbitals = molecular_orbitals(basis, mesh, C)
     density = jnp.sum(basis.occupancy * orbitals * orbitals, axis=-1)
     return density
+
+
+def molecular_orbitals(
+    basis: Basis, mesh: FloatNx3, C: Optional[FloatNxN] = None
+) -> FloatN:
+    C = jnp.eye(basis.num_orbitals) if C is None else C
+    orbitals = basis(mesh) @ C
+    return orbitals

pyscf_ipu/experimental/plot.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+import numpy as np
+from numpy.typing import NDArray
+
+from .structure import Structure
+from .types import MeshAxes
+from .units import to_angstrom
+
+
+def plot_volume(structure: Structure, value: NDArray, axes: MeshAxes):
+    """plots volumetric data value with molecular structure.
+
+    Args:
+        structure (Structure): molecular structure
+        value (NDArray): the volume data to render
+        axes (MeshAxes): the axes over which the data was sampled.
+
+    Returns:
+        py3DMol View object
+    """
+    v = structure.view()
+    v.addVolumetricData(cube_data(value, axes), "cube", build_transferfn(value))
+    return v
+
+
+def cube_data(value: NDArray, axes: MeshAxes) -> str:
+    """Generate the cube file format as a string.  See:
+
+      https://paulbourke.net/dataformats/cube/
+
+    Args:
+        value (NDArray): the volume data to serialise in the cube format
+        axes (MeshAxes): the axes over which the data was sampled
+
+    Returns:
+        str: cube format representation of the volumetric data.
+    """
+    axes = [to_angstrom(ax) for ax in axes]
+    fmt = "cube format\n\n"
+    x, y, z = axes
+    nx, ny, nz = [ax.shape[0] for ax in axes]
+    fmt += "0 " + " ".join([f"{v:12.6f}" for v in [x[0], y[0], z[0]]]) + "\n"
+    fmt += f"{nx} " + " ".join([f"{v:12.6f}" for v in [x[1] - x[0], 0.0, 0.0]]) + "\n"
+    fmt += f"{ny} " + " ".join([f"{v:12.6f}" for v in [0.0, y[1] - y[0], 0.0]]) + "\n"
+    fmt += f"{nz} " + " ".join([f"{v:12.6f}" for v in [0.0, 0.0, z[1] - z[0]]]) + "\n"
+
+    line = ""
+    for i in range(len(value)):
+        line += f"{value[i]:12.6f}"
+
+        if i % 6 == 0:
+            fmt += line + "\n"
+            line = ""
+
+    return fmt
+
+
+def build_transferfn(value: NDArray) -> dict:
+    """Generate the 3dmol.js transferfn argument for a particular value.
+
+    Tries to set isovalues to capture main features of the volume data.
+
+    Args:
+        value (NDArray): the volume data.
+
+    Returns:
+        dict: containing transferfn
+    """
+    v = np.percentile(value, [99.9, 75])
+    a = [0.02, 0.0005]
+    return {
+        "transferfn": [
+            {"color": "blue", "opacity": a[0], "value": -v[0]},
+            {"color": "blue", "opacity": a[1], "value": -v[1]},
+            {"color": "white", "opacity": 0.0, "value": 0.0},
+            {"color": "red", "opacity": a[1], "value": v[1]},
+            {"color": "red", "opacity": a[0], "value": v[0]},
+        ]
+    }

pyscf_ipu/experimental/structure.py

@@ -5,7 +5,6 @@
 import numpy as np
 from periodictable import elements
 from py3Dmol import view
-from pyscf import gto
 
 from .types import FloatNx3, IntN
 from .units import to_angstrom, to_bohr
@@ -48,19 +47,6 @@ def view(self) -> "view":
         return view(data=self.to_xyz(), style={"stick": {"radius": 0.06}})
 
 
-def to_pyscf(
-    structure: Structure, basis_name: str = "sto-3g", unit: str = "Bohr"
-) -> "gto.Mole":
-    mol = gto.Mole(unit=unit, spin=structure.num_electrons % 2, cart=True)
-    mol.atom = [
-        (symbol, pos)
-        for symbol, pos in zip(structure.atomic_symbol, structure.position)
-    ]
-    mol.basis = basis_name
-    mol.build(unit=unit)
-    return mol
-
-
 def molecule(name: str):
     name = name.lower()
 

pyscf_ipu/experimental/types.py

@@ -1,4 +1,6 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+from typing import Tuple
+
 from jaxtyping import Array, Float, Int
 
 Float3 = Float[Array, "3"]
@@ -8,3 +10,5 @@
 FloatNxM = Float[Array, "N M"]
 Int3 = Int[Array, "3"]
 IntN = Int[Array, "N"]
+
+MeshAxes = Tuple[FloatN, FloatN, FloatN]

test/test_experimental.py

@@ -18,9 +18,10 @@
     overlap_basis,
     overlap_primitives,
 )
+from pyscf_ipu.experimental.interop import to_pyscf
 from pyscf_ipu.experimental.mesh import electron_density, uniform_mesh
 from pyscf_ipu.experimental.primitive import Primitive
-from pyscf_ipu.experimental.structure import molecule, to_pyscf
+from pyscf_ipu.experimental.structure import molecule
 
 
 @pytest.mark.parametrize("basis_name", ["sto-3g", "6-31g**"])
@@ -38,7 +39,7 @@ def test_gto(basis_name):
     # Atomic orbitals
     structure = molecule("water")
     basis = basisset(structure, basis_name)
-    mesh = uniform_mesh()
+    mesh, _ = uniform_mesh()
     actual = basis(mesh)
 
     mol = to_pyscf(structure, basis_name)