Rotating 3D Sphere with Light Source in Python + Pygame

import pygame
import numpy as np

# Initialize Pygame
pygame.init()

# Set up the display
WIDTH = 800
HEIGHT = 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Optimized Rotating Sphere with Smooth Shading")

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
YELLOW = (255, 255, 0)

# Set up font for display
font = pygame.font.Font(None, 24)

# Load point cloud data
# Calculate the number of points per ring
POINTS_PER_RING = 10
NUM_RINGS = len(points) // POINTS_PER_RING

# Visualization parameters
scale = 200
z_offset = 500
rotation_speed = 0.02
point_size = 2

# Light parameters
light_distance = 800  # Distance from center
light_height = 400    # Height above sphere
light_pos = np.array([light_distance, light_height, 0.])
ambient_intensity = 0.2
diffuse_intensity = 0.7
specular_intensity = 0.3
specular_power = 32

def create_triangles():
    """Create triangle indices for connecting points"""
    triangles = []
    for ring in range(NUM_RINGS - 1):
        for point in range(POINTS_PER_RING):
            current = ring * POINTS_PER_RING + point
            next_point = ring * POINTS_PER_RING + ((point + 1) % POINTS_PER_RING)
            next_ring_current = (ring + 1) * POINTS_PER_RING + point
            next_ring_next = (ring + 1) * POINTS_PER_RING + ((point + 1) % POINTS_PER_RING)

            # Original winding order might be inconsistent
            # Let's define triangles with consistent counter-clockwise order
            triangles.append([current, next_point, next_ring_next])
            triangles.append([current, next_ring_next, next_ring_current])
    return np.array(triangles)

triangles = create_triangles()

# Precompute vertex normals (normalized position vectors for a sphere)
vertex_normals = points / np.linalg.norm(points, axis=1)[:, np.newaxis]

def rotate_y(points, angle):
    """Rotate points around Y axis"""
    cos_a = np.cos(angle)
    sin_a = np.sin(angle)
    rotation_matrix = np.array([
        [cos_a, 0, sin_a],
        [0, 1, 0],
        [-sin_a, 0, cos_a]
    ])
    return np.dot(points, rotation_matrix.T)

def project_points(points, scale_factor):
    """Project 3D points to 2D coordinates"""
    x, y, z = (points * scale_factor).T
    factor = z_offset / (z + z_offset)
    x_projected = x * factor + WIDTH // 2
    y_projected = y * factor + HEIGHT // 2
    projected_points = np.stack([x_projected, y_projected], axis=-1)
    return projected_points, z

def calculate_lighting(normals, vertex_positions, viewer_pos, light_pos):
    """Vectorized calculation of lighting using Phong reflection model"""
    # View direction
    view_dir = viewer_pos - vertex_positions
    view_dir /= np.linalg.norm(view_dir, axis=1)[:, np.newaxis]

    # Light direction
    light_dir = light_pos - vertex_positions
    light_dir /= np.linalg.norm(light_dir, axis=1)[:, np.newaxis]

    # Ambient component
    ambient = ambient_intensity

    # Diffuse component
    diff = np.maximum(0, np.einsum('ij,ij->i', normals, light_dir))
    diffuse = diff * diffuse_intensity

    # Specular component
    reflect_dir = 2 * np.einsum('ij,ij->i', normals, light_dir)[:, np.newaxis] * normals - light_dir
    spec = np.maximum(0, np.einsum('ij,ij->i', view_dir, reflect_dir)) ** specular_power
    specular = spec * specular_intensity

    # Total intensity
    total_intensity = np.clip(ambient + diffuse + specular, 0, 1)
    return total_intensity

def draw_shaded_mesh(screen, points, normals, triangles, scale_factor, light_pos, angle):
    """Draw the triangular mesh with smooth shading"""
    viewer_pos = np.array([0, 0, z_offset / scale_factor])
    light_pos_scaled = light_pos / scale_factor

    # Project points
    projected_points, z_values = project_points(points, scale_factor)

    # Calculate lighting per vertex
    vertex_lighting = calculate_lighting(normals, points, viewer_pos, light_pos_scaled)

    # Prepare triangle data
    triangle_indices = triangles
    v0 = points[triangle_indices[:, 0]]
    v1 = points[triangle_indices[:, 1]]
    v2 = points[triangle_indices[:, 2]]
    edge1 = v1 - v0
    edge2 = v2 - v0
    face_normals = np.cross(edge1, edge2)
    face_normals /= np.linalg.norm(face_normals, axis=1)[:, np.newaxis]

    # Backface culling based on face normal Z component
    # Since camera looks along negative Z, faces with normals pointing in +Z direction are culled
    visible = face_normals[:, 2] <= 0

    # Filter visible triangles
    triangle_indices = triangle_indices[visible]
    triangle_z = -np.mean(z_values[triangle_indices], axis=1)
    triangle_points_2d = projected_points[triangle_indices]
    triangle_intensities = vertex_lighting[triangle_indices]

    # Compute triangle colors
    triangle_colors = (np.mean(triangle_intensities, axis=1) * 255).astype(np.uint8)

    # Sort triangles by depth
    sort_indices = np.argsort(triangle_z)

    # Draw triangles
    for idx in sort_indices:
        points_2d = triangle_points_2d[idx]
        color_intensity = triangle_colors[idx]
        color = (color_intensity, color_intensity, color_intensity)
        pygame.draw.polygon(screen, color, points_2d.astype(int))

    # Draw light source
    light_projected, _ = project_points(light_pos_scaled[np.newaxis, :], scale_factor)
    pygame.draw.circle(screen, YELLOW, light_projected[0].astype(int), 5)


# Initialize light parameters
light_distance = 800
light_height = 0
light_pos = np.array([0., light_height, -light_distance])
light_angle = 0.0
light_orbit = True  # Ensure light orbit is enabled

# Game loop
running = True
angle = 0.0
clock = pygame.time.Clock()
paused = False
show_points = True
show_mesh = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                paused = not paused
            elif event.key == pygame.K_UP:
                scale *= 1.1
            elif event.key == pygame.K_DOWN:
                scale /= 1.1
            elif event.key == pygame.K_LEFT:
                rotation_speed /= 1.2
            elif event.key == pygame.K_RIGHT:
                rotation_speed *= 1.2
            elif event.key == pygame.K_r:
                scale = 200
                rotation_speed = 0.02
                angle = 0
                light_distance = 800
                light_height = 0
                light_pos = np.array([0., light_height, -light_distance])
                light_angle = 0.0
                light_orbit = True  # Ensure light orbit is enabled
            elif event.key == pygame.K_p:
                show_points = not show_points
            elif event.key == pygame.K_m:
                show_mesh = not show_mesh
            elif event.key == pygame.K_l:
                light_orbit = not light_orbit
            elif event.key == pygame.K_w:
                light_height += 50
                light_pos[1] = light_height
            elif event.key == pygame.K_s:
                light_height -= 50
                light_pos[1] = light_height
            elif event.key == pygame.K_a:
                if light_distance > 100:
                    light_distance *= 0.9
                    light_pos[2] = -light_distance
            elif event.key == pygame.K_d:
                light_distance *= 1.1
                light_pos[2] = -light_distance

    # Clear screen
    screen.fill(BLACK)

    # Rotate the points and normals
    if not paused:
        angle += rotation_speed
        rotated_points = rotate_y(points, angle)
        rotated_normals = rotate_y(vertex_normals, angle)
        if light_orbit:
            light_angle += rotation_speed * 5
    else:
        rotated_points = rotate_y(points, angle)
        rotated_normals = rotate_y(vertex_normals, angle)

    # Update light position
    if light_orbit:
        # Rotate light around Y-axis
        rotated_light = rotate_y(light_pos[np.newaxis, :], light_angle)[0]
    else:
        rotated_light = light_pos  # Use static position if not orbiting

    # Draw mesh
    if show_mesh:
        draw_shaded_mesh(screen, rotated_points, rotated_normals, triangles, scale, rotated_light, angle)

    # Draw points
    if show_points:
        projected_points, _ = project_points(rotated_points, scale)
        for point in projected_points.astype(int):
            pygame.draw.circle(screen, WHITE, point, point_size)

    # Display controls and info
    info_texts = [
        f"FPS: {int(clock.get_fps())}",
        "Controls:",
        "SPACE: Pause/Resume",
        "UP/DOWN: Scale",
        "LEFT/RIGHT: Rotation Speed",
        "P: Toggle Points",
        "M: Toggle Mesh",
        "L: Toggle Light Orbit",
        "R: Reset",
        f"Scale: {scale:.1f}",
        f"Rotation Speed: {rotation_speed:.3f}"
    ]

    for i, text in enumerate(info_texts):
        text_surface = font.render(text, True, GREEN)
        screen.blit(text_surface, (10, 10 + i * 20))

    pygame.display.flip()
    clock.tick(60)

pygame.quit()