本文记录一个立体几何相关的基于距离求解坐标的函数。
先直接看效果:
该部分代码使用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 pd
import 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渲染的一些脚本控件啥的还是要联网。