3D projection of Antarctica with original claimants of the region

south_pole_3d.py

# %%
"""
Antarctic Territorial Claims — 3D Relief Map

Renders a tilted 3D disk of Antarctica using real ETOPO 2022 elevation
data from NOAA, with territorial claim overlays and country flags.

Data source: ETOPO 2022 v1 (60-arc-second surface elevation)
  https://www.ngdc.noaa.gov/mgg/global/global.html
"""

import cartopy.crs as ccrs
import matplotlib.path as mpath
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LightSource
from matplotlib.offsetbox import AnnotationBbox, OffsetImage
from mpl_toolkits.mplot3d import proj3d
from PIL import Image, ImageDraw, ImageFont, ImageOps
from scipy.interpolate import RegularGridInterpolator
from scipy.ndimage import gaussian_filter

# ── Configuration ─────────────────────────────────────────────────────

ETOPO_PATH = "/tmp/etopo_antarctic_hires.npz"
FLAG_DIR = "/tmp/flags"
OUTPUT = "south_pole_3d.png"

# Antarctic territorial claims: (lon_start, lon_end, color, country_code, label)
# Use None for label when a country has a split territory (e.g. Australia)
TERRITORIES = [
    (-74, -25, "#4e79a7", "ar", "Argentina"),
    (45, 136, "#f28e2b", "au", "Australia"),
    (142, 160, "#f28e2b", "au", None),  # Australia's eastern sector
    (-90, -53, "#e15759", "cl", "Chile"),
    (136, 142, "#b07aa1", "fr", "France"),
    (160, 210, "#59a14f", "nz", "New Zealand"),
    (-20, 45, "#76b7b2", "no", "Norway"),
    (-80, -20, "#ff9da7", "gb", "United Kingdom"),
]

# Override the default midpoint angle for overlapping claims,
# and control how high + far out each flag floats
FLAG_LAYOUT = {
    "ar": {"angle": -30, "z": 0.35, "r": 0.95},
    "gb": {"angle": -55, "z": 0.35, "r": 0.95},
    "cl": {"angle": -80, "z": 0.35, "r": 0.95},
    "no": {"angle": None, "z": 0.35, "r": 0.95},
    "au": {"angle": None, "z": 0.35, "r": 0.95},
    "fr": {"angle": None, "z": 0.35, "r": 0.95},
    "nz": {"angle": None, "z": 0.35, "r": 0.95},
}

# 3D surface parameters
GRID_RADIAL = 350  # radial resolution
GRID_ANGULAR = 500  # angular resolution
HEIGHT_EXAGGERATION = 0.12  # ~50x real scale
VIEW_ELEVATION = 40
VIEW_AZIMUTH = -70

PLATE_CARREE = ccrs.PlateCarree()


def render_polar_map(territories: list) -> np.ndarray:
    """Render a flat South Polar Stereographic map with territory overlays
    and return it as an RGBA image array."""
    proj = ccrs.SouthPolarStereo()
    fig = plt.figure(figsize=(10, 10), dpi=150)
    ax = fig.add_axes([0, 0, 1, 1], projection=proj)
    ax.set_extent([-180, 180, -90, -50], PLATE_CARREE)

    # Circular boundary
    theta = np.linspace(0, 2 * np.pi, 100)
    verts = np.vstack([np.sin(theta), np.cos(theta)]).T
    ax.set_boundary(mpath.Path(verts * 0.5 + 0.5), transform=ax.transAxes)

    ax.stock_img()
    ax.coastlines(lw=0.8, color="#3a3a3a", zorder=2)
    ax.gridlines(draw_labels=False, color="white", alpha=0.15, lw=0.5)
    ax.patch.set_visible(False)
    ax.set_facecolor("none")
    fig.set_facecolor("none")

    for lon0, lon1, color, _, _ in territories:
        draw_territory_wedge(ax, lon0, lon1, color)

    fig.canvas.draw()
    img = np.array(fig.canvas.renderer.buffer_rgba())
    plt.close(fig)
    return img


def draw_territory_wedge(
    ax,
    lon_start: float,
    lon_end: float,
    color: str,
    lat_inner: float = -89.5,
    lat_outer: float = -60,
    alpha: float = 0.2,
):
    """Draw a filled wedge between two longitudes from the pole to 60°S."""
    n_points = max(int(abs(lon_end - lon_start)) * 2, 10)
    lons_arc = np.linspace(lon_start, lon_end, n_points)
    lons = np.concatenate([lons_arc, lons_arc[::-1], [lons_arc[0]]])
    lats = np.concatenate(
        [
            np.full_like(lons_arc, lat_outer),
            np.full_like(lons_arc, lat_inner),
            [lat_outer],
        ]
    )
    ax.fill(lons, lats, transform=PLATE_CARREE, color=color, alpha=alpha, zorder=3)


def build_elevation_surface(etopo_path: str, n_r: int, n_t: int):
    """Load ETOPO data and interpolate it onto a polar disk grid.

    Returns X, Y, Z_scaled (disk coordinates), R (radius grid),
    and the raw grid dimensions for later sampling.
    """
    data = np.load(etopo_path)
    elev, lat, lon = data["z"], data["lat"], data["lon"]

    r_vals = np.linspace(0, 1, n_r)
    t_vals = np.linspace(0, 2 * np.pi, n_t)
    R, T = np.meshgrid(r_vals, t_vals)

    X = R * np.cos(T)
    Y = R * np.sin(T)

    # Map polar grid → geographic coordinates
    # r=0 is the South Pole (-90°), r=1 is the disk edge (-50°)
    # Math angle convention: +X is 0°, +Y is 90°
    # In SouthPolarStereo: +Y points toward 0° longitude
    # So geographic longitude = 90° - math_angle_degrees
    lat_grid = -90 + R * 40
    lon_grid = ((90 - np.rad2deg(T) + 180) % 360) - 180

    interp = RegularGridInterpolator(
        (lat, lon),
        elev,
        method="linear",
        bounds_error=False,
        fill_value=0,
    )
    Z = interp(np.column_stack([lat_grid.ravel(), lon_grid.ravel()])).reshape(R.shape)

    # Flatten the ocean and apply light smoothing
    Z = gaussian_filter(np.maximum(Z, 0), sigma=1)

    # Scale height and fade at edges for a clean boundary
    z_max = max(Z.max(), 1)
    Z_scaled = Z / z_max * HEIGHT_EXAGGERATION
    Z_scaled *= np.clip((1 - R) / 0.05, 0, 1)

    return X, Y, Z_scaled, R, n_r, n_t


def sample_map_texture(
    map_img: np.ndarray, X: np.ndarray, Y: np.ndarray, R: np.ndarray
):
    """Sample the rendered polar map image at each surface point."""
    h, w = map_img.shape[:2]
    px = ((X + 1) / 2 * (w - 1)).astype(int).clip(0, w - 1)
    py = ((1 - (Y + 1) / 2) * (h - 1)).astype(int).clip(0, h - 1)
    colors = map_img[py, px, :3] / 255.0
    colors[R > 0.98] = 1.0  # white outside the disk
    return colors


def longitude_to_disk(lon_deg: float, radius: float) -> tuple[float, float]:
    """Convert geographic longitude + radius to X, Y on the polar disk."""
    angle = np.deg2rad(90 - lon_deg)
    return radius * np.cos(angle), radius * np.sin(angle)


def sample_surface_z(
    Z_scaled, lon_deg: float, radius: float, n_r: int, n_t: int
) -> float:
    """Look up the surface height at a given longitude and radius."""
    angle = np.deg2rad(90 - lon_deg) % (2 * np.pi)
    ri = int(radius * (n_r - 1))
    ti = int(angle / (2 * np.pi) * (n_t - 1)) % n_t
    return float(Z_scaled[ti, ri])


def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
    """Parse '#rrggbb' to (r, g, b) integers."""
    return int(hex_color[1:3], 16), int(hex_color[3:5], 16), int(hex_color[5:7], 16)


def build_flag_badge(flag_code: str, label: str, color: str) -> Image.Image:
    """Create a combined image with a colored label banner above the flag."""
    flag = Image.open(f"{FLAG_DIR}/{flag_code}_hires.png").convert("RGBA")
    flag = ImageOps.expand(flag, border=4, fill="white")
    fw, fh = flag.size

    # Measure text
    try:
        font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 52)
    except (OSError, IOError):
        font = ImageFont.load_default()

    tmp = ImageDraw.Draw(Image.new("RGBA", (1, 1)))
    text_box = tmp.textbbox((0, 0), label, font=font)
    tw = text_box[2] - text_box[0]
    th = text_box[3] - text_box[1]

    # Size the banner to fit the text or the flag width, whichever is larger
    banner_w = max(fw, tw + 40)
    banner_h = th + 30
    gap = 8

    # Assemble: banner on top, flag below
    combined = Image.new("RGBA", (banner_w, fh + banner_h + gap), (0, 0, 0, 0))

    r, g, b = hex_to_rgb(color)
    banner = Image.new("RGBA", (banner_w, banner_h), (r, g, b, 230))
    draw = ImageDraw.Draw(banner)
    draw.text(
        ((banner_w - tw) // 2, (banner_h - th) // 2 - 2),
        label,
        fill=(255, 255, 255, 255),
        font=font,
    )
    combined.paste(banner, (0, 0), banner)
    combined.paste(flag, ((banner_w - fw) // 2, banner_h + gap), flag)

    # Drop shadow
    shadow_px = 2
    final = Image.new(
        "RGBA",
        (combined.width + shadow_px * 2, combined.height + shadow_px * 2),
        (0, 0, 0, 0),
    )
    shadow = Image.new("RGBA", combined.size, (0, 0, 0, 35))
    final.paste(shadow, (shadow_px * 2, shadow_px * 2))
    final.paste(combined, (0, 0), combined)
    return final


def nearest_edge_alignment(flag_fig, anchor_fig) -> tuple[float, float]:
    """Return the box_alignment that places the connector at the
    nearest edge center (left, right, top, or bottom) of the flag badge."""
    dx = flag_fig[0] - anchor_fig[0]
    dy = flag_fig[1] - anchor_fig[1]
    if abs(dx) > abs(dy):
        return (0.0, 0.5) if dx > 0 else (1.0, 0.5)
    return (0.5, 0.0) if dy > 0 else (0.5, 1.0)


def project_3d_to_fig(ax3d, fig, x, y, z):
    """Project a 3D point to figure-fraction coordinates."""
    x2, y2, _ = proj3d.proj_transform(x, y, z, ax3d.get_proj())
    disp = ax3d.transData.transform((x2, y2))
    return fig.transFigure.inverted().transform(disp)


# Step 1: Render the flat polar map as a texture image
map_img = render_polar_map(TERRITORIES)

# Step 2: Build the 3D elevation surface
X, Y, Z_s, R, n_r, n_t = build_elevation_surface(ETOPO_PATH, GRID_RADIAL, GRID_ANGULAR)
facecolors = sample_map_texture(map_img, X, Y, R)

# Step 3: Plot the 3D surface
fig = plt.figure(figsize=(13, 11), facecolor="white")
ax3 = fig.add_axes([-0.05, -0.05, 1.1, 1.1], projection="3d")
ax3.set_facecolor("white")
ax3.view_init(elev=VIEW_ELEVATION, azim=VIEW_AZIMUTH)

ax3.plot_surface(
    X,
    Y,
    Z_s,
    facecolors=facecolors,
    rstride=1,
    cstride=1,
    shade=True,
    lightsource=LightSource(azdeg=315, altdeg=40),
    antialiased=False,
)

ax3.set_xlim(-1.2, 1.2)
ax3.set_ylim(-1.2, 1.2)
ax3.set_zlim(-0.2, 1.0)
ax3.set_box_aspect([1, 1, 0.5])
ax3.set_axis_off()

# Step 4: Add flag badges with connector sticks
fig.canvas.draw()  # needed for accurate 3D → 2D projection

placed = set()
for lon0, lon1, color, code, label in TERRITORIES:
    if label is None or code in placed:
        continue
    placed.add(code)

    layout = FLAG_LAYOUT.get(code, {})
    mid_lon = layout.get("angle") or (lon0 + lon1) / 2
    z_offset = layout.get("z", 0.35)
    r_flag = layout.get("r", 0.95)

    # Anchor: where the stick meets the surface (near the 60°S circle)
    r_anchor = 0.75
    xa, ya = longitude_to_disk(mid_lon, r_anchor)
    za = sample_surface_z(Z_s, mid_lon, r_anchor, n_r, n_t)

    # Flag position: further out and floating above
    xf, yf = longitude_to_disk(mid_lon, r_flag)
    zf = za + z_offset

    # Draw connector stick and anchor dot in 3D
    ax3.plot([xa, xf], [ya, yf], [za, zf], color=color, lw=1.5, alpha=0.7, zorder=10)
    ax3.scatter([xa], [ya], [za], color=color, s=20, zorder=11)

    # Project flag position to 2D figure coordinates
    flag_fig = project_3d_to_fig(ax3, fig, xf, yf, zf)
    flag_fig = np.clip(flag_fig, [0.08, 0.06], [0.92, 0.94])

    # Align the badge so the stick touches its nearest edge
    anchor_fig = project_3d_to_fig(ax3, fig, xa, ya, za)
    alignment = nearest_edge_alignment(flag_fig, anchor_fig)

    badge = build_flag_badge(code, label, color)
    ab = AnnotationBbox(
        OffsetImage(np.array(badge), zoom=0.12),
        flag_fig,
        xycoords="figure fraction",
        frameon=False,
        box_alignment=alignment,
    )
    fig.add_artist(ab)

fig.savefig(OUTPUT, dpi=200, transparent=True, bbox_inches="tight")
print(f"Saved {OUTPUT}")