app.py

# 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