Cosmic Microwave Background v2

The successor of Cosmic Microwave Background v1. Made with Chatgpt 5, which yielded a much better result much more quickly.

#!/usr/bin/env python3
"""
Generate a 4096×4096 image of the Cosmic Microwave Background (CMB) that uses
EVERY 24-bit RGB color exactly once, using real Planck data (Planck 2018 PR3).

Colab-friendly version:
- No argparse (hardcoded settings below)
- Uses os.getcwd() instead of __file__
- Automatically downloads from working IRSA URLs with fallback options

Requires:
    pip install numpy healpy astropy pillow
"""

import os
import sys
import math
import time
import urllib.request
import numpy as np
from urllib.error import HTTPError
from PIL import Image
import healpy as hp
from astropy.io import fits

# ---------- Settings ----------
W = H = 4096
CHUNK_ROWS = 256

# Output file
OUT_FILE = "cmb_allrgb_4096.png"

# FITS file path (empty = auto-download)
FITS_PATH = ""  # leave blank for auto-download

# URLs (working IRSA paths, component-maps/cmb/)
BASE_URL = "https://irsa.ipac.caltech.edu/data/Planck/release_3/all-sky-maps/maps/component-maps/cmb/"
CANDIDATES = [
    "COM_CMB_IQU-smica-nosz_2048_R3.00_full.fits",  # preferred (no tSZ)
    "COM_CMB_IQU-smica_2048_R3.00_full.fits",       # with tSZ
    "COM_CMB_IQU-commander_2048_R3.00_full.fits",   # alternative
]

# ---------- Download & FITS Reading ----------
def download_with_fallback(local_path: str):
    for fname in CANDIDATES:
        url = BASE_URL + fname
        try:
            print(f"Attempting download: {url}")
            urllib.request.urlretrieve(url, local_path)
            print(f"Downloaded successfully: {local_path}")
            return local_path
        except HTTPError as e:
            print(f"Failed: {e}")
    raise RuntimeError("All download attempts failed.")

def detect_ordering_nest_from_fits(path: str, extname="I_STOKES") -> bool:
    with fits.open(path, memmap=True) as hdul:
        hdu = None
        for h in hdul:
            if (h.header.get("EXTNAME", "").strip().upper() == extname) or (h.name.strip().upper() == extname):
                hdu = h
                break
        if hdu is None:
            hdu = hdul[1] if len(hdul) > 1 else hdul[0]
        ordering = (hdu.header.get("ORDERING", "") or hdu.header.get("ORDER", "")).strip().upper()
        if ordering in ("NESTED", "NEST"):
            return True
        elif ordering in ("RING",):
            return False
        return False

def read_smica_temperature_map(path: str, force_nest: bool | None = None):
    if force_nest is None:
        is_nest = detect_ordering_nest_from_fits(path, extname="I_STOKES")
    else:
        is_nest = bool(force_nest)
    m = hp.read_map(path, field=0, nest=is_nest, dtype=np.float64, verbose=False)
    unseen = hp.UNSEEN
    if np.any(~np.isfinite(m)) or np.any(m == unseen):
        good = np.isfinite(m) & (m != unseen)
        med = np.median(m[good])
        m = m.copy()
        m[~good] = med
    return m, is_nest

# ---------- Sampling & Color Ordering ----------
def sample_cmb_to_square(map_1d: np.ndarray, is_nest: bool, w: int = W, h: int = H) -> np.ndarray:
    out = np.empty((h, w), dtype=np.float32)
    xs = (np.arange(w, dtype=np.float64) + 0.5) / w
    phi = xs * (2.0 * math.pi)
    for y0 in range(0, h, CHUNK_ROWS):
        y1 = min(y0 + CHUNK_ROWS, h)
        ys = (np.arange(y0, y1, dtype=np.float64) + 0.5) / h
        theta = ys[:, None] * math.pi
        phi_grid = np.broadcast_to(phi[None, :], (y1 - y0, w))
        theta_grid = np.broadcast_to(theta, (y1 - y0, w))
        vals = hp.get_interp_val(map_1d, theta_grid.ravel(), phi_grid.ravel(), nest=is_nest)
        out[y0:y1, :] = vals.reshape((y1 - y0, w)).astype(np.float32)
        print(f"Sampled rows {y0:4d}–{y1-1:4d} / {h}")
    return out

def build_color_order_indices(N: int) -> np.ndarray:
    c = np.arange(N, dtype=np.uint32)
    r = (c >> 16) & 0xFF
    g = (c >> 8)  & 0xFF
    b = c & 0xFF
    r_i = r.astype(np.int16)
    g_i = g.astype(np.int16)
    b_i = b.astype(np.int16)
    score_primary = (r_i - b_i).astype(np.int32)
    score_secondary = (r_i + g_i + b_i).astype(np.int32)
    score_tertiary = np.abs(g_i - r_i).astype(np.int32)
    idx = np.lexsort((score_secondary, score_tertiary, score_primary))
    return idx.astype(np.uint32)

def assign_allrgb_by_rank(vals2d: np.ndarray, color_order_idx: np.ndarray) -> np.ndarray:
    h, w = vals2d.shape
    N = h * w
    assert N == 256 * 256 * 256, "Image must have exactly 16,777,216 pixels."
    flat = vals2d.ravel()
    print("Sorting CMB values…")
    pix_order = np.argsort(flat, kind="mergesort").astype(np.uint32)
    c_sorted = color_order_idx
    out = np.empty((N, 3), dtype=np.uint8)
    r = (c_sorted >> 16).astype(np.uint8)
    g = ((c_sorted >> 8) & 0xFF).astype(np.uint8)
    b = (c_sorted & 0xFF).astype(np.uint8)
    out[pix_order, 0] = r
    out[pix_order, 1] = g
    out[pix_order, 2] = b
    return out.reshape((h, w, 3))

# ---------- Main ----------
def main():
    t0 = time.time()

    fits_path = FITS_PATH.strip()
    if not fits_path:
        fname = CANDIDATES[0]  # smica-nosz preferred
        fits_path = os.path.join(os.getcwd(), fname)
        if not os.path.exists(fits_path):
            download_with_fallback(fits_path)

    print("Reading Planck map…")
    cmb_map, is_nest = read_smica_temperature_map(fits_path)
    print(f"  nside={hp.get_nside(cmb_map)}, ordering={'NESTED' if is_nest else 'RING'}")

    print(f"Sampling to {W}×{H} grid (equirectangular)…")
    vals = sample_cmb_to_square(cmb_map, is_nest, W, H)
    vals -= np.mean(vals, dtype=np.float64)

    print("Preparing color ordering (all 16,777,216 colors)…")
    color_idx = build_color_order_indices(W * H)

    print("Assigning colors by CMB rank…")
    img = assign_allrgb_by_rank(vals, color_idx)

    Image.fromarray(img, mode="RGB").save(OUT_FILE, optimize=True)
    t1 = time.time()
    print(f"Done → {OUT_FILE}   ({(t1 - t0):.1f}s)")
    print("Pixels used:", img.shape[0] * img.shape[1])

if __name__ == "__main__":
    main()