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.")
