Create app.py

app.py

@@ -0,0 +1,186 @@
+# Streamlit App: PyTorch Geometric Structure Visualization
+
+import streamlit as st
+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)
+
+# Function Definitions
+
+def generate_sierpinski_triangle(depth):
+    # Generate the vertices of the initial triangle
+    vertices = np.array([
+        [0, 0, 0],
+        [1, 0, 0],
+        [0.5, np.sqrt(3)/2, 0]
+    ])
+
+    # Function to recursively generate points
+    def recurse_triangle(v1, v2, v3, depth):
+        if depth == 0:
+            return [v1, v2, v3]
+        else:
+            # Calculate midpoints
+            m12 = (v1 + v2) / 2
+            m23 = (v2 + v3) / 2
+            m31 = (v3 + v1) / 2
+            # Recursively subdivide
+            return (recurse_triangle(v1, m12, m31, depth - 1) +
+                    recurse_triangle(m12, v2, m23, depth - 1) +
+                    recurse_triangle(m31, m23, v3, depth - 1))
+
+    points = recurse_triangle(vertices[0], vertices[1], vertices[2], depth)
+    pos = np.array(points)
+    # Create edges between points
+    edge_index = []
+    for i in range(0, len(pos), 3):
+        edge_index.extend([
+            [i, i+1],
+            [i+1, i+2],
+            [i+2, i]
+        ])
+    edge_index = np.array(edge_index).T
+    return pos, edge_index
+
+def generate_spiral(turns, points_per_turn):
+    total_points = turns * points_per_turn
+    theta_max = 2 * np.pi * turns
+    theta = np.linspace(0, theta_max, total_points)
+    z = np.linspace(0, 1, total_points)
+    r = z  # Spiral expanding in radius
+    x = r * np.cos(theta)
+    y = r * np.sin(theta)
+    pos = np.vstack((x, y, z)).T
+    # Edges connect sequential points
+    edge_index = np.array([np.arange(total_points - 1), np.arange(1, total_points)])
+    return pos, edge_index
+
+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):
+            new_axiom = ""
+            for ch in axiom:
+                new_axiom += rules.get(ch, ch)
+            axiom = new_axiom
+        return axiom
+
+    final_axiom = expand_axiom(axiom, rules, iterations)
+
+    stack = []
+    pos_list = []
+    edge_list = []
+    current_pos = np.array([0, 0, 0])
+    pos_list.append(current_pos.copy())
+    idx = 0
+    direction = np.array([0, 1, 0])
+
+    for command in final_axiom:
+        if command == 'F':
+            next_pos = current_pos + direction
+            pos_list.append(next_pos.copy())
+            edge_list.append([idx, idx + 1])
+            current_pos = next_pos
+            idx += 1
+        elif command == '+':
+            theta = np.radians(angle)
+            rotation_matrix = rotation_matrix_3d(np.array([0, 0, 1]), theta)
+            direction = rotation_matrix @ direction
+        elif command == '-':
+            theta = np.radians(-angle)
+            rotation_matrix = rotation_matrix_3d(np.array([0, 0, 1]), theta)
+            direction = rotation_matrix @ direction
+        elif command == '[':
+            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
+    return pos, edge_index
+
+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]])
+
+def plot_graph_3d(pos, edge_index):
+    x, y, z = pos[:, 0], pos[:, 1], pos[:, 2]
+    edge_x = []
+    edge_y = []
+    edge_z = []
+
+    for i in range(edge_index.shape[1]):
+        src = edge_index[0, i]
+        dst = edge_index[1, i]
+        edge_x += [x[src], x[dst], None]
+        edge_y += [y[src], y[dst], None]
+        edge_z += [z[src], z[dst], None]
+
+    edge_trace = go.Scatter3d(
+        x=edge_x, y=edge_y, z=edge_z,
+        line=dict(width=2, color='gray'),
+        hoverinfo='none',
+        mode='lines')
+
+    node_trace = go.Scatter3d(
+        x=x, y=y, z=z,
+        mode='markers',
+        marker=dict(
+            size=4,
+            color='red',
+        ),
+        hoverinfo='none'
+    )
+
+    fig = go.Figure(data=[edge_trace, node_trace])
+    fig.update_layout(
+        scene=dict(
+            xaxis_title='X',
+            yaxis_title='Y',
+            zaxis_title='Z',
+            aspectmode='data'
+        ),
+        showlegend=False
+    )
+    return fig