sunset_simulator/skydome.py

#!/usr/bin/env python
# coding: utf-8

# Low pressure: 720mmHg
# High pressure: 780 mmHg

import numpy as np
import math
import imageio
import time
import random
import cython


"""
This class allows us to manipulate vectors
Not all of the functions are used yet - but they might be later when we start adding things
"""


#typed Vec3 class with appropriate functions, so we dont have to mess with ndarray[3] all the time
@cython.cclass
class Vec3:
    #cython specific declaration of the init and the v vactor
    v = cython.declare(cython.float[3],visibility="public")
    def __cinit__(self, x:cython.float =0, y:cython.float=0, z:cython.float=0):
        self.v:cython.float[3] = [x, y, z]

    def __init__(self, x:cython.float =0, y:cython.float=0, z:cython.float=0):
        self.v:cython.float[3] = [x, y, z]

    def __add__(self, other ):
        return Vec3(self.v[0] + other.v[0],self.v[1] + other.v[1],self.v[2] + other.v[2])

    def __sub__(self, other):
        return Vec3(self.v[0] - other.v[0],self.v[1] - other.v[1],self.v[2] - other.v[2])

    def __mul__(self, other):
        if isinstance(other, Vec3):
            return Vec3(self.v[0] * other.v[0],self.v[1] * other.v[1],self.v[2] * other.v[2])
        return Vec3(self.v[0] * other,self.v[1] * other,self.v[2] * other)

    def __truediv__(self, scalar):
        return Vec3(self.v[0] / scalar,self.v[1] / scalar,self.v[2] / scalar)

    def length(self)->cython.float:
        return math.sqrt(self.v[0]**2 + self.v[1]**2 + self.v[2]**2)

    def normalize(self):
        length = self.length()
        return self / length if length > 0 else Vec3()

    def to_tuple(self):
        return (self.v[0],self.v[1],self.v[2])

    def __repr__(self):
        return f"Vec3({self.v[0]}, {self.v[1]}, {self.v[2]})"

    @cython.cfunc
    def dot(self, other):
        return self.v[0]* other.v[0]+self.v[1]* other.v[1]+self.v[2]* other.v[2]


"""
Here we will be calculating what the value for n is based on T, P, but we will
also be taking dispersion into account - remember, we are considering 440 nm, 550 nm, 680 nm

These equations match the tabulated results made by Penndorf
"""

wavelengths = Vec3(0.68, 0.55, 0.44)  # These are the wavelengths in um

# The n_s values gives us the different refractive indices for the different wavelenghts
# The n values give us the new refractive indices depending on the temperature and pressure

@cython.cfunc
def refraction_calculator(T:cython.int, P:cython.int) -> tuple[Vec3, Vec3]:
    alpha:cython.float = 0.00367  # in inverse Celsius
    t_s:cython.int = 25  # Temperature in degrees Celsius
    p_s:cython.int = 760  # Pressure in mmHg


    n_s_values = []  # for debugging / comparing to the internet values
    n_values = []
    n_s:cython.double
    n:cython.double
    i:cython.int
    for i in [wavelengths.v[0], wavelengths.v[1], wavelengths.v[2]]:
        n_s = 1 + (0.05792105 / (238.0185 - i**-2) + 0.00167917 / (57.362 - i**-2))
        
        n = ((n_s - 1) * ((1 + alpha * t_s) / (1 + alpha * T)) * (P / p_s)) + 1
        
        n_s_values.append(n_s)
        n_values.append(n)

    return Vec3(*n_s_values), Vec3(*n_values)

@cython.cfunc
def beta(n_values:Vec3):
    
    N_s:cython.double = 2.54743e7  # um^-3
    
    beta_values = []

    for i, n in enumerate(n_values.v):
        wavelength = wavelengths.v[i]  # Corresponding wavelength value
        beta_1 = ((8 * math.pi * (n**2 - 1)**2) / ((2 * N_s) * (wavelength)**4))*1e6  # The whole thing should be in m^-1 to match with the rest of the code
       
        beta = (beta_1*10)
        beta_values.append(beta)

    return Vec3(beta_values[0],beta_values[1],beta_values[2])
    

# Example usage, normal is 25, 750
#t_s:cython.int = 25  # Temperature in degrees Celsius
#p_s:cython.int = 760  # Pressure in mmHg
#n_s_values:Vec3
#n_values:cython.double
#n_s_values, n_values = refraction_calculator(T, P)
#beta_values = beta(n_values)

#print("n_s values:", refraction_calculator_result.v)
#print("n values:", n_values.v)
#print("Beta values:", beta_values.v)


# We will have a scattering enhancement factor for our scattering values

kInfinity = float('inf')
angle_degrees:cython.int = 0  # Sun direction in degrees (0 to 90, noon to sunset)
angle_radians:cython.double = math.radians(angle_degrees)  # Convert to radians
#sunDir:Vec3 = Vec3(0, math.cos(angle_radians), -math.sin(angle_radians))

"""
The values for betaR depend on many things
For one, they scale proportionally to barometric pressure
It also varies according to temperature

"""

#betaR = Vec3(*beta_values.v)  # Keeps it as a Vec3
#betaR = Vec3(3.8e-6, 13.5e-6, 33.1e-6)

# R, G, B (the smaller the value, the less it will scatter during the day) in m-1
# 3.8 e-6, 13.5 e-6, 33.1 e-6 - default
# The order should be 680 nm, 550 nm, 440 nm


"""
There are many "categories" of aerosols and each has their own extinction coefficient (I removed the factor of 1.1)

Continental Clean - 26e-6
Continental Average - 75e-6
Continental Polluted - 175e-6

Urban - 353e-6

Desert - 145e-6

Maritime Clean - 90e-6
Maritime Polluted - 115e-6
Maritime Tropic - 43e-6

Arctic - 23e-6

Antarctic - 11e-6


"""
#betaM = Vec3(21e-6, 21e-6, 21e-6) # We need to have the Mie scattering be the same in all the directions
# The greater the value of beta, the smaller the Mie scattering point (responsable for the halo around the sun)
# If there is more pollution, we get a larger halo and the colors of the sunset become desaturated (more hazy)
# Default is 21e-6


earthRadius:cython.int = 6360e3 #m
atmosphereRadius:cython.int = 6420e3 #m (60 km higher than earth)
Hr:cython.int = 7994
Hm:cython.int = 1200

# The direction will change for each pixel
@cython.cfunc
def computeIncidentLight(direction:Vec3,betaM,betaR,g,observerEarthRadius,sunDir) -> tuple[float,float,float]:
    tmin:cython.float=0
    tmax:cython.float=kInfinity

    # We can change the origin position if we want, but for now it is 1 meter above the surface
    orig:Vec3 = Vec3(0, observerEarthRadius + 1, 0)
    t0:cython.float
    t1:cython.float
    t0, t1 = raySphereIntersect(orig, direction, atmosphereRadius)
    if t1 < 0:
        return Vec3(0, 0, 0)
    if t0 > tmin and t0 > 0:
        tmin = t0
    if t1 < tmax:
        tmax = t1

    numSamples:int = 16
    numSamplesLight:int = 8
    segmentLength = (tmax - tmin) / numSamples
    tCurrent = tmin
    sumR = Vec3()
    sumM = Vec3()
    opticalDepthR = 0
    opticalDepthM = 0
    mu:cython.float = direction.dot(sunDir) # This is the cosine of the angle between the direction vector (V) and the sun Direction

    phaseR = (3 * (1 + (mu * mu))) / (16 * math.pi)
    #g = 0.76
    phaseM = 3 / (8 * math.pi) * ((1 - g * g) * (1 + mu * mu)) / ((2 + g * g) * ((1 + g * g - 2 * g * mu) ** 1.5))

    i:cython.int
    for i in range(numSamples):
        #first ray in the ray walking algorithm
        samplePosition = orig + direction * (tCurrent + segmentLength * 0.5)
        height = samplePosition.length() - earthRadius
        hr = math.exp(-height / Hr) * segmentLength
        hm = math.exp(-height / Hm) * segmentLength
        opticalDepthR += hr
        opticalDepthM += hm

        # Sample position is the start, but it should be in the direction of the sun
        t0Light, t1Light = raySphereIntersect(samplePosition, sunDir, atmosphereRadius)

        if t1Light < 0:
            continue

        segmentLengthLight = (t1Light - t0Light) / numSamplesLight
        tCurrentLight = 0
        opticalDepthLightR = 0
        opticalDepthLightM = 0

        j:cython.int
        for j in range(numSamplesLight):
            #second ray in the ray walking
            samplePositionLight = samplePosition + (sunDir * (tCurrentLight + segmentLengthLight * 0.5))
            heightLight = samplePositionLight.length() - earthRadius
            if heightLight < 0:
                break
            opticalDepthLightR += (math.exp(-heightLight / Hr) * segmentLengthLight)
            opticalDepthLightM += (math.exp(-heightLight / Hm) * segmentLengthLight)
            tCurrentLight += segmentLengthLight

        if j == numSamplesLight - 1:
            tau = (betaR * (opticalDepthR + opticalDepthLightR)) + (betaM * (opticalDepthM + opticalDepthLightM))
            attenuation = Vec3(math.exp(-tau.v[0]), math.exp(-tau.v[1]), math.exp(-tau.v[2]))
            #separate contributions of Rayleigh and Mie contributions, with attentuation dependant on the BetaM/R, and altitude and pressure (through tau)
            sumR += (attenuation * hr)
            sumM += (attenuation * hm)
        tCurrent += (segmentLength)

    # Changing the * 20 number just changes the intensity os the light, it does not much change the colors themselves
    # Well the colors kinda change, but more they just get super pale or super not pale. 20 should be fine I guess
    #print("sumR: {}, betaR: {}, phaseR {}\n, sumM {},betaM {}, phaseM {}\n".format(sumR,betaR,phaseR,sumM,betaM,phaseM))
    final_color = (sumR * betaR * phaseR + sumM * betaM * phaseM) * 20 
    return final_color.to_tuple()

#intersection of a ray and a sphere
@cython.cfunc
def raySphereIntersect(orig:Vec3, direction:Vec3, radius:cython.double) -> tuple[cython.float,cython.float]:
    A:cython.float = direction.v[0]**2 + direction.v[1]**2 +direction.v[2]**2
    B:cython.float = 2 * (direction.v[0]*orig.v[0] +direction.v[1]*orig.v[1] +direction.v[2]*orig.v[2] )
    C:cython.float = orig.v[0]**2 + orig.v[1]**2 +orig.v[2]**2 - radius ** 2

    t0: cython.float
    t1: cython.float
    t0, t1 = solveQuadratic(A, B, C)
    if t1 < 0:
        return (kInfinity, kInfinity)
    return (t0, t1)

#quadratic equation solved.
@cython.cfunc
def solveQuadratic(a:cython.float, b:cython.float, c:cython.float) -> tuple[cython.float,cython.float]:
    if b == 0:
        if a == 0:
            return (0, 0)
        x1 = 0
        x2 = math.sqrt(-c / a)
        return (x1, x2)
    discr = b ** 2 - 4 * a * c
    if discr < 0:
        return (kInfinity, kInfinity)

    discr:cython.float
    q:cython.float
    discrs = math.sqrt(discr)
    q = (-0.5 * (b - discrs)) if b < 0 else (-0.5 * (b + discrs))
    return (q/a, c/q)



@cython.cfunc
def renderSkydome(filename,betaM,g,altitude,T,P,sunDir):
    width:cython.int
    height:cython.int
    width, height = 256,256
    image = np.zeros((height, width, 3), dtype=float)

    observerEarthRadius = earthRadius + altitude
    start_time = time.time()  # Start timing

    n_s_values, n_values = refraction_calculator(T, P)
    beta_values = beta(n_values)
    betaR = Vec3(*beta_values.v)  # Keeps it as a Vec3

    j:cython.int
    i:cython.int
    x:cython.float
    y:cython.float
    z2:cython.float
    phi:float
    theta:float
    direction:Vec3
    for j in range(height):
        for i in range(width):
            x = 2 * (i + 0.5) / (width - 1) - 1
            y = 2 * (j + 0.5) / (height - 1) - 1
            z2 = x * x + y * y
            if z2 <= 1:
                phi = math.atan2(y, x)
                theta = math.acos(1 - z2)
                # This changes for each pixel
                direction = Vec3(math.sin(theta) * math.cos(phi), math.cos(theta), math.sin(theta) * math.sin(phi))
                color = computeIncidentLight(direction,betaM,betaR,g,observerEarthRadius,sunDir)
                #color = computeIncidentLight(direction)


                # Assign the clipped color directly to the image array
                image[j][i][0] = np.clip(color, 0, 1)[0]
                image[j][i][1] = np.clip(color, 0, 1)[1]
                image[j][i][2] = np.clip(color, 0, 1)[2]

        # Print elapsed time after each row
        elapsed_time = time.time() - start_time
        print(f"Rendering row {j + 1}/{height}, elapsed time: {elapsed_time:.2f} seconds")
        #print(f"Rendering row {j + 1}/{height}")

    image = np.clip(image, 0, 1) * 255 # change 255
    return image.astype(np.uint8)



def renderFromCamera(filename,betaM,g,altitude,T,P,sunDir):

    print(wavelengths)
    n_s_values, n_values = refraction_calculator(T, P)
    beta_values = beta(n_values)
    betaR = Vec3(*beta_values.v)  # Keeps it as a Vec3

    observerEarthRadius = earthRadius + altitude
    width:cython.int
    height:cython.int
    width, height = 100, 100

    x:cython.int
    y:cython.int
    rayx:cython.float
    rayy:cython.float
    image = np.zeros((height, width, 3), dtype=np.float32)
    
    aspectRatio = width / float(height)
    fov = 65 # field of view
    angle = math.tan(fov * np.pi / 180 * 0.5) # How much the camera can see

    start_time = time.time()  # Start timing

    numPixelSamples:cython.int = 4
    for y in range(height):
        for x in range(width):
            if y > 70:
                # horrible horrible hack, mainly because below the horizon its dark amnyway?
                color = (np.nan,np.nan,np.nan)
                image[y][x] = color
            else:

                rayx = (2 * x  / float(width) - 1) * aspectRatio * angle
                rayy = (1 - y / float(height) * 2) * angle

                #This changes for each pixel, it is the direction we are looking in
                direction = Vec3(rayx, rayy, -1).normalize()
                color = computeIncidentLight(direction,betaM,betaR,g,observerEarthRadius,sunDir)
                image[y][x] = np.clip(color, 0, 1)

        elapsed_time = time.time() - start_time
        print(f"Renderinggg row {y + 1}/{height}, elapsed time: {elapsed_time:.2f} seconds")  

    image = np.clip(image, 0, 1) * 255
    return image.astype(np.uint8)