import glob
import math
import time

import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import npy2cube  # slightly modified: addedd map to list conversions
import numpy as np
from scipy.special import genlaguerre, sph_harm

# Переход в сферические координаты
def get_r(x, y, z):
    """Радиальное расстояние от начала координат до точки"""
    return np.sqrt(x ** 2 + y ** 2 + z ** 2)


def get_theta(x, y, z):
    """Зенитный угол точки"""
    return np.arccos(z / get_r(x, y, z))


def get_phi(x, y):
    """Азимутальный угол точки"""
    return np.arctan(y / x)


def get_radial_wf(r, n, l, a0):  # Former R
    """Радиальная составляющая волновой функции, рассчитанная через полином Лагранжа"""
    return (
        (2 * r / n / a0) ** l
        * np.exp(-r / n / a0)
        * genlaguerre(n - l - 1, 2 * l + 1)(2 * r / n / a0)
    )


def wf_norm(n, l, a0):
    """Нормировочный коэффициент волновой функции"""
    return np.sqrt(
        (2 / n / a0) ** 3 * math.factorial(n - l - 1) / (2 * n * math.factorial(n + l))
    )


def wave_function(r, theta, phi, n, l, m, a0):
    """Собственно волновая функция"""
    return wf_norm(n, l, a0) * get_radial_wf(r, n, l, a0) * sph_harm(m, l, phi, theta)


def w(n, l, m, d):
    x, y, z = np.mgrid[-d:d:30j, -d:d:30j, -d:d:30j]
    # Боровский радиус, наиболее вероятное расстояние между ядром и электроном
    a0 = 1
    result = wave_function(
        get_r(x, y, z), get_theta(x, y, z), get_phi(x, y), n, l, m, a0
    )
    # Изменена с квадрата модуля на это, иначе орбитали не получаются похожими на реальность совсем
    return np.real(result) + np.imag(result)

Рассчет плотностей

d = 20
step = 0.5
for n in range(1, 4):
    
    for l in range(0, n):
        for m in range(-l, l + 1):
            grid = w(n, l, m, d)
            npy2cube.npy2cube(
                grid, (-d, -d, -d), (step, step, step), f"{n}_{l}_{m}.cube"
            )

Визуализация в Pymol

import __main__

__main__.pymol_argv = ["pymol", "-x"]

import pymol

pymol.finish_launching()
from pymol import cmd, stored

cmd.reinitialize()
# Палитра
cmd.volume_ramp_new(
    "ramp007",
    [
        -0.015,
        1.00,
        0.00,
        0.00,
        0.00,
        -0.01,
        0.60,
        2.00,
        1.00,
        0.20,
        -0.005,
        0.10,
        0.50,
        1.00,
        0.00,
        0.005,
        0.50,
        0.20,
        0.50,
        0.00,
        0.01,
        0.30,
        1.00,
        1.00,
        0.50,
        0.015,
        1.00,
        1.00,
        2.00,
        0.00,
    ],
)

# Картинки плотностей
for n in range(1, 4):
    for l in range(0, n):
        for m in range(-l, l + 1):
            cmd.load(f"{n}_{l}_{m}.cube")
            cmd.volume(f"{n}_{l}_{m}_volume", f"{n}_{l}_{m}", ramp="ramp007")
            cmd.hide()
            cmd.show("volume", f"{n}_{l}_{m}_volume")
            
            # Slab view
            cmd.reset()
            cmd.turn("x", 90)
            cmd.clip("slab", 0)
            time.sleep(1)
            cmd.png(f"{n}_{l}_{m}.png")
            
            # Full volume
            cmd.reset()
            cmd.turn("x", 45)
            cmd.turn("y", 45)
            cmd.turn("z", 45)
            time.sleep(1)
            cmd.png(f"full_{n}_{l}_{m}.png")

Срезы орбиталей

orbs = {0: "s", 1: "p", 2: "d"}

fig, ax = plt.subplots(7, 2, figsize=(15, 28))

i, j = 0, 0

for n in range(1, 4):
    for l in range(0, n):
        for m in range(-l, l + 1):
            f = plt.imread(f"{n}_{l}_{m}.png")
            ax[i, j].axis("off")
            ax[i, j].imshow(f)
            ax[i, j].set_title(f"{n}{orbs[l]}", {"fontsize": 18})
            j += 1

            if j == 2:
                i += 1
                j = 0
plt.tight_layout()

Внешний вид орбиталей

orbs = {0: "s", 1: "p", 2: "d"}

fig, ax = plt.subplots(7, 2, figsize=(15, 28))

i, j = 0, 0

for n in range(1, 4):
    for l in range(0, n):
        for m in range(-l, l + 1):
            f = plt.imread(f"full_{n}_{l}_{m}.png")
            ax[i, j].axis("off")
            ax[i, j].imshow(f)
            ax[i, j].set_title(f"{n}{orbs[l]}", {"fontsize": 18})
            j += 1

            if j == 2:
                i += 1
                j = 0
plt.tight_layout()

ORCA

cmd.reinitialize()

cmd.volume_ramp_new(
    "ramp007",
    [
        -0.015,
        1.00,
        0.00,
        0.00,
        0.00,
        -0.01,
        0.60,
        2.00,
        1.00,
        0.20,
        -0.005,
        0.10,
        0.50,
        1.00,
        0.00,
        0.005,
        0.50,
        0.20,
        0.50,
        0.00,
        0.01,
        0.30,
        1.00,
        1.00,
        0.50,
        0.015,
        1.00,
        1.00,
        2.00,
        0.00,
    ],
)

# Картинки плотностей
for i in range(14):
    cmd.load(f"H-{i}.cube")
    cmd.volume(f"H-{i}_volume", f"H-{i}", ramp="ramp007")
    cmd.hide()
    cmd.show("volume", f"H-{i}_volume")

    # Slab view
    cmd.reset()
    cmd.turn("x", 90)
    cmd.clip("slab", 0)
    time.sleep(4)
    cmd.png(f"H-{i}.png")

    # Full volume
    cmd.reset()
    cmd.turn("x", 45)
    cmd.turn("y", 45)
    cmd.turn("z", 45)
    time.sleep(4)
    cmd.png(f"full_H-{i}.png")

Отображение орбиталей из ORCA

Срезы

orbs = [
    "1s",
    "2s",
    "2p",
    "2p",
    "2p",
    "3s",
    "3d",
    "3d",
    "3d",
    "3d",
    "3d",
    "3p",
    "3p",
    "3p",
]

fig, ax = plt.subplots(7, 2, figsize=(15, 28))

i, j = 0, 0

for n in range(14):
    f = plt.imread(f"H-{n}.png")
    ax[i, j].axis("off")
    ax[i, j].imshow(f)
    ax[i, j].set_title(f"{orbs[n]}", {"fontsize": 18})
    j += 1

    if j == 2:
        i += 1
        j = 0
plt.tight_layout()

Внешний вид

fig, ax = plt.subplots(7, 2, figsize=(15, 28))

i, j = 0, 0

for n in range(14):
    f = plt.imread(f"full_H-{n}.png")
    ax[i, j].axis("off")
    ax[i, j].imshow(f)
    ax[i, j].set_title(f"{orbs[n]}", {"fontsize": 18})
    j += 1

    if j == 2:
        i += 1
        j = 0
plt.tight_layout()