python画出水分子势能面

import plotly.graph_objects as go
import pandas as pd
from ase import Atoms
import numpy as np
from ase.calculators.emt import EMT
import math


sita = math.radians(100)
H1 = np.linspace(0.5, 1.5, 100)
H2 = H1
e_surface = {}
for H1_position in H1:
    e = []
    for H2_position in H2:
        water = Atoms('HOH', positions=[[H1_position, 0, 0], [0, 0, 0], [H2_position * np.cos(sita), H2_position * np.sin(sita), 0]])
        water.calc = EMT()
        e.append(water.get_potential_energy())
    e_surface_temp = {H1_position: e}
    e_surface.update(e_surface_temp)
e_surface_df = pd.DataFrame(e_surface)
fig = go.Figure(data=[go.Surface(z=e_surface_df, x=H1, y=H2)])
fig.update_traces(contours_z=dict(show=True, usecolormap=True,
                                  highlightcolor="limegreen", project_z=True))
fig.update_layout(title='sita=100', autosize=True)
fig.show()

结果如下：

夹角100°不变，两氢原子与氧原子的距离是变化的。鼠标放在势能面上可以看到两距离和对应的能量。