import numpy
import scipy.special
import scipy.misc
import npy2cube
from pathlib import Path
import psi4
import numpy as np
import matplotlib.pyplot as plt

Волновая фукция

def w(n,l,m,d):
    
    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.arccos(z/r(x,y,z)) # Полярный угол
    phi = lambda x,y,z: numpy.arctan(y/x) # Азимутальный угол
    
    # Боровский радиус
    a0 = 1.
    
    # Радиальная функция: n -- главное, l -- орбитальное квантовые числа
    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)
    # Волновая функция
    WF = lambda r,theta,phi,n,l,m: 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

    ### может тут чего то не хватает?  
    # Можно еще складывать реальную и мнимую части волновой функцииы
    #a = WF(r(x,y,z),theta(x,y,z),phi(x,y,z),n,l,m)
    #a = np.real(a) + np.imag(a)
    #
    return WF(r(x,y,z),theta(x,y,z),phi(x,y,z),n,l,m)

outdir = Path("./cube_files")
outdir.mkdir(parents=True, exist_ok=True)

# Зададим цикл по перебору квантовых чисел

for n in range(1,4):
    d = 10*n
    step = n
    for l in range(0,n):
        for m in range(0,l+1):
            grid = w(n,l,m,d)
            name=f'{n}-{l}-{m}' 
            # для сохранения нужно задать координаты старта grid и шаг по каждому направлению
            npy2cube.npy2cube(grid,
                              (-d,-d,-d),
                              (step,step,step), 
                              outdir / f"{name}.cube")

import py3Dmol
view = py3Dmol.view()
alpha = open(outdir / '1-0-0.cube','r').read()
view.addVolumetricData(alpha, "cube", {'isoval': 0.04, 'color': "red", 'opacity': 0.5})
view.zoomTo()
view.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

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

Этот скрипт использовала для генерации изображений:

# volume_ramp_new принимает список, в котором числа сгруппированы по 5: [Value, R, G, B, A]
# то есть значение плотности WF - цвет в формате RGB - прозрачность

cmd.volume_ramp_new('colors', 
    [-0.015, 0.00, 0.5, 1.00, 0.00, # Отрицательная фаза синенькая (+-)
      -0.01,  0.20, 0.70, 1.00, 0.20, # Добавила чуть красного
      -0.005, 0.30, 0.80, 1.00, 0.00, 
      0, 0, 0, 0, 0, 
      0.005, 1.00, 0.80, 0.30, 0.00, # положительная фаза красненькая
      0.01,  1.00, 0.70, 0.20, 0.20, 
      0.015, 1.00, 0.50, 0.00, 0.0]) 
names = []
cmd.set("specular", 1.0)
for n in range(1, 4):
    for l in range(0, n):
        for m in range(0, l+1):
            
            def_name = f'{n}-{l}-{m}'
            names.append(def_name)
            
            cmd.load(cube_path / f'{def_name}.cube')
            
            cmd.volume(f'{def_name}-volume', def_name, ramp='colors')

            cmd.hide()
            cmd.show('volume', f'{def_name}-volume')
            cmd.reset()
            cmd.turn('x', 45)
            cmd.turn('y', 45)
            cmd.turn('z', 45)
            cmd.draw(600)
            cmd.save(cube_path /f'{def_name}.png')
            cmd.delete('all')

import matplotlib.pyplot as plt
from pathlib import Path

pics = Path("./cube_files/pics")

Вот такие изображения получились:

orbitales = {0: 's', 1: 'p', 2: 'd'}

_, axarr = plt.subplots(5, 2, figsize=(10, 21))

row, col = 0, 0

for n in range(1,4):
    for l in range(0,n):
        for m in range(0,l+1):
            name=f"{n}-{l}-{m}"
            f = plt.imread(pics / f'{name}.png')
            axarr[row, col].axis('off')
            axarr[row, col].imshow(f)
            axarr[row, col].set_title(str(n)+orbitales[l], {'fontsize': 35}, pad=10)
            col += 1
            if col == 2:
                row += 1
                col = 0

plt.tight_layout()

Расчет орбиталей в программе psi4

psi4.core.set_output_file('output.dat')

g = '''
         0 1
         C     0.000000     0.000000     0.000000
'''

m = psi4.geometry(g)
psi4.set_options({"maxiter": 200, "fail_on_maxiter" :  True})
ener,wfn=psi4.energy('scf/cc-pvtz' , return_wfn=True, molecule = m )

ener

-37.60218711612774

psi4.set_options({
    'cubeprop_tasks': ['orbitals'],  
    'cubeprop_orbitals': [6, 7, 8, 9], # Номера орбиталей
    'cubic_grid_spacing': [0.2, 0.2, 0.2] # Плотность сетки 
})
psi4.cubeprop(wfn)

Ура, ничего не упало (первый раз засорила кодомо при генерации .cube файлов)

Такие орбитали получились:

pics = Path("./psi_pics")
_, axarr = plt.subplots(3, 3, figsize=(10, 21))

row=0
col=0

for  file in pics.glob("*"):
   
    f = plt.imread(file)
    axarr[row, col].axis('off')
    axarr[row, col].imshow(f)
    axarr[row, col].set_title(file.name[:-4])
    
    col+=1
    if col>=3:
        col=0
        row+=1
    

plt.tight_layout()

Не придумала как убрать обрезанные края

/usr/bin/bash: line 1: ipython: command not found