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


H1 = np.linspace(0.5, 3, 100)
H2 = H1
def scan_energy_surface(sita):
    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)
    return e_surface_df


Sita = 90
fig = go.Figure(data=[go.Surface(z=scan_energy_surface(math.radians(Sita)), x=H1, y=H2)])
for i in range(10, 40, 10):
    Sita = Sita + i
    fig.add_traces(go.Surface(z=scan_energy_surface(math.radians(Sita)) + i * 3, x=H1, y=H2, opacity=0.8))
fig.update_layout(title="Water molecule energy surface with the angle varied from 90 degrees to 120 degrees", scene=dict(xaxis_title='H_1-O bond distance', yaxis_title='H_2-O bond distance', zaxis_title='Energy'))
fig.show()

结果如下：

难点：增加3D图形时要用add_traces而不是update_traces
如果前面已经出现过一次surface，这次再用update_traces，会将前面结果覆盖，只保留最后的结果。