Penrose Diagram v2

Having fun with Google ai Studio (Gemini Pro 2.5 Preview 05-06).

import numpy as np
from PIL import Image
from scipy.special import lambertw
import math

# Image dimensions
WIDTH = 4096
HEIGHT = 4096

# Schwarzschild radius parameter, 2M. We set 2M=1 for simplicity in formulas.
TWOM = 1.0 
M = TWOM / 2.0 # M = 0.5

# --- Feature Definition Parameters ---
# Thresholds for identifying singularities and horizons based on r-value
R_THRESH_SINGULARITY = 0.03 * TWOM  # r near 0 is singularity-like
R_THRESH_HORIZON = 0.03 * TWOM      # r near 2M (i.e., 1.0 here) is horizon-like

# Parameters for r=const lines
NUM_R_LINES = 10  # Number of r-lines to draw
# Define r-values for these lines, from near singularity to outside horizon
R_LINE_VALS_INNER = np.linspace(R_THRESH_SINGULARITY * 1.5, TWOM - R_THRESH_HORIZON * 1.5, NUM_R_LINES // 2)
R_LINE_VALS_OUTER = np.linspace(TWOM + R_THRESH_HORIZON * 1.5, 3.0 * TWOM, NUM_R_LINES - (NUM_R_LINES // 2) )
R_LINE_VALUES = np.concatenate((R_LINE_VALS_INNER, R_LINE_VALS_OUTER))
R_LINE_WIDTH_PARAM = 0.025 * TWOM # Visual thickness of r-lines (in r-units)

# Parameters for t=const lines
NUM_T_LINES = 24  # Number of t-lines (spokes)
# t-lines are based on normalized effective t or angle. Width is relative.
T_LINE_WIDTH_PARAM = 0.03 # Visual thickness of t-lines (fraction of inter-line spacing)
# --- End Feature Definition Parameters ---

# Helper function to calculate Schwarzschild r from Kruskal-related product P
# P = (1 - r/2M) * exp(r/2M)
def get_r_from_P(P_val):
    # Handle P_val near special points carefully
    if math.isclose(P_val, 0.0): # Horizon (r = 2M)
        return TWOM
    if math.isclose(P_val, 1.0): # Singularity (r = 0 for 2M=1 convention of P)
        return 0.0

    # For r < 2M (Regions II, IV), P_val is in (0, 1]. (1-r/2M) > 0.
    if P_val > 0:
        # We solve (1-x)e^x = P_val where x = r/2M.
        # Let y = 1-x. Then y * e^(1-y) = P_val => y * e * e^(-y) = P_val
        # => y * e^(-y) = P_val / e. So, y = W_0(P_val / e).
        # r/2M = 1 - y  => r = 2M * (1 - W_0(P_val / e))
        # Ensure argument to lambertw is within its domain [-1/e, inf)
        arg_lambert = P_val / np.e
        if arg_lambert < -1.0/np.e: # Clamp if slightly out due to precision
            arg_lambert = -1.0/np.e + 1e-12
        
        y = lambertw(arg_lambert, k=0).real # Principal branch W_0
        r = TWOM * (1.0 - y)
        return np.clip(r, 0, TWOM) # r should be in [0, 2M]

    # For r > 2M (Regions I, III), P_val is < 0. (1-r/2M) < 0.
    else: # P_val < 0
        # We solve (1-x)e^x = P_val where x = r/2M.
        # Let u = x-1 > 0. Then x = u+1.
        # -u * e^(u+1) = P_val => -u * e * e^u = P_val => u * e^u = -P_val / e.
        # u = W_0(-P_val / e).
        # r/2M = 1 + u => r = 2M * (1 + W_0(-P_val / e))
        arg_lambert = -P_val / np.e # This is > 0
        u = lambertw(arg_lambert, k=0).real
        r = TWOM * (1.0 + u)
        return np.clip(r, TWOM, None) # r should be >= 2M


# Helper function to calculate effective Schwarzschild t-like coordinate
def get_t_eff(U_k, V_k, region_id):
    # Avoid division by zero or log of non-positive for edge cases
    if math.isclose(U_k, 0.0) or math.isclose(V_k, 0.0):
        return 0.0 # On a horizon, t is not well-defined this way or is infinite

    # Kruskal T and R coordinates
    T_k = (U_k + V_k) / 2.0
    R_k = (V_k - U_k) / 2.0

    # Regions I, III (exterior, r > 2M, R_k is spacelike, T_k is timelike)
    # Schwarzschild t = 2*(2M) * arctanh(T_k/R_k)
    if region_id == 1 or region_id == 3:
        if math.isclose(R_k, 0.0) : return np.inf * np.sign(T_k) # Approximate
        ratio = T_k / R_k
        # Ratio should be in (-1, 1). Clamp for stability.
        ratio = np.clip(ratio, -1.0 + 1e-9, 1.0 - 1e-9)
        return 2.0 * TWOM * np.arctanh(ratio)

    # Regions II, IV (interior, r < 2M, T_k is spacelike, R_k is timelike)
    # Schwarzschild r (acting as time) = 2*(2M) * arctanh(R_k/T_k)
    # The 't_eff' here is this time-like Schwarzschild r.
    elif region_id == 2 or region_id == 4:
        if math.isclose(T_k, 0.0): return np.inf * np.sign(R_k) # Approximate
        ratio = R_k / T_k
        ratio = np.clip(ratio, -1.0 + 1e-9, 1.0 - 1e-9)
        return 2.0 * TWOM * np.arctanh(ratio)
    
    return 0.0 # Default

# Store pixel properties and sort keys
pixel_data = [] # List of dictionaries

print("Step 1: Calculating pixel properties and sort keys...")
for iy in range(HEIGHT):
    if iy % (max(1, HEIGHT // 20)) == 0: # Print progress roughly 20 times
        print(f"  Processing row {iy+1}/{HEIGHT}...")
    
    for ix in range(WIDTH):
        original_idx = iy * WIDTH + ix # Unique index for each pixel

        # Map pixel (ix, iy) to Penrose coordinates (rho, tau)
        # rho (horizontal) from -π/2 to π/2
        # tau (vertical, increases upwards) from -π/2 to π/2
        rho = (ix / (WIDTH - 1.0)) * np.pi - np.pi / 2.0
        tau = ((HEIGHT - 1.0 - iy) / (HEIGHT - 1.0)) * np.pi - np.pi / 2.0

        # Kruskal U, V from Penrose coordinates:
        # arctan(V_k) = tau + rho, arctan(U_k) = tau - rho
        arg_V = tau + rho
        arg_U = tau - rho
        
        # Epsilon for boundary checks (where tan arguments are ±π/2)
        eps_boundary_arg = 1e-7 
        is_scri = False
        if abs(abs(arg_U) - np.pi/2.0) < eps_boundary_arg or \
           abs(abs(arg_V) - np.pi/2.0) < eps_boundary_arg:
            is_scri = True

        # Clip arguments to tan to avoid infinity, but allow large values near boundaries
        # This makes U_k, V_k large but finite at scri, good enough for r, t calc.
        tan_clip_arg = np.pi/2.0 - eps_boundary_arg 
        U_k = np.tan(np.clip(arg_U, -tan_clip_arg, tan_clip_arg))
        V_k = np.tan(np.clip(arg_V, -tan_clip_arg, tan_clip_arg))

        # Schwarzschild r from -U_k * V_k = (1 - r/2M)e^(r/2M)
        P_val = -U_k * V_k
        r_val = -1.0 # Default to invalid r
        try:
            r_val = get_r_from_P(P_val)
        except (ValueError, OverflowError, ZeroDivisionError): 
            # Catch issues with LambertW or extreme inputs
            if P_val > 1.0 + 1e-6 : r_val = 0.0 # Effectively singularity if P too large
            elif is_scri : r_val = 1000 * TWOM # Large r for scri if P is negative infinite like
            # Other calculation failures might leave r_val = -1

        # Determine region ID based on signs of U_k, V_k (or arg_U, arg_V for robustness)
        # Region I (Exterior, right): V_k > 0, U_k < 0 => arg_V > 0, arg_U < 0 => rho > |tau|
        # Region II (Black Hole, top): V_k > 0, U_k > 0 => arg_V > 0, arg_U > 0 => tau > |rho|
        # Region III (Other Exterior, left): V_k < 0, U_k > 0 => arg_V < 0, arg_U > 0 => rho < -|tau|
        # Region IV (White Hole, bottom): V_k < 0, U_k < 0 => arg_V < 0, arg_U < 0 => tau < -|rho|
        region_id = 0 # Default/Ambiguous
        # Use small epsilon for comparing rho, tau magnitudes to define regions clearly
        eps_region_comp = 1e-6 
        if tau > abs(rho) - eps_region_comp: region_id = 2
        elif rho > abs(tau) - eps_region_comp: region_id = 1
        elif tau < -abs(rho) + eps_region_comp: region_id = 4
        elif rho < -abs(tau) + eps_region_comp: region_id = 3
        
        # Effective Schwarzschild t-like coordinate
        t_eff_val = 0.0
        if r_val >= 0: # Only if r is valid
            try:
                t_eff_val = get_t_eff(U_k, V_k, region_id)
            except (ValueError, OverflowError, ZeroDivisionError):
                 t_eff_val = np.inf if region_id in [1,3] else 0 # Default t based on region behavior

        # --- Identify features for sorting key ---
        is_sing = (r_val >= 0 and r_val < R_THRESH_SINGULARITY)
        is_hor = (r_val >= 0 and abs(r_val - TWOM) < R_THRESH_HORIZON)
        
        # t=const lines
        is_t_line = False
        t_line_idx = 0
        # Normalize t_eff to something periodic or bounded for line checks
        # For simplicity, check against discrete t_eff values if they are somewhat bounded by arctanh
        # Or, use angle in Kruskal plane: angle = np.arctan2(U_k + V_k, V_k - U_k) (for T_k, R_k)
        # This angle ranges -pi to pi.
        kruskal_T_k = (U_k + V_k) / 2.0
        kruskal_R_k = (V_k - U_k) / 2.0
        # Use different angles depending on region (where T or R is timelike)
        if region_id == 1 or region_id == 3: # T_k timelike (effectively)
            line_angle_rad = np.arctan2(kruskal_T_k, kruskal_R_k) if not (math.isclose(kruskal_T_k,0) and math.isclose(kruskal_R_k,0)) else 0
        else: # R_k timelike (effectively)
            line_angle_rad = np.arctan2(kruskal_R_k, kruskal_T_k) if not (math.isclose(kruskal_T_k,0) and math.isclose(kruskal_R_k,0)) else 0

        # Check if current angle is close to one of NUM_T_LINES angles
        target_angles = np.linspace(-np.pi, np.pi, NUM_T_LINES, endpoint=False)
        angle_diff_min = np.pi # Smallest difference to a target angle
        for i, target_ang in enumerate(target_angles):
            diff = abs(line_angle_rad - target_ang)
            diff = min(diff, 2*np.pi - diff) # Account for angle wrap-around
            if diff < angle_diff_min:
                angle_diff_min = diff
                t_line_idx = i
        
        if angle_diff_min < (np.pi / NUM_T_LINES) * T_LINE_WIDTH_PARAM * 5: # Width factor
            is_t_line = True

        # r=const lines
        is_r_line = False
        r_line_idx = 0
        if r_val >= 0: # Only if r is valid
            for i, r_center in enumerate(R_LINE_VALUES):
                if abs(r_val - r_center) < R_LINE_WIDTH_PARAM / 2.0:
                    is_r_line = True
                    r_line_idx = i
                    break
        
        # --- Construct sort key: (priority, param1, param2) ---
        # Lower priority value means "drawn earlier" / gets earlier block of colors.
        priority = 10.0 # Default fill
        param1 = 0.0    # Primary sort within priority group
        param2 = 0.0    # Secondary sort / tie-breaker

        if is_sing:
            priority = 0.0
            param1 = tau # Position along singularity (future vs past via tau)
        elif is_hor:
            priority = 1.0
            param1 = tau # Position along horizon (future vs past via tau)
        elif is_t_line:
            priority = 2.0
            param1 = t_line_idx # Which t-line
            param2 = r_val      # Position along the t-line (by r_val)
        elif is_r_line:
            priority = 3.0
            param1 = r_line_idx # Which r-line
            # Sort along r-line by angle (more consistent than t_eff which can be inf)
            param2 = line_angle_rad 
        elif is_scri:
            priority = 4.0
            # Distinguish scri segments, e.g. using which boundary it is.
            # (tau+rho) near +-pi/2, (tau-rho) near +-pi/2
            # Simplified: sort by tau, then rho.
            param1 = tau 
            param2 = rho
        else: # General region fill
            priority = 5.0 + region_id # Separate fill for each region
            # Shading within region based on r and t (or angle)
            # Example: make r variation primary, angle secondary
            param1 = r_val
            param2 = line_angle_rad
            
        sort_key_tuple = (priority, param1, param2)
        pixel_data.append({'key': sort_key_tuple, 'idx': original_idx})

print("Step 2: Sorting pixel data...")
# Python's sort is stable. If keys are identical, original order (by idx) is preserved.
# This is fine as idx itself is unique.
pixel_data.sort(key=lambda p: p['key'])

print("Step 3: Assigning color ranks...")
# pixel_ranks[original_pixel_idx] stores the final color rank for that pixel.
pixel_ranks = np.empty(WIDTH * HEIGHT, dtype=np.uint32)
for k_rank, p_info in enumerate(pixel_data):
    pixel_ranks[p_info['idx']] = k_rank

print("Step 4: Creating image from color ranks...")
img_array = np.empty((HEIGHT, WIDTH, 3), dtype=np.uint8) # HxWxRGB
total_pixels = WIDTH * HEIGHT

for iy_img in range(HEIGHT):
    if iy_img % (max(1, HEIGHT // 20)) == 0:
        print(f"  Generating image row {iy_img+1}/{HEIGHT}...")
    for ix_img in range(WIDTH):
        original_pixel_idx = iy_img * WIDTH + ix_img
        final_color_rank = pixel_ranks[original_pixel_idx]
        
        # Map final_color_rank (0 to 2^24-1) to unique R,G,B color
        r_byte = (final_color_rank >> 16) & 0xFF
        g_byte = (final_color_rank >> 8) & 0xFF
        b_byte = final_color_rank & 0xFF
        
        img_array[iy_img, ix_img, 0] = r_byte
        img_array[iy_img, ix_img, 1] = g_byte
        img_array[iy_img, ix_img, 2] = b_byte

final_image = Image.fromarray(img_array, 'RGB')
output_filename = "penrose_all_colors_4096.png"
final_image.save(output_filename)
print(f"Image saved as {output_filename}")

print("Done.")