Absorption Spectra v2

Improved version of Absorption Spectra v1 Made with help from Chatgpt and Google Colab. Here are the Python script and sample output:

#!/usr/bin/env python3
"""
generate_spectra_with_labels.py

Generate a 4096×4096-pixel image of the absorption spectra for elements Z=1–128,
laid out in 4 columns × 32 rows, using every 24-bit RGB color exactly once.

Modifications:
1. Rainbow-style background for known-spectrum tiles goes from purple→red (hue rotated +60°).
2. 1-pixel “white” border around each tile and 2-pixel “white” border around the entire image,
   all drawn with the lightest available colors.
3. Each tile displays the chemical name (DejaVuSans, 32 px, no anti-aliasing) at its top-left,
   positioned 2 px from outer & tile borders, using the next-lightest available colors.
   Labels are drawn by masking, not textboxes, so no clipping or anti-aliasing.
4. The IUPAC systematic names are used for Z > 118.

- Absorption lines from NIST (380–780 nm) use the darkest colors (excluding borders & labels).
- Known-spectrum background: smooth rotated-hue gradient (violet→red) right-to-left.
- Unknown-spectrum background: flat gray.
- Every 24-bit RGB color is used exactly once, including borders and labels.

Dependencies:
    pip install astroquery astropy pillow numpy pandas
"""
import numpy as np
from PIL import Image, ImageFont
from astroquery.nist import Nist
from astropy import units as u
import pandas as pd

# --- 1. Image geometry ---
tile_w, tile_h = 1024, 128
ncols, nrows = 4, 32
width, height = tile_w * ncols, tile_h * nrows  # 4096 × 4096
total_pixels = width * height
print(f"Total pixels: {total_pixels}")

# Precompute flat indices, x/y coords, and tile IDs
indices = np.arange(total_pixels, dtype=np.int64)
y = indices // width
x = indices % width
tile_row = y // tile_h
tile_col = x // tile_w
tile_id = tile_row * ncols + tile_col + 1  # 1-based tile IDs

# --- 2. Element names Z=1..128 ---
element_names = [
    # Z=1–10
    "Hydrogen","Helium","Lithium","Beryllium","Boron",
    "Carbon","Nitrogen","Oxygen","Fluorine","Neon",
    # Z=11–20
    "Sodium","Magnesium","Aluminum","Silicon","Phosphorus",
    "Sulfur","Chlorine","Argon","Potassium","Calcium",
    # Z=21–30
    "Scandium","Titanium","Vanadium","Chromium","Manganese",
    "Iron","Cobalt","Nickel","Copper","Zinc",
    # Z=31–40
    "Gallium","Germanium","Arsenic","Selenium","Bromine",
    "Krypton","Rubidium","Strontium","Yttrium","Zirconium",
    # Z=41–50
    "Niobium","Molybdenum","Technetium","Ruthenium","Rhodium",
    "Palladium","Silver","Cadmium","Indium","Tin",
    # Z=51–60
    "Antimony","Tellurium","Iodine","Xenon","Cesium",
    "Barium","Lanthanum","Cerium","Praseodymium","Neodymium",
    # Z=61–70
    "Promethium","Samarium","Europium","Gadolinium","Terbium",
    "Dysprosium","Holmium","Erbium","Thulium","Ytterbium",
    # Z=71–80
    "Lutetium","Hafnium","Tantalum","Tungsten","Rhenium",
    "Osmium","Iridium","Platinum","Gold","Mercury",
    # Z=81–90
    "Thallium","Lead","Bismuth","Polonium","Astatine",
    "Radon","Francium","Radium","Actinium","Thorium",
    # Z=91–100
    "Protactinium","Uranium","Neptunium","Plutonium","Americium",
    "Curium","Berkelium","Californium","Einsteinium","Fermium",
    # Z=101–110
    "Mendelevium","Nobelium","Lawrencium","Rutherfordium","Dubnium",
    "Seaborgium","Bohrium","Hassium","Meitnerium","Darmstadtium",
    # Z=111–118
    "Roentgenium","Copernicium","Nihonium","Flerovium","Moscovium",
    "Livermorium","Tennessine","Oganesson",
    # Z=119–128 (IUPAC systematic temporary names)
    "Ununennium","Unbinilium","Unbiunium","Unbibium","Unbitrium",
    "Unbiquadium","Unbipentium","Unbihexium","Unbiseptium","Unbioctium"
]
assert len(element_names) == 128

# --- 3. Build border masks ---
# Outer border: 2 px wide at each edge
outer_mask = (x < 2) | (x >= width - 2) | (y < 2) | (y >= height - 2)

# Tile border: 1 px frame around each tile’s 1024×128 block
x_tile = x % tile_w
y_tile = y % tile_h
tile_border_mask = (
    (x_tile == 0) | (x_tile == tile_w - 1)
    | (y_tile == 0) | (y_tile == tile_h - 1)
)

# Combined border mask
border_mask = outer_mask | tile_border_mask
B_border = int(border_mask.sum())
print(f"Total border pixels (2px outer + 1px/tile): {B_border}")

# --- 4. Fetch NIST spectral lines (380–780 nm) for Z=1–118 ---
element_symbols = [
    'H','He','Li','Be','B','C','N','O','F','Ne',
    'Na','Mg','Al','Si','P','S','Cl','Ar','K','Ca',
    'Sc','Ti','V','Cr','Mn','Fe','Co','Ni','Cu','Zn',
    'Ga','Ge','As','Se','Br','Kr','Rb','Sr','Y','Zr',
    'Nb','Mo','Tc','Ru','Rh','Pd','Ag','Cd','In','Sn',
    'Sb','Te','I','Xe','Cs','Ba','La','Ce','Pr','Nd',
    'Pm','Sm','Eu','Gd','Tb','Dy','Ho','Er','Tm','Yb',
    'Lu','Hf','Ta','W','Re','Os','Ir','Pt','Au','Hg',
    'Tl','Pb','Bi','Po','At','Rn','Fr','Ra','Ac','Th',
    'Pa','U','Np','Pu','Am','Cm','Bk','Cf','Es','Fm',
    'Md','No','Lr','Rf','Db','Sg','Bh','Hs','Mt','Ds',
    'Rg','Cn','Nh','Fl','Mc','Lv','Ts','Og'
]
spectral_lines = {}
nist = Nist()
for Z, sym in enumerate(element_symbols, start=1):
    try:
        table = nist.query(380 * u.nm, 780 * u.nm, linename=sym)
        df = table.to_pandas() if table is not None else pd.DataFrame()
        if 'Observed' in df.columns:
            col = df['Observed']
        elif 'Ritz' in df.columns:
            col = df['Ritz']
        elif 'Wavelength' in df.columns:
            col = df['Wavelength']
        else:
            spectral_lines[Z] = []
            continue
        vals = pd.to_numeric(col, errors='coerce').dropna().astype(float)
        # Convert Å→nm if entries exceed 1000
        if len(vals) and vals.max() > 1000:
            vals /= 10.0
        spectral_lines[Z] = [w for w in vals if 380 <= w <= 780]
    except Exception as e:
        print(f"Warning Z={Z} ({sym}) fetch error: {e}")
        spectral_lines[Z] = []

known_ids = [Z for Z in range(1, 129) if spectral_lines.get(Z)]
unknown_ids = [Z for Z in range(1, 129) if Z not in known_ids]
print(f"Known spectra: {len(known_ids)}, Unknown: {len(unknown_ids)}")

# --- 5. Build label mask and list of label pixel indices ---
# Load DejaVuSans at 32 px (no anti-aliasing)
font = ImageFont.truetype("DejaVuSans.ttf", 32)

label_mask = np.zeros((total_pixels,), dtype=bool)
label_pixels_list = []

for Z in range(1, 129):
    name = element_names[Z - 1]
    mask = font.getmask(name, mode="1")
    w_text, h_text = mask.size  # getmask size returns (width, height)
    mask_arr = np.array(mask).reshape((h_text, w_text))  # 0 or 255

    # Position at 6 px from outer & tile borders: (x0 + 6, y0 + 6)
    r, c = divmod(Z - 1, ncols)
    y0, x0 = r * tile_h, c * tile_w
    for ty in range(h_text):
        for tx in range(w_text):
            if mask_arr[ty, tx]:  # pixel belongs to glyph
                global_x = x0 + 6 + tx
                global_y = y0 + 6 + ty
                pix_idx = global_y * width + global_x
                label_pixels_list.append(pix_idx)

label_flat = np.unique(np.array(label_pixels_list, dtype=np.int64))
B_labels = int(label_flat.size)
print(f"Total label pixels: {B_labels}")

label_mask[label_flat] = True

# --- 6. Identify raw absorption-line pixels (pre-filter) ---
raw_line_pixels = []
for Z in known_ids:
    r, c = divmod(Z - 1, ncols)
    y0, x0 = r * tile_h, c * tile_w
    for wl in spectral_lines[Z]:
        xi = x0 + int((wl - 380) / 400 * (tile_w - 1))
        ys = np.arange(y0, y0 + tile_h)
        raw_line_pixels.extend((ys * width + xi).tolist())

raw_line_flat = np.unique(np.array(raw_line_pixels, dtype=np.int64))
# Exclude any that fall on a border or a label
keep_line_mask = (~border_mask[raw_line_flat]) & (~label_mask[raw_line_flat])
line_flat = raw_line_flat[keep_line_mask]
M = int(line_flat.size)
print(f"Absorption-line pixels (excluding borders/labels): {M}")

# --- 7. Background masks (excluding borders, labels, and lines) ---
occupied = border_mask.copy()
occupied[line_flat] = True
occupied[label_flat] = True

mask_unk_bg = (~occupied) & np.isin(tile_id, unknown_ids)
G_count = int(mask_unk_bg.sum())

mask_known_bg = (~occupied) & np.isin(tile_id, known_ids)
R_count = int(mask_known_bg.sum())

assert (M + G_count + R_count + B_border + B_labels) == total_pixels, (
    f"Counts mismatch: M={M}, G={G_count}, R={R_count}, B_border={B_border}, "
    f"B_labels={B_labels}, sum={M+G_count+R_count+B_border+B_labels}, total={total_pixels}"
)
print(f"Counts → Lines: {M}, Gray bg: {G_count}, Rainbow bg: {R_count}, "
      f"Border: {B_border}, Labels: {B_labels}")

# --- 8. Sort all 24-bit colors by brightness ---
colors = np.arange(total_pixels, dtype=np.uint32)
R_vals = (colors >> 16) & 0xFF
G_vals = (colors >> 8) & 0xFF
B_vals = colors & 0xFF

brightness = 0.299 * R_vals + 0.587 * G_vals + 0.114 * B_vals
sorted_idx = np.argsort(brightness)  # darkest → lightest

# Reserve lightest B_border colors for borders
border_pool = sorted_idx[-B_border:]
remaining_after_border = sorted_idx[: total_pixels - B_border]

# Reserve next-lightest B_labels colors for labels
label_pool = remaining_after_border[-B_labels:]
remaining_after_border_labels = remaining_after_border[: total_pixels - B_border - B_labels]

# Next M darkest → absorption lines
line_pool = remaining_after_border_labels[:M]
remaining_after_lines = remaining_after_border_labels[M:]

# Build grayness metric for the rest (for flat gray background)
rem_R = (remaining_after_lines >> 16) & 0xFF
rem_G = (remaining_after_lines >> 8) & 0xFF
rem_B = remaining_after_lines & 0xFF
grayness = (
    np.abs(rem_R - rem_G).astype(np.int32)
    + np.abs(rem_G - rem_B).astype(np.int32)
    + np.abs(rem_R - rem_B).astype(np.int32)
)
g_idx = np.argsort(grayness)[:G_count]
gray_pool = remaining_after_lines[g_idx]

# The remainder → rainbow background
rainbow_pool = np.delete(remaining_after_lines, g_idx)
print(f"Pools sizes → Lines: {len(line_pool)}, Gray: {len(gray_pool)}, "
      f"Rainbow: {len(rainbow_pool)}, Border: {len(border_pool)}, Labels: {len(label_pool)}")

# --- 9. Hue-based sorting for rainbow_pool with +60° rotation ---
R3 = R_vals[rainbow_pool] / 255.0
G3 = G_vals[rainbow_pool] / 255.0
B3 = B_vals[rainbow_pool] / 255.0

maxc = np.maximum.reduce([R3, G3, B3])
minc = np.minimum.reduce([R3, G3, B3])
diff = maxc - minc

hue = np.zeros_like(maxc)
maskv = diff > 1e-6
rm = maskv & (maxc == R3)
gm = maskv & (maxc == G3)
bm = maskv & (maxc == B3)

hue[rm] = (60 * ((G3[rm] - B3[rm]) / diff[rm]) + 360) % 360
hue[gm] = (60 * ((B3[gm] - R3[gm]) / diff[gm]) + 120) % 360
hue[bm] = (60 * ((R3[bm] - G3[bm]) / diff[bm]) + 240) % 360

# Rotate by +60° → shift from violet→red
hue_rot = (hue + 60.0) % 360.0
hue_rot_norm = hue_rot / 360.0

hue_order = np.argsort(hue_rot_norm)
rainbow_sorted = rainbow_pool[hue_order]

# --- 10. Assemble final image ---
img_flat = np.zeros((total_pixels, 3), dtype=np.uint8)

# 10a. Draw borders (lightest colors)
border_positions = np.where(border_mask)[0]
border_positions.sort()
for i, pix in enumerate(border_positions):
    c = border_pool[i]
    img_flat[pix] = [(c >> 16) & 0xFF, (c >> 8) & 0xFF, c & 0xFF]

# 10b. Draw labels (next-lightest colors)
label_positions = label_flat.copy()
label_positions.sort()
for i, pix in enumerate(label_positions):
    c = label_pool[i]
    img_flat[pix] = [(c >> 16) & 0xFF, (c >> 8) & 0xFF, c & 0xFF]

# 10c. Draw absorption lines (darkest M colors)
line_positions = line_flat.copy()
line_positions.sort()
for i, pix in enumerate(line_positions):
    c = line_pool[i]
    img_flat[pix] = [(c >> 16) & 0xFF, (c >> 8) & 0xFF, c & 0xFF]

# 10d. Draw flat gray background for unknown-spectrum tiles
unk_positions = np.where(mask_unk_bg)[0]
unk_positions.sort()
for i, pix in enumerate(unk_positions):
    c = gray_pool[i]
    img_flat[pix] = [(c >> 16) & 0xFF, (c >> 8) & 0xFF, c & 0xFF]

# 10e. Draw rainbow background for known-spectrum tiles (right→left gradient)
k = len(known_ids)
n = len(rainbow_sorted)

# Compute per-element count of background pixels
tile_positions = {}
for idx, Z in enumerate(known_ids):
    count = int(np.count_nonzero(mask_known_bg & (tile_id == Z)))
    pos = (np.arange(count) * k + idx) % n
    tile_positions[Z] = pos.tolist()

# Resolve any duplicates among the hue indices
all_pos = np.concatenate([tile_positions[Z] for Z in known_ids])
unique_pos, counts = np.unique(all_pos, return_counts=True)
missing = list(set(np.arange(n)) - set(unique_pos))

def resolve_dupes(tile_positions_dict, counts_map, missing_list):
    for Z in known_ids:
        lst = tile_positions_dict[Z]
        for i, p in enumerate(lst):
            if counts_map[p] > 1 and missing_list:
                newp = missing_list.pop()
                counts_map[p] -= 1
                counts_map[newp] = 1
                lst[i] = newp
        if not missing_list:
            break

counts_map = {p: c for p, c in zip(unique_pos, counts)}
resolve_dupes(tile_positions, counts_map, missing)

# Assign actual rainbow colors per tile
for idx, Z in enumerate(known_ids):
    pos = np.array(tile_positions[Z], dtype=np.int64)
    cols = rainbow_sorted[pos]
    local_h = hue_rot_norm[hue_order][pos]
    grad_idx = np.argsort(local_h)
    cols_sorted = cols[grad_idx]
    cols_sorted = cols_sorted[::-1]  # reverse for right→left

    mask_k = mask_known_bg & (tile_id == Z)
    idxs = np.where(mask_k)[0]
    order_x = np.argsort(idxs % width)
    for j, pix in enumerate(idxs[order_x]):
        c = cols_sorted[j]
        img_flat[pix] = [(c >> 16) & 0xFF, (c >> 8) & 0xFF, c & 0xFF]

# --- 11. Save output ---
print("Saving image with labels and borders…")
img = Image.fromarray(img_flat.reshape((height, width, 3)), mode='RGB')
img.save('absorption_spectra_with_labels.png')
print("Done. Output file: absorption_spectra_with_labels.png")

Total pixels: 16777216
Total border pixels (2px outer + 1px/tile): 310636
Warning Z=85 (At) fetch error: Result did not contain a table
Warning Z=100 (Fm) fetch error: Result did not contain a table
Warning Z=101 (Md) fetch error: Result did not contain a table
Warning Z=102 (No) fetch error: Result did not contain a table
Warning Z=103 (Lr) fetch error: Result did not contain a table
Warning Z=104 (Rf) fetch error: Result did not contain a table
Warning Z=105 (Db) fetch error: Result did not contain a table
Warning Z=106 (Sg) fetch error: Result did not contain a table
Warning Z=107 (Bh) fetch error: Result did not contain a table
Warning Z=108 (Hs) fetch error: Result did not contain a table
Warning Z=109 (Mt) fetch error: Result did not contain a table
Warning Z=110 (Ds) fetch error: Result did not contain a table
Warning Z=111 (Rg) fetch error: Result did not contain a table
Warning Z=112 (Cn) fetch error: Result did not contain a table
Warning Z=113 (Nh) fetch error: Result did not contain a table
Warning Z=114 (Fl) fetch error: Result did not contain a table
Warning Z=115 (Mc) fetch error: Result did not contain a table
Warning Z=116 (Lv) fetch error: Result did not contain a table
Warning Z=117 (Ts) fetch error: Result did not contain a table
Warning Z=118 (Og) fetch error: Result did not contain a table
Known spectra: 98, Unknown: 30
Total label pixels: 136729
Absorption-line pixels (excluding borders/labels): 3054128
Counts → Lines: 3054128, Gray bg: 3818243, Rainbow bg: 9457480, Border: 310636, Labels: 136729
Pools sizes → Lines: 3054128, Gray: 3818243, Rainbow: 9457480, Border: 310636, Labels: 136729
Saving image with labels and borders…
Done. Output file: absorption_spectra_with_labels.png