sunset_simulator/Midterm_submittion.py

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


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
"""

@cython.cclass
class Vec3:
    v = cython.declare(cython.float[3],visibility="public")
    def __cinit__(self, x:cython.double =0, y:cython.double=0, z:cython.double=0):
        self.v:cython.float[3] = [x, y, z]

    def __init__(self, x:cython.double =0, y:cython.double=0, z:cython.double=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],)
        elif isinstance(other, (int, float)):
            return Vec3(self.v[0] * other,self.v[1] * other,self.v[2] * other,)
        raise TypeError("Can only multiply by a Vec3 or scalar (int or float).")

    @cython.cfunc
    def __truediv__(self, scalar):
        if scalar == 0:
            raise ValueError("Cannot divide by zero.")
        return Vec3(*(self.v[i] / scalar for i in range(3)))

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

    @cython.cfunc
    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 = 15  # degrees Celsius
    p_s:cython.int = 760  # 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:cython.int = 5  # Temperature in degrees Celsius
P: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)


# In[7]:


"""
Here we are able to change the sun's direction - once the GUI is in place, this will be done from there
"""

kInfinity = float('inf')
angle_degrees:cython.int = 90  # 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 betaR values come from the previous cell
"""

#betaR = Vec3(3.8e-6, 13.5e-6, 33.1e-6), this is what the Nishita paper used
betaR = Vec3(*beta_values.v)  # Keeps it as a Vec3


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

Continental Clean - 26e-6     g = 0.709
Continental Average - 75e-6   g = 0.703
Continental Polluted - 175e-6 g = 0.698

Urban - 353e-6                g = 0.689

Desert - 145e-6               g = 0.729

Maritime Clean - 90e-6        g = 0.772
Maritime Polluted - 115e-6    g = 0.756
Maritime Tropic - 43e-6       g = 0.774

Arctic - 23e-6                g = 0.721

Antarctic - 11e-6             g = 0.784


"""
betaM = Vec3(11e-6, 11e-6, 11e-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

# For the value of betaM, we also need to change the anisotropy factor. This is in the same paper. 
g = 0.784

"""
The atmosphere model we use is as follows, with the radii and the scale heights
"""
earthRadius:cython.int = 6360e3 #m
atmosphereRadius:cython.int = 6420e3 #m (60 km higher than earth)
Hr:cython.int = 7994
Hm:cython.int = 1200


# In[8]:


"""
This is the compute incident light funciton from Nishita's paper
We translated it to Python and then are using cython to make it go faster

This is all explained in more detail in the submitted PDF
"""

@cython.cfunc
def computeIncidentLight(direction:Vec3) -> 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, earthRadius + 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

    """
    Phase functions - the anisotropy factor depends on the Mie scattering we have chosen, it is defined in the previous cell
    """
    phaseR = (3 * (1 + (mu * mu))) / (16 * math.pi)
    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):
        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):
            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]))
            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, this is the "magic number" Nishita used
    final_color = (sumR * betaR * phaseR + sumM * betaM * phaseM) * 20 
    return final_color.to_tuple()


# In[9]:


"""
These are complimentary functions that are used to solve the intensity of light. Again, from nishita's paper
"""

@cython.cfunc
def raySphereIntersect(orig:Vec3, direction:Vec3, radius:cython.double) -> tuple[cython.float,cython.float]:
    A:cython.double = direction.v[0]**2 + direction.v[1]**2 +direction.v[2]**2
    B:cython.double = 2 * (direction.v[0]*orig.v[0] +direction.v[1]*orig.v[1] +direction.v[2]*orig.v[2] )
    C:cython.double = 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)
    if t0 > t1:
        t0, t1 = t1, t0
    return (t0, t1)

@cython.cfunc
def solveQuadratic(a:cython.double, b:cython.double, c:cython.double) -> 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)

    q = (-0.5 * (b - np.sqrt(discr))) if b < 0 else (-0.5 * (b + np.sqrt(discr)))
    x1 = q / a
    x2 = c / q
    return (x1, x2)


# In[10]:


"""
Rendering Skydome image
"""

@cython.cfunc
def renderSkydome(filename):
    width:cython.int
    height:cython.int
    width, height = 256,256
    image = np.zeros((height, width, 3), dtype=float)

    start_time = time.time()  # Start timing

    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)
                #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}")

    # Save result to a PNG image
    image = np.clip(image, 0, 1) * 255 # change 255
    imageio.imwrite(filename, image.astype(np.uint8))



# In[11]:


"""
Rendering camera image
"""


def renderFromCamera(filename: str):
    width: cython.int = 252
    height: cython.int = 252
    image = np.zeros((height, width, 3), dtype=np.float32)

    aspectRatio: cython.float = width / float(height)
    fov: cython.float = 65.0  # Field of view
    angle: cython.float = np.tan(fov * np.pi / 180 * 0.5)  # Camera's view angle

    start_time = time.time()  # Start timing

    numPixelSamples: cython.int = 4
    y: cython.int
    x: cython.int
    m: cython.int
    n: cython.int
    rayx: cython.float
    rayy: cython.float
    direction: Vec3
    color: tuple  # Assuming computeIncidentLight returns a tuple

    for y in range(height):
        for x in range(width):
            sum_color = np.zeros(3, dtype=np.float32)  # Use numpy array for accumulation

            for m in range(numPixelSamples):
                for n in range(numPixelSamples):
                    # Compute ray direction with jitter
                    rayx = (2 * (x + (m + random.uniform(0, 1)) / numPixelSamples) / float(width) - 1) * aspectRatio * angle
                    rayy = (1 - (y + (n + random.uniform(0, 1)) / numPixelSamples) / float(height) * 2) * angle
                    
                    # Create the direction vector and normalize it
                    direction = Vec3(rayx, rayy, -1).normalize()
                    color = computeIncidentLight(direction)  # Assuming this returns a tuple

                    # Accumulate color values
                    sum_color[0] += color[0]
                    sum_color[1] += color[1]
                    sum_color[2] += color[2]

            # Average the accumulated color and clip to [0, 1]
            image[y, x] = np.clip(sum_color / (numPixelSamples * numPixelSamples), 0, 1)

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

    # Save the result to a PNG image
    image = np.clip(image, 0, 1) * 255
    imageio.imwrite(filename, image.astype(np.uint8))

     


# In[12]:


"""
Testing
"""

renderSkydome("test.png")


# In[ ]: