
import numpy
import math
import scipy.special
import scipy.misc

from IPython.display import display,Image


def npy2cube(grid, start, step, cube_f):
    '''
    PARAMETERS:
        grid : numpy array
            3-dimentional array, containing grid data
        start : tuple
            format: (x, y, z), coordinates of cube start point
        step: tuple
            format: (x, y, z), step size on 3 axes 
        cube_f: string
             name of output .cube file
    
    RETURNS:
        void    
    '''
    
    cube_string = ""
    bohr_to_angs = 0.529177210859 #const
    start = list(map(lambda x: x/bohr_to_angs, start))
    step = list(map(lambda x: x/bohr_to_angs, step))

    with open(cube_f, "w") as cube_file:
        ###HEADER###
        cube_file.write(" CPMD CUBE FILE.\nOUTER 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]))

        ###DATA###
        i = 0
        for x in range(grid.shape[0]):
            for y in range(grid.shape[1]):
                for z in range(grid.shape[2]):
                    if i < 5:
                        cube_file.write("%f " %(float(grid[x, y, z])))
                        i += 1
                    elif i == 5:
                        cube_file.write("%f\n" %(float(grid[x, y, z])))
                        i = 0
    return 0


def w(n,l,m,d):
    #n,l,m – квантовые числа

    
    #координаты x, y, z
    x,y,z = numpy.mgrid[-d:d:30j,-d:d:30j,-d:d:30j]
    
    #переход к сферическим координатам
    #радиус
    r = lambda x,y,z: numpy.sqrt(x**2+y**2+z**2)
    #угол тета
    theta = lambda x,y,z: -numpy.arcsin(z/r(x,y,z)) + math.pi/2
    #угол фи
    phi = lambda x,y,z: numpy.arctan2(y, x) + math.pi
    
    #радиус Бора
    a0 = 1
    
    #R – радиус функция
    #scipy.special.genlaguerre - обобщённый полином Лагерра
    R = lambda r,n,l: (2*r/n/a0)**l * numpy.exp(-r/n/a0) * scipy.special.genlaguerre(n-l-1,2*l+1)(2*r/n/a0)
    
    #нормировочный коэффициент
    k = lambda n,l: numpy.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: k(n, l) * R(r,n,l) * scipy.special.sph_harm(m,l,phi,theta)
    
    #модуль волновой функции
    absWF = lambda r,theta,phi,n,l,m: numpy.absolute(WF(r,theta,phi,n,l,m))**2
    
    final = WF(r(x,y,z),theta(x,y,z),phi(x,y,z),n,l,m)
    
    return numpy.real(final) + numpy.imag(final)


#шаг grid при заданном диапазоне от -d до d


for n in range(0,5):
    d = 10*n
    step = d/10
    for l in range(0,n):
        for m in range(0,l+1,1):
            grid= w(n, l, m, d)
            name='%s-%s-%s' % (n,l,m)
            npy2cube(grid,(-d, -d, -d),
                     (step, step, step),
                     name+'.cube')


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

import pymol
pymol.finish_launching()

from pymol import cmd,stored


from IPython.display import Image

Ab initio вычисление атомных орбиталей водорода¶

Вычисление по формулам¶

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