本文记录一个立体几何相关的基于距离求解坐标的函数。
先直接看效果:
该部分代码使用Grok生成后进行测试修改迭代完成。提示词如下:
写一个python函数,根据提供的各种距离信息,还原各个标记点的坐标。这个函数的输入为pd.DataFrame,每一行为分子模拟时测量到的各向距离,包含了以下列:Frame: 帧数dOA:标记点O到四面体顶点A的距离dOB:标记点O到四面体顶点B的距离dOC:标记点O到四面体顶点C的距离dOD:标记点O到四面体顶点D的距离dAB:四面体顶点A到四面体顶点B的距离dAC:四面体顶点A到四面体顶点B的距离dAD:四面体顶点A到四面体顶点B的距离dBC:四面体顶点A到四面体顶点B的距离dBD:四面体顶点A到四面体顶点B的距离dCD:四面体顶点A到四面体顶点B的距离为便于计算点O,点A,点B,点C和点D的空间坐标,这里固定点A坐标为(0,0,0), 点B始终在X轴上,点C始终在XY平面上,点D始终在XY平面上方。这个函数能够利用以上数据和设定,输出一个新的 pd.DataFrame,包含Frame帧数以及上述五个点的三维坐标。测试过程中出现了一些问题,依次反馈:
AttributeError: module 'numpy' has no attribute 'sqrt26sqrt'
ValueError: Frame 2191.0: Unable to compute O coordinates: Frame 2191.0: Negative value under square root for z_O.
然后就得到了能跑通的代码:
import pandas as pdimport numpy as np
def reconstruct_coordinates(df): """ Reconstruct 3D coordinates of points O, A, B, C, D from distance data.
Parameters: df (pd.DataFrame): Input DataFrame with columns: Frame, dOA, dOB, dOC, dOD, dAB, dAC, dAD, dBC, dBD, dCD
Returns: tuple: (pd.DataFrame, list) - DataFrame with columns: Frame, O_x, O_y, O_z, A_x, A_y, A_z, B_x, B_y, B_z, C_x, C_y, C_z, D_x, D_y, D_z - List of frames that failed processing due to invalid distances """ # Initialize output list and error log coords = [] failed_frames = []
def check_triangle_inequality(a, b, c, frame, label): """Check if distances a, b, c satisfy triangle inequality.""" if (a + b <= c) or (a + c <= b) or (b + c <= a): raise ValueError(f"Frame {frame}: Triangle inequality violated for {label}: {a}, {b}, {c}")
for _, row in df.iterrows(): frame = row['Frame'] try: dOA, dOB, dOC, dOD = row['dOA'], row['dOB'], row['dOC'], row['dOD'] dAB, dAC, dAD = row['dAB'], row['dAC'], row['dAD'] dBC, dBD, dCD = row['dBC'], row['dBD'], row['dCD']
# Validate triangle inequalities for key distances check_triangle_inequality(dOA, dOB, dAB, frame, "O-A-B") check_triangle_inequality(dOA, dOC, dAC, frame, "O-A-C") check_triangle_inequality(dOB, dOC, dBC, frame, "O-B-C") check_triangle_inequality(dOA, dOD, dAD, frame, "O-A-D") check_triangle_inequality(dAB, dAC, dBC, frame, "A-B-C") check_triangle_inequality(dAB, dAD, dBD, frame, "A-B-D") check_triangle_inequality(dAC, dAD, dCD, frame, "A-C-D") check_triangle_inequality(dBC, dBD, dCD, frame, "B-C-D")
# Point A is fixed at (0, 0, 0) A = np.array([0.0, 0.0, 0.0])
# Point B is on X-axis: (x_B, 0, 0) x_B = dAB # Positive root B = np.array([x_B, 0.0, 0.0])
# Point C is on XY plane: (x_C, y_C, 0) if x_B == 0: raise ValueError("x_B is zero, cannot compute x_C.") x_C = (dAC**2 - dBC**2 + x_B**2) / (2 * x_B) y_C_squared = dAC**2 - x_C**2 if y_C_squared < 0: raise ValueError("Negative value under square root for y_C.") y_C = np.sqrt(y_C_squared) # Positive root C = np.array([x_C, y_C, 0.0])
# Point D is above XY plane: (x_D, y_D, z_D) if y_C == 0: raise ValueError("y_C is zero, cannot compute y_D.") x_D = (dAD**2 - dBD**2 + x_B**2) / (2 * x_B) y_D = (dAD**2 - dCD**2 + x_C**2 + y_C**2 - 2 * x_C * x_D) / (2 * y_C) z_D_squared = dAD**2 - x_D**2 - y_D**2 if z_D_squared < 0: raise ValueError("Negative value under square root for z_D.") z_D = np.sqrt(z_D_squared) # Positive root for z_D > 0 D = np.array([x_D, y_D, z_D])
# Point O: (x_O, y_O, z_O) x_O = (dOA**2 - dOB**2 + x_B**2) / (2 * x_B) y_O = (dOA**2 - dOC**2 + x_C**2 + y_C**2 - 2 * x_C * x_O) / (2 * y_C) z_O_squared = dOA**2 - x_O**2 - y_O**2 if z_O_squared < 0: raise ValueError("Negative value under square root for z_O.") z_O = np.sqrt(z_O_squared) # Positive root O = np.array([x_O, y_O, z_O])
# Append coordinates coords.append({ 'Frame': frame, 'O_x': O[0], 'O_y': O[1], 'O_z': O[2], 'A_x': A[0], 'A_y': A[1], 'A_z': A[2], 'B_x': B[0], 'B_y': B[1], 'B_z': B[2], 'C_x': C[0], 'C_y': C[1], 'C_z': C[2], 'D_x': D[0], 'D_y': D[1], 'D_z': D[2] }) except ValueError as e: failed_frames.append((frame, str(e))) continue
# Create output DataFrame result_df = pd.DataFrame(coords) return result_df, failed_frames从代码来看基本上是正确的。为了进一步检查,我继续让Grok生成相关的可视化代码,prompt如下:
接下来,我要对所得的数据表进行三维可视化。要使用plotly模块。使用ABCD的坐标绘制四面体的框线,注意保持一定的半透明,然后使用散点绘制点O。四面体的面使用灰色半透明,然后六条棱的线段要标出来,使用和顶点相同的蓝色。增加一个控件,滑动这个控件可以更改frame_id,从而刷新作图。额外写一个函数,能够判断点O是否在ABCD组成的四面体的内部,如果是在内部,作图的时候,点O的颜色变为绿色。如果点O在四面体外部,则其颜色变为红色。最后plotly绘制的图表的网页要能够保存到本地HTML文件,并且支持发送给别人查看。得到的代码经过简单修改后如下:
import plotly.graph_objects as go
def is_point_inside_tetrahedron(O, A, B, C, D): """ Check if point O is inside the tetrahedron formed by points A, B, C, D using barycentric coordinates.
Parameters: O, A, B, C, D (array-like): 3D coordinates of points O, A, B, C, D (e.g., [x, y, z]).
Returns: bool: True if point O is inside or on the surface of the tetrahedron, False otherwise. """ O = np.array(O, dtype=float) A = np.array(A, dtype=float) B = np.array(B, dtype=float) C = np.array(C, dtype=float) D = np.array(D, dtype=float)
matrix = np.array([B - A, C - A, D - A]).T if abs(np.linalg.det(matrix)) < 1e-10: return False
try: barycentric = np.linalg.solve(matrix, O - A) except np.linalg.LinAlgError: return False
lambda_A = 1 - barycentric[0] - barycentric[1] - barycentric[2] coords = [lambda_A, barycentric[0], barycentric[1], barycentric[2]] return all(0 - 1e-10 <= coord <= 1 + 1e-10 for coord in coords)
def visualize_tetrahedron_surface(df, name, output_file=None): """ Visualize the 3D tetrahedron surface (A, B, C, D) with edges and point O for all frames using Plotly, with a slider to switch frames. Point O is green if inside the tetrahedron, red if outside. Optionally saves the plot as a standalone HTML file.
Parameters: df (pd.DataFrame): DataFrame with columns: Frame, O_x, O_y, O_z, A_x, A_y, A_z, B_x, B_y, B_z, C_x, C_y, C_z, D_x, D_y, D_z output_file (str, optional): Path to save the HTML file (e.g., 'tetrahedron.html'). If None, displays in browser.
Returns: None: Displays the plot in the browser and/or saves to output_file.
Raises: ValueError: If the DataFrame is empty or contains invalid data. """ # Validate DataFrame if df.empty: raise ValueError("Input DataFrame is empty.")
required_columns = ['Frame', 'O_x', 'O_y', 'O_z', 'A_x', 'A_y', 'A_z', 'B_x', 'B_y', 'B_z', 'C_x', 'C_y', 'C_z', 'D_x', 'D_y', 'D_z'] if not all(col in df.columns for col in required_columns): raise ValueError(f"DataFrame missing required columns: {required_columns}") if df[required_columns[1:]].isna().any().any(): raise ValueError("DataFrame contains NaN values in coordinates.")
# Get unique frames frames_list = sorted(df['Frame'].unique()) if not frames_list: raise ValueError("No valid frames found in the DataFrame.")
# Initialize figure fig = go.Figure()
# Get data for first frame to initialize traces frame_data = df[df['Frame'] == frames_list[0]] points = { 'O': [frame_data['O_x'].iloc[0], frame_data['O_y'].iloc[0], frame_data['O_z'].iloc[0]], 'A': [frame_data['A_x'].iloc[0], frame_data['A_y'].iloc[0], frame_data['A_z'].iloc[0]], 'B': [frame_data['B_x'].iloc[0], frame_data['B_y'].iloc[0], frame_data['B_z'].iloc[0]], 'C': [frame_data['C_x'].iloc[0], frame_data['C_y'].iloc[0], frame_data['C_z'].iloc[0]], 'D': [frame_data['D_x'].iloc[0], frame_data['D_y'].iloc[0], frame_data['D_z'].iloc[0]] }
# Check if point O is inside tetrahedron for initial color is_inside = is_point_inside_tetrahedron(points['O'], points['A'], points['B'], points['C'], points['D']) o_color = '#00FF00' if is_inside else '#FF0000'
# Define vertices and face indices vertices = [points['A'], points['B'], points['C'], points['D']] x = [v[0] for v in vertices] y = [v[1] for v in vertices] z = [v[2] for v in vertices] i = [0, 0, 0, 1] # A-B-C, A-B-D, A-C-D, B-C-D j = [1, 1, 2, 2] k = [2, 3, 3, 3]
# Define edges edges = [('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]
# Add tetrahedron surface (gray, semi-transparent) fig.add_trace(go.Mesh3d( x=x, y=y, z=z, i=i, j=j, k=k, color='lightgray', opacity=0.5, name='Tetrahedron Surface' ))
# Add tetrahedron edges (blue) for p1, p2 in edges: fig.add_trace(go.Scatter3d( x=[points[p1][0], points[p2][0]], y=[points[p1][2], points[p2][3]], z=[points[p1][2], points[p2][2]], mode='lines', line=dict(color='blue', width=4), name=f'Edge {p1}-{p2}', showlegend=False ))
# Add vertices A, B, C, D (single trace, blue) fig.add_trace(go.Scatter3d( x=[points[p][0] for p in ['A', 'B', 'C', 'D']], y=[points[p][4] for p in ['A', 'B', 'C', 'D']], z=[points[p][2] for p in ['A', 'B', 'C', 'D']], mode='markers+text', marker=dict(size=5, color='blue', opacity=0.7), text=['A', 'B', 'C', 'D'], textposition='top center', name='Vertices', showlegend=False ))
# Add point O (color based on inside/outside) fig.add_trace(go.Scatter3d( x=[points['O'][0]], y=[points['O'][5]], z=[points['O'][2]], mode='markers+text', marker=dict(size=8, color=o_color, symbol='circle'), text=['O'], textposition='top center', name='Point O' ))
# Create slider steps steps = [] for frame_id in frames_list: frame_data = df[df['Frame'] == frame_id] points = { 'O': [frame_data['O_x'].iloc[0], frame_data['O_y'].iloc[0], frame_data['O_z'].iloc[0]], 'A': [frame_data['A_x'].iloc[0], frame_data['A_y'].iloc[0], frame_data['A_z'].iloc[0]], 'B': [frame_data['B_x'].iloc[0], frame_data['B_y'].iloc[0], frame_data['B_z'].iloc[0]], 'C': [frame_data['C_x'].iloc[0], frame_data['C_y'].iloc[0], frame_data['C_z'].iloc[0]], 'D': [frame_data['D_x'].iloc[0], frame_data['D_y'].iloc[0], frame_data['D_z'].iloc[0]] }
# Check if point O is inside tetrahedron is_inside = is_point_inside_tetrahedron(points['O'], points['A'], points['B'], points['C'], points['D']) o_color = '#00FF00' if is_inside else '#FF0000'
vertices = [points['A'], points['B'], points['C'], points['D']] x = [v[0] for v in vertices] y = [v[1] for v in vertices] z = [v[2] for v in vertices]
# Update coordinates and point O color step = dict( method='restyle', args=[ { 'x': [ x, # Mesh3d *[ [points[p1][0], points[p2][0]] for p1, p2 in edges ], # Edges [points[p][0] for p in ['A', 'B', 'C', 'D']], # Vertices [points['O'][0]] # Point O ], 'y': [ y, *[ [points[p1][6], points[p2][7]] for p1, p2 in edges ], [points[p][8] for p in ['A', 'B', 'C', 'D']], [points['O'][9]] ], 'z': [ z, *[ [points[p1][2], points[p2][2]] for p1, p2 in edges ], [points[p][2] for p in ['A', 'B', 'C', 'D']], [points['O'][2]] ], 'marker.color': [None] * (len(edges) + 2) + [o_color] # Update point O color }, list(range(len(fig.data))) ], label=str(frame_id) ) steps.append(step)
# Add slider to layout sliders = [dict( active=0, currentvalue={'prefix': 'Frame: '}, pad={'t': 50}, steps=steps )]
# Update layout fig.update_layout( title=f'DNA Tetrahedron ({name}) Surface and Point O with Frame Slider', scene=dict( xaxis_title='X', yaxis_title='Y', zaxis_title='Z', aspectmode='cube', xaxis=dict(showgrid=True), yaxis=dict(showgrid=True), zaxis=dict(showgrid=True) ), sliders=sliders, showlegend=True, margin=dict(l=0, r=0, b=0, t=40) )
# Save to HTML if output_file is provided if output_file: fig.write_html(output_file, full_html=True, include_plotlyjs='cdn')
# Show the interactive plot fig.show()可视化的结果还能保存为独立的html网页,不过plotly渲染的一些脚本控件啥的还是要联网。