Update app.py

app.py

@@ -5,36 +5,7 @@ import torch
 from torch_geometric.data import Data
 import numpy as np
 import plotly.graph_objs as go
-
-st.title("PyTorch Geometric Structure Visualization")
-
-structure_type = st.sidebar.selectbox(
-    "Select Structure Type",
-    ("Sierpinski Triangle", "Spiral", "Plant Structure")
-)
-
-if structure_type == "Sierpinski Triangle":
-    depth = st.sidebar.slider("Recursion Depth", 0, 5, 3)
-    pos, edge_index = generate_sierpinski_triangle(depth)
-    data = Data(pos=torch.tensor(pos, dtype=torch.float), edge_index=torch.tensor(edge_index, dtype=torch.long))
-    fig = plot_graph_3d(pos, edge_index)
-    st.plotly_chart(fig)
-
-elif structure_type == "Spiral":
-    turns = st.sidebar.slider("Number of Turns", 1, 20, 5)
-    points_per_turn = st.sidebar.slider("Points per Turn", 10, 100, 50)
-    pos, edge_index = generate_spiral(turns, points_per_turn)
-    data = Data(pos=torch.tensor(pos, dtype=torch.float), edge_index=torch.tensor(edge_index, dtype=torch.long))
-    fig = plot_graph_3d(pos, edge_index)
-    st.plotly_chart(fig)
-
-elif structure_type == "Plant Structure":
-    iterations = st.sidebar.slider("L-system Iterations", 1, 5, 3)
-    angle = st.sidebar.slider("Branching Angle", 15, 45, 25)
-    pos, edge_index = generate_plant(iterations, angle)
-    data = Data(pos=torch.tensor(pos, dtype=torch.float), edge_index=torch.tensor(edge_index, dtype=torch.long))
-    fig = plot_graph_3d(pos, edge_index)
-    st.plotly_chart(fig)
+import math  # Moved import to the top
 
 # Function Definitions
 
@@ -62,13 +33,16 @@ def generate_sierpinski_triangle(depth):
 
     points = recurse_triangle(vertices[0], vertices[1], vertices[2], depth)
     pos = np.array(points)
+    # Remove duplicate points
+    pos = np.unique(pos, axis=0)
     # Create edges between points
     edge_index = []
     for i in range(0, len(pos), 3):
+        idx = i % len(pos)
         edge_index.extend([
-            [i, i+1],
-            [i+1, i+2],
-            [i+2, i]
+            [idx, (idx+1)%len(pos)],
+            [(idx+1)%len(pos), (idx+2)%len(pos)],
+            [(idx+2)%len(pos), idx]
         ])
     edge_index = np.array(edge_index).T
     return pos, edge_index
@@ -89,7 +63,6 @@ def generate_spiral(turns, points_per_turn):
 def generate_plant(iterations, angle):
     axiom = "F"
     rules = {"F": "F[+F]F[-F]F"}
-    import math
 
     def expand_axiom(axiom, rules, iterations):
         for _ in range(iterations):
@@ -104,10 +77,10 @@ def generate_plant(iterations, angle):
     stack = []
     pos_list = []
     edge_list = []
-    current_pos = np.array([0, 0, 0])
+    current_pos = np.array([0.0, 0.0, 0.0])
     pos_list.append(current_pos.copy())
     idx = 0
-    direction = np.array([0, 1, 0])
+    direction = np.array([0.0, 1.0, 0.0])
 
     for command in final_axiom:
         if command == 'F':
@@ -128,8 +101,6 @@ def generate_plant(iterations, angle):
             stack.append((current_pos.copy(), direction.copy(), idx))
         elif command == ']':
             current_pos, direction, idx = stack.pop()
-            pos_list.append(current_pos.copy())
-            idx += 1
 
     pos = np.array(pos_list)
     edge_index = np.array(edge_list).T
@@ -138,11 +109,19 @@ def generate_plant(iterations, angle):
 def rotation_matrix_3d(axis, theta):
     # Return the rotation matrix associated with rotation about the given axis by theta radians.
     axis = axis / np.linalg.norm(axis)
-    a = np.cos(theta / 2)
-    b, c, d = -axis * np.sin(theta / 2)
-    return np.array([[a*a + b*b - c*c - d*d, 2*(b*c - a*d),     2*(b*d + a*c)],
-                     [2*(b*c + a*d),     a*a + c*c - b*b - d*d, 2*(c*d - a*b)],
-                     [2*(b*d - a*c),     2*(c*d + a*b),     a*a + d*d - b*b - c*c]])
+    a = np.cos(theta)
+    b, c, d = axis * np.sin(theta)
+    return np.array([
+        [a + (1 - a) * axis[0] * axis[0],
+         (1 - a) * axis[0] * axis[1] - axis[2] * np.sin(theta),
+         (1 - a) * axis[0] * axis[2] + axis[1] * np.sin(theta)],
+        [(1 - a) * axis[1] * axis[0] + axis[2] * np.sin(theta),
+         a + (1 - a) * axis[1] * axis[1],
+         (1 - a) * axis[1] * axis[2] - axis[0] * np.sin(theta)],
+        [(1 - a) * axis[2] * axis[0] - axis[1] * np.sin(theta),
+         (1 - a) * axis[2] * axis[1] + axis[0] * np.sin(theta),
+         a + (1 - a) * axis[2] * axis[2]]
+    ])
 
 def plot_graph_3d(pos, edge_index):
     x, y, z = pos[:, 0], pos[:, 1], pos[:, 2]
@@ -181,6 +160,43 @@ def plot_graph_3d(pos, edge_index):
             zaxis_title='Z',
             aspectmode='data'
         ),
-        showlegend=False
+        showlegend=False,
+        margin=dict(l=0, r=0, b=0, t=0)  # Optional: adjust margins
     )
     return fig
+
+# Main App Code
+
+def main():
+    st.title("PyTorch Geometric Structure Visualization")
+
+    structure_type = st.sidebar.selectbox(
+        "Select Structure Type",
+        ("Sierpinski Triangle", "Spiral", "Plant Structure")
+    )
+
+    if structure_type == "Sierpinski Triangle":
+        depth = st.sidebar.slider("Recursion Depth", 0, 5, 3)
+        pos, edge_index = generate_sierpinski_triangle(depth)
+        data = Data(pos=torch.tensor(pos, dtype=torch.float), edge_index=torch.tensor(edge_index, dtype=torch.long))
+        fig = plot_graph_3d(pos, edge_index)
+        st.plotly_chart(fig)
+
+    elif structure_type == "Spiral":
+        turns = st.sidebar.slider("Number of Turns", 1, 20, 5)
+        points_per_turn = st.sidebar.slider("Points per Turn", 10, 100, 50)
+        pos, edge_index = generate_spiral(turns, points_per_turn)
+        data = Data(pos=torch.tensor(pos, dtype=torch.float), edge_index=torch.tensor(edge_index, dtype=torch.long))
+        fig = plot_graph_3d(pos, edge_index)
+        st.plotly_chart(fig)
+
+    elif structure_type == "Plant Structure":
+        iterations = st.sidebar.slider("L-system Iterations", 1, 5, 3)
+        angle = st.sidebar.slider("Branching Angle", 15, 45, 25)
+        pos, edge_index = generate_plant(iterations, angle)
+        data = Data(pos=torch.tensor(pos, dtype=torch.float), edge_index=torch.tensor(edge_index, dtype=torch.long))
+        fig = plot_graph_3d(pos, edge_index)
+        st.plotly_chart(fig)
+
+if __name__ == "__main__":
+    main()