potato-universe.py

#!/usr/bin/env python3
import numpy as np
from PIL import Image, ImageDraw
import plotly.graph_objects as go


def molleweide_to_equirectagular():
    """
    converts Mollweide projection image to equirectangular projection
    """
    mollweide = Image.open("universe-mollweide.png")

    width, height = mollweide.size
    equi = Image.new("RGBA", (width, height), "black")

    mol_pixels = mollweide.load()
    equi_pixels = equi.load()

    epsilon = 0.000000001
    sqrt2 = np.sqrt(2)

    def aux_angle(latitude):
        """
        solves Newton-Raphson to compute auxillary angle for Mollweide projection
        """
        if np.abs(latitude - np.pi / 2) < epsilon:
            return np.pi / 2
        elif np.abs(latitude + np.pi / 2) < epsilon:
            return -np.pi / 2

        aux = latitude
        aux_prev = np.inf
        while np.abs(aux - aux_prev) > epsilon:
            aux_prev = aux
            aux = aux - (2 * aux + np.sin(2 * aux) - np.pi * np.sin(latitude)) / (
                2 + 2 * np.cos(2 * aux)
            )
        return aux

    for j in range(0, height):
        latitude = (np.pi) * (j / height - 0.5)
        aux = aux_angle(latitude)
        for i in range(0, width):
            longitude = np.pi * (i / width - 0.5)

            x = (2 * sqrt2 * (longitude - 0) * np.cos(aux)) / np.pi
            y = sqrt2 * np.sin(aux)

            x = (x / sqrt2 + 1) * height
            y = (y / sqrt2 + 1) * height / 2

            try:
                equi_pixels[i, j] = mol_pixels[np.round(x), np.round(y)]
            except IndexError:
                print("index out of bounds", i, j)
                print("index out of bounds", x, y)

    equi.show()
    equi.save("universe-equi-rect.png", "PNG")


# molleweide_to_equirectagular()

scale = Image.open("universe-scale.png")
scale_pixels = scale.load()


def get_value(pixel):
    u = np.array(pixel)
    min_norm = 256
    min_idx = 0
    for i in range(scale.size[1]):
        v = np.array(scale_pixels[5, i])
        w = u - v
        norm = np.sqrt(w.dot(w))
        if norm < min_norm:
            min_norm = norm
            min_idx = i
    return min_idx / scale.size[1]


img = Image.open("universe-equi-rect.png")
img_pixels = img.load()
width, height = img.size

rows = 50
cols = 50
vertices = (rows + 1) * (cols + 1)

Xs = [0] * vertices
Ys = [0] * vertices
Zs = [0] * vertices
intensities = [0] * vertices
Is = []
Js = []
Ks = []

for i in range(vertices):
    col = i % (cols + 1)
    row = i // (cols + 1)

    az = 2 * np.pi * col / cols
    alt = np.pi * row / rows

    if row == 0 or row == rows:
        pixel = img_pixels[width / 2, height * row / (rows + 1)]
    elif col == cols:
        pixel = img_pixels[0, height * row / (rows + 1)]
    else:
        pixel = img_pixels[(width * col / cols) % width, height * row / (rows + 1)]

    val = get_value(pixel)
    intensities[i] = val

    r = 1 + 2 * val

    Xs[i] = r * np.sin(alt) * np.cos(az)
    Ys[i] = r * np.sin(alt) * np.sin(az)
    Zs[i] = r * np.cos(alt)

    if i // (cols + 1) < cols:
        if i % (cols + 1) < cols:
            Is.append(i)
            Js.append(i + 1)
            Ks.append(i + cols + 1)
        if i % (cols + 1) > 0:
            Is.append(i)
            Js.append(i + cols + 1)
            Ks.append(i + cols)


fig = go.Figure(
    data=[
        go.Mesh3d(
            x=Xs,
            y=Ys,
            z=Zs,
            colorscale=[[0, "indigo"], [0.5, "red"], [1, "yellow"]],
            intensity=intensities,
            i=Is,
            j=Js,
            k=Ks,
            showscale=True,
        )
    ]
)

fig.show()
# fig.write_html("potato-universe.html")