diff --git a/notebooks/nanoDFT-demo.ipynb b/notebooks/nanoDFT-demo.ipynb index aed87eb..4b0db13 100644 --- a/notebooks/nanoDFT-demo.ipynb +++ b/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, diff --git a/pyscf_ipu/experimental/interop.py b/pyscf_ipu/experimental/interop.py new file mode 100644 index 0000000..0b57574 --- /dev/null +++ b/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 diff --git a/pyscf_ipu/experimental/mesh.py b/pyscf_ipu/experimental/mesh.py index b75956f..aa2c917 100644 --- a/pyscf_ipu/experimental/mesh.py +++ b/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 diff --git a/pyscf_ipu/experimental/plot.py b/pyscf_ipu/experimental/plot.py new file mode 100644 index 0000000..adb1cd6 --- /dev/null +++ b/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]}, + ] + } diff --git a/pyscf_ipu/experimental/structure.py b/pyscf_ipu/experimental/structure.py index c1ba34d..7977151 100644 --- a/pyscf_ipu/experimental/structure.py +++ b/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() diff --git a/pyscf_ipu/experimental/types.py b/pyscf_ipu/experimental/types.py index e7fee8f..63077ca 100644 --- a/pyscf_ipu/experimental/types.py +++ b/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] diff --git a/test/test_experimental.py b/test/test_experimental.py index 26fbbe0..5571f6d 100644 --- a/test/test_experimental.py +++ b/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)