import numpy
import scipy.special
import scipy.misc
import npy2cube
import math
import time
import glob
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import os

Задание функций

Прежде всего, зададим волновую функцию

def w(n,l,m,d): #На вход она принимает квантовые числа n, l и m
    x,y,z = np.mgrid[-d:d:30j,-d:d:30j,-d:d:30j]
    
# Делаем переход к сферическим координатам
    r = lambda x,y,z: np.sqrt(x**2+y**2+z**2)
    theta = lambda x,y,z: np.arccos(z/r(x,y,z))
    phi = lambda x,y,z: np.arctan2(y, x) + math.pi

# Задаем боровский радиус - наиболее вероятная дистанция между ядром и электроном
    a0 = 1
    
# R – радиальная часть волновой функции, n, l – квантовые числа, genlaguerre - подсчет полинома Лагерра
    R = lambda r,n,l: (2*r/n/a0)**l * np.exp(-r/n/a0) * scipy.special.genlaguerre(n-l-1,2*l+1)(2*r/n/a0)
    
# Нормировочный коэффициент
    norm = lambda n,l: np.sqrt((2/n/a0)**3 * math.factorial(n-l-1)/(2*n*math.factorial(n+l)))
    
# Наша волновая функция, scipy.special.sph_harm – сферическая гармоника
    WF = lambda r,theta,phi,n,l,m: norm(n, l) * R(r,n,l) * scipy.special.sph_harm(m,l,phi,theta)
    
# Итоговая плотность вероятности WF* 
    result = WF(r(x,y,z),theta(x,y,z),phi(x,y,z),n,l,m) 
    
    return np.real(result) + np.imag(result)

* Если выводить непосредственно квадрат модуля, то полученные электронные орбитали будут не вполне верными. По итогу, было решено выводить сумму действительной и мнимой частей WF (помогли консультации с коллегами).

Воспользуемся функцией от Андрея Демкива, немного ее изменив (раньше при запуске python ругался на попытки взять элемент map по индексу - мы преобразовали его к списку).

import numpy as np

def npy2cube(grid, start, step, cube_f):

    bohr_to_angs = 0.529177210859  # const
    start = [x / bohr_to_angs for x in start]
    step = [x / bohr_to_angs for x in step]

    with open(cube_f, 'w') as cube_file:
        cube_file.write('CPMD CUBE FILE.\n'
                        'OUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n')
        cube_file.write("    1 %f %f %f\n" % (start[0], start[1], start[2]))
        cube_file.write("  %i    %f   0.000000     0.000000\n" %
                        (grid.shape[0], step[0]))
        cube_file.write("  %i    0.000000    %f    0.000000\n" %
                        (grid.shape[1], step[1]))
        cube_file.write("  %i    0.000000    0.000000    %f\n" %
                        (grid.shape[2], step[2]))
        cube_file.write("    1    0.000000    %f %f %f\n" %
                        (start[0], start[1], start[2]))

        i = 0
        for grid_item in np.nditer(grid, order='C'):
            if i < 5:
                cube_file.write('%f ' % (float(grid_item)))
                i += 1
            elif i == 5:
                cube_file.write('%f\n' % (float(grid_item)))
                i = 0

Считаем

Для каждого главного квантового числа n 1..3 перебираем орбитальные квантовые числа l 0..(n-1), для них, в свою очередь, перебираем магнитные m -l..l

for n in range(1, 4):
    d = 10 * n
    step = d / 10
    for l in range(0, n):
        for m in range(-l, l+1):
            grid = w(n, l, m, d)
            name = '{}_{}_{}'.format(n, l, m)
            npy2cube(grid, (-d,-d,-d), (step,step,step), '{}.cube'.format(name))

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

import __main__
__main__.pymol_argv = [ 'pymol', '-x' ]

import pymol
pymol.finish_launching()
from pymol import cmd,stored

os.chdir('/home/julybel/study/biopol/')

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])

#Для каждого набора из трех квантовых чисел - (n, l, m) таких чисел загружаем соответствующий .cube,
# визуализируем его (срез/поверхность) и сохраняем
for n in range(1, 4):
    for l in range(0, n):
        for m in range(-l, l+1):
            name = f'{n}_{l}_{m}'
            cmd.load(f'{name}.cube')
            cmd.volume(f'{name}-vol', name, ramp='ramp007')
            cmd.hide()
            cmd.show('volume', f'{name}-vol')
            cmd.reset()
            cmd.turn('x', 90)
            cmd.clip('slab', 0)
            time.sleep(0.5)
            cmd.do(f'png clip_{name}, width={480}, height={270}')
for n in range(1, 4):
    for l in range(0, n):
        for m in range(-l, l+1):
            name = f'{n}_{l}_{m}'
            cmd.load(f'{name}.cube')
            cmd.volume(f'{name}-vol', name, ramp='ramp007')
            cmd.hide()
            cmd.show('volume', f'{name}-vol')
            cmd.reset()
            cmd.turn('x', 45)
            cmd.turn('y', 45)
            cmd.turn('z', 45)
            time.sleep(0.5)
            cmd.do(f'png {name}, width={480}, height={270}')

Визуализация срезов орбиталей

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):
            name = f'{n}_{l}_{m}'
            f = plt.imread(f'clip_{name}.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):
            name = f'{n}_{l}_{m}'
            f = plt.imread(f'{name}.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()

os.chdir('./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 n in range(14):
    cmd.load(f'H-{n}.cube')
    cmd.volume(f'H-{n}-vol', f'H-{n}', ramp='ramp007')
    cmd.hide()
    cmd.show('volume', f'H-{n}-vol')
    cmd.reset()
    cmd.turn('x', 90)
    cmd.clip('slab', 0)
    time.sleep(0.5)
    cmd.do(f'png clip_orca_{n}, width={480}, height={270}')
    
for n in range(14):
    cmd.load(f'H-{n}.cube')
    cmd.volume(f'H-{n}-vol', f'H-{n}', ramp='ramp007')
    cmd.hide()
    cmd.show('volume', f'H-{n}-vol')
    cmd.reset()
    cmd.turn('x', 45)
    cmd.turn('y', 45)
    cmd.turn('z', 45)
    time.sleep(0.5)
    cmd.do(f'png orca_{n}, width={480}, height={270}')

Визуализация срезов орбиталей, рассчитаных с помощью Orca

orbs = ['1s', '2s', '2p', '2p', '2p', '3s', '3d', '3d', '3d', '3d', '3d', '3p', '3p', '3p']
orbs = (n for n in orbs)

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

i, j = 0, 0
for n in range(14):
    name = f'clip_orca_{n}'
    f = plt.imread(f'{name}.png')
    ax[i, j].axis('off')
    ax[i, j].imshow(f)
    ax[i, j].set_title(f'{next(orbs)}', {'fontsize': 35}, pad=10)
    j += 1
    if j == 2: i += 1; j = 0
        
plt.tight_layout()

Визуализация поверхностей орбиталей, рассчитаных с помощью Orca

orbs = ['1s', '2s', '2p', '2p', '2p', '3s', '3d', '3d', '3d', '3d', '3d', '3p', '3p', '3p']
orbs = (n for n in orbs)

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

i, j = 0, 0
for n in range(14):
    name = f'orca_{n}'
    f = plt.imread(f'{name}.png')
    ax[i, j].axis('off')
    ax[i, j].imshow(f)
    ax[i, j].set_title(f'{next(orbs)}', {'fontsize': 35}, pad=10)
    j += 1
    if j == 2: i += 1; j = 0
        
plt.tight_layout()

В целом, можно сказать, что вид орбиталей, рассчитанных нами и Orca не сильно различается. Вспомним, что мы брали, по итогу, сумму действительной и мнимой частей волновой функции, а также, без особых уточнений, приняли боровский радиус равным 1. Из визуальных соображений, можно отметить, что орбитали, построенные по результатам Orca выглядят более "пузатенькими" и крупными по размеру.