{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"%matplotlib inline"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Plot Atomic Orbitals {#plot_atomic_orbitals_example}\n\nVisualize the wave functions (orbitals) of the hydrogen atom.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Import\n\nImport the applicable libraries.\n\n::: note\n::: title\nNote\n:::\n\nThis example is modeled off of [Matplotlib: Hydrogen Wave\nFunction](http://staff.ustc.edu.cn/~zqj/posts/Hydrogen-Wavefunction/).\n\nThis example requires [sympy](https://www.sympy.org/). Install it with:\n\n``` python\npip install sympy\n```\n:::\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import numpy as np\n\nimport pyvista as pv\nfrom pyvista import examples"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Generate the Dataset\n\nGenerate the dataset by evaluating the analytic hydrogen wave function\nfrom `sympy`.\n\n$$\\begin{equation}\n\\psi_{n\\ell m}(r,\\theta,\\phi)\n=\n\\sqrt{\n\\left(\\frac{2}{na_0}\\right)^3\\, \\frac{(n-\\ell-1)!}{2n[(n+\\ell)!]}\n}\ne^{-r / na_0}\n\\left(\\frac{2r}{na_0}\\right)^\\ell\nL_{n-\\ell-1}^{2\\ell+1} \\cdot Y_\\ell^m(\\theta, \\phi)\n\\end{equation}$$\n\nSee [Hydrogen atom](https://en.wikipedia.org/wiki/Hydrogen_atom) for\nmore details.\n\nThis dataset evaluates this function for the hydrogen orbital $3d_{xy}$,\nwith the following quantum numbers:\n\n- Principal quantum number: `n=3`\n- Azimuthal quantum number: `l=2`\n- Magnetic quantum number: `m=-2`\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"grid = examples.load_hydrogen_orbital(3, 2, -2)\ngrid"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Plot the Orbital\n\nPlot the orbital using\n`add_volume() <pyvista.Plotter.add_volume>`{.interpreted-text\nrole=\"func\"} and using the default scalars contained in `grid`,\n`real_wf`. This way we can plot more than just the probability of the\nelectron, but also the phase of the electron wave function.\n\n::: note\n::: title\nNote\n:::\n\nSince the real value of evaluated wave function for this orbital varies\nbetween `[-<value>, <value>]`, we cannot use the default opacity\n`opacity='linear'`. Instead, we use `[1, 0, 1]` since we would like the\nopacity to be proportional to the absolute value of the scalars.\n:::\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"pl = pv.Plotter()\nvol = pl.add_volume(grid, cmap='magma', opacity=[1, 0, 1])\nvol.prop.interpolation_type = 'linear'\npl.camera.zoom(2)\npl.show_axes()\npl.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Plot the Orbital Contours as an Isosurface\n\nGenerate the contour plot for the orbital by determining when the\norbital equals 10% the maximum value of the orbital. This effectively\ncaptures the most likely locations of the electron for this orbital.\n\nNote how we use the absolute value of the scalars when evaluating\n`contour() <pyvista.DataSetFilters.contour>`{.interpreted-text\nrole=\"func\"} to capture where the positive and negative phases cross\n`eval_at`.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"eval_at = grid['real_wf'].max() * 0.1\ncontours = grid.contour(\n [eval_at],\n scalars=np.abs(grid['real_wf']),\n method='marching_cubes',\n)\ncontours = contours.interpolate(grid)\ncontours.plot(\n smooth_shading=True,\n show_scalar_bar=False,\n)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Volumetric Plot: Plot the Orbitals using RGBA\n\nLet\\'s now combine some of the best parts of the two above plots. The\nvolumetric plot is great for showing the probability of the \\\"electron\ncloud\\\" orbitals, but the colormap doesn\\'t quite match reality as well\nas the isosurface plot.\n\nFor this example we\\'re going to use an RGBA colormap to tightly control\nthe way the orbitals are plotted. For this, the opacity will be mapped\nto the probability of the electron being at a location in the grid,\nwhich we can do by taking the absolute value squared of the orbital\\'s\nwave function. We can set the color of the orbital based on the phase,\nwhich we can get simply with `orbital['real_wf'] < 0`.\n\nLet\\'s start with a simple one, the $3p_z$ orbital.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"def plot_orbital(orbital, cpos='iso', clip_plane='x'):\n \"\"\"Plot an electron orbital using an RGBA colormap.\"\"\"\n neg_mask = orbital['real_wf'] < 0\n rgba = np.zeros((orbital.n_points, 4), np.uint8)\n rgba[neg_mask, 0] = 255\n rgba[~neg_mask, 1] = 255\n\n # normalize opacity\n opac = np.abs(orbital['real_wf']) ** 2\n opac /= opac.max()\n rgba[:, -1] = opac * 255\n\n orbital['plot_scalars'] = rgba\n\n pl = pv.Plotter()\n vol = pl.add_volume(\n orbital,\n scalars='plot_scalars',\n )\n vol.prop.interpolation_type = 'linear'\n if clip_plane:\n pl.add_volume_clip_plane(\n vol,\n normal=clip_plane,\n normal_rotation=False,\n )\n pl.camera_position = cpos\n pl.camera.zoom(1.5)\n pl.show_axes()\n return pl.show()\n\n\nhydro_orbital = examples.load_hydrogen_orbital(3, 1, 0)\nplot_orbital(hydro_orbital, clip_plane='-x')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Volumetric Plot: $4d_{z^2}$ orbital\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"hydro_orbital = examples.load_hydrogen_orbital(4, 2, 0)\nplot_orbital(hydro_orbital, clip_plane='-y')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Volumetric Plot: $4d_{xz}$ orbital\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"hydro_orbital = examples.load_hydrogen_orbital(4, 2, -1)\nplot_orbital(hydro_orbital, clip_plane='-y')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Plot an Orbital Using a Density Plot\n\nWe can also plot atomic orbitals using a 3D density plot. For this, we\nwill use `numpy.random.choice`{.interpreted-text role=\"func\"} to sample\nall the points of our `pyvista.ImageData`{.interpreted-text\nrole=\"class\"} based on the probability of the electron being at that\ncoordinate.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"# Generate the orbital and sample based on the square of the probability of an\n# electron being within a particular volume of space.\nhydro_orbital = examples.load_hydrogen_orbital(4, 2, 0, zoom_fac=0.5)\nprob = np.abs(hydro_orbital['real_wf']) ** 2\nprob /= prob.sum()\nindices = np.random.default_rng().choice(hydro_orbital.n_points, 10000, p=prob)\n\n# add a small amount of noise to these coordinates to remove the \"grid like\"\n# structure present in the underlying ImageData\npoints = hydro_orbital.points[indices]\npoints += np.random.default_rng().random(points.shape) - 0.5\n\n# Create a point cloud and add the phase as the active scalars\npoint_cloud = pv.PolyData(points)\npoint_cloud['phase'] = hydro_orbital['real_wf'][indices] < 0\n\n# Turn the point cloud into individual spheres. We do this so we can improve\n# the plot by enabling surface space ambient occlusion (SSAO)\ndplot = point_cloud.glyph(\n geom=pv.Sphere(theta_resolution=8, phi_resolution=8),\n scale=False,\n orient=False,\n)\n\n# be sure to enable SSAO here. This makes the \"points\" that are deeper within\n# the density plot darker.\npl = pv.Plotter()\npl.add_mesh(\n dplot,\n smooth_shading=True,\n show_scalar_bar=False,\n cmap=['red', 'green'],\n ambient=0.2,\n)\npl.enable_ssao(radius=10)\npl.enable_anti_aliasing()\npl.camera.zoom(2)\npl.background_color = 'w'\npl.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Density Plot - Gaussian Points Representation\n\nFinally, let\\'s plot the same data using the \\\"Gaussian points\\\"\nrepresentation.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"point_cloud.plot(\n style='points_gaussian',\n render_points_as_spheres=False,\n point_size=3,\n emissive=True,\n background='k',\n show_scalar_bar=False,\n cpos='xz',\n zoom=2,\n)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.12.2"
}
},
"nbformat": 4,
"nbformat_minor": 0
}