Cubemap TBN Matrix

# written by gehtsiegarnixan

import numpy as np
import math
import datetime
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from pysr import PySRRegressor
import pandas as pd
from sympy import symbols, diff, sqrt


def rnd_normal_vectors(num_points=2 ** 16):
    # Generate random uv values
    uv_grid = np.random.rand(num_points, 2)

    # Initialize an array to store the output vectors
    vectors = []

    # Calculate the inf_cubemap for each uv in the grid
    for uv in uv_grid:
        vector = inf_cubemap(uv, correction=True)
        vectors.append(vector)

    return np.array(vectors)


def grid_normal_vectors(num_points):
    # Calculate the number of points along each dimension
    num_points_dim = int(np.sqrt(num_points))

    # Create a grid of uv values
    u = np.linspace(0, 1, num_points_dim)
    v = np.linspace(0, 1, num_points_dim)
    uv_grid = np.array(np.meshgrid(u, v)).T.reshape(-1, 2)

    # Initialize an array to store the output vectors
    vectors = []

    # Calculate the inf_cubemap for each uv in the grid
    for uv in uv_grid:
        vector = inf_cubemap(uv, correction=True)
        vectors.append(vector)

    return np.array(vectors)


def wsgrid_normal_vectors(num_samples=2 ** 16):
    # bounds of the plot
    v_max = math.sqrt(2) / 2

    # Create a grid of UV values
    uv_values = np.linspace(-v_max, v_max, num_samples)
    uv_grid = np.meshgrid(uv_values, uv_values)
    uv_grid = np.dstack(uv_grid)

    # Initialize an array to store RGB colors
    normals = []

    # Calculate sphere normals for each UV value and convert to RGB
    for row in uv_grid:
        for uv in row:
            # Clamp the value inside the square root to be no less than 0
            w_world_value = max(0.00001, -uv[0] ** 2 - uv[1] ** 2 + 1)
            w_world = math.sqrt(w_world_value)
            normal = np.array([uv[0], uv[1], w_world])

            # Append the normalized vector to the array
            normals.append(normal)

    return np.array(normals)


def cubemap(uvw, correction=True):
    # project into face
    cubemap_coord = uvw[:2] / uvw[2]

    if correction:
        # Cass Everitt's piecewise quadratic warp
        distort = 1.45109572583 - 0.451095725826 * np.abs(cubemap_coord)
        cubemap_coord *= distort

    # rescale to 0-1 range
    uv = cubemap_coord * 0.5 + 0.5

    return uv


def inf_cubemap(uv, correction=True):
    # rescale -1 to 1
    uv = uv * 2. - 1.

    if correction:
        # reverse Cass Everitt's piecewise quadratic warp
        arg_sqrt = 2.105679 - 1.804383 * np.abs(uv)
        arg_sqrt = np.maximum(arg_sqrt, 0)  # Ensure the argument is non-negative
        uv = np.sign(uv) * -1.108412 * (-1.451096 + np.sqrt(arg_sqrt))

    # recreate normal vector with side being up
    partial = np.sqrt((uv[0] ** 2) + (uv[1] ** 2) + 1.)
    normal = np.array([uv[0] / partial, uv[1] / partial, 1. / partial])

    return normal


def test_cubemap_inversion(num_samples=2 ** 16, correction=False):
    # Convert the list of vectors into a NumPy array
    input_values = rnd_normal_vectors(num_samples)

    for i in range(num_samples):
        uvw = input_values[i]

        # Apply cubemap and then inf_cubemap
        uv = cubemap(uvw, correction)
        uvw_reconstructed = inf_cubemap(uv, correction)

        # Check if the reconstructed values match the original input
        if not np.allclose(uvw_reconstructed, uvw, atol=1e-6):
            print("Test failed!")
            print("Original Input:", uvw)
            print("Reconstructed Input:", uvw_reconstructed)
            return

    print("All tests passed!")


def tangent_finite_difference(normal, correction=True):
    # cubemapping
    uv = cubemap(normal, correction)

    # inverse cubemapping with a tiny offset in x direction
    offset_uv = uv + np.array([0.00001, 0])
    normal_offset = inf_cubemap(offset_uv, correction)

    # generate the tangent vector approximation
    tangent_vec = normal_offset - normal
    tangent_vec /= np.linalg.norm(tangent_vec)  # normalize the tangent vector

    return tangent_vec


def tangent_derivative(normal, correction=True):
    # Generate the tangent vector approximation
    tangent_vec = np.array([1 + normal[1] ** 2,
                            - normal[0] * normal[1],
                            - normal[0]])
    # skipping the correction for this as it basically looks the same but is way more complicated.
    tangent_vec /= np.linalg.norm(tangent_vec)  # normalize the tangent vector

    return tangent_vec


def bitangent(normal, correction=True):
    # Generate the tangent vector approximation
    tangent_vec = np.array([- normal[0] * normal[1],
                            1 + normal[0] ** 2,
                            - normal[1]])
    tangent_vec /= np.linalg.norm(tangent_vec)  # normalize the tangent vector

    return tangent_vec


def z_plot(num_points_dim=255, correction=True):
    # Generate a grid of vectors
    vectors_array = wsgrid_normal_vectors(num_points_dim)

    # Initialize arrays to store RGB colors for both methods
    rgb_colors_tangent = []
    rgb_colors_approximation = []

    # Calculate sphere normals for each UV value and convert to RGB
    for uvw in vectors_array:
        uv = cubemap(uvw, correction=False)
        if np.all((0 <= uv) & (uv <= 1)):
            # Calculate tangent vector for cubemap using both methods
            tangent_vec = tangent_derivative(uvw, correction)
            approximation_vec = tangent_finite_difference(uvw, correction)

            rgb_colors_tangent.append(tangent_vec)
            rgb_colors_approximation.append(approximation_vec)
        else:
            rgb_colors_tangent.append([0, 0, 0])
            rgb_colors_approximation.append([0, 0, 0])

    # Convert the RGB colors to NumPy arrays
    rgb_colors_tangent = np.array(rgb_colors_tangent)
    rgb_colors_approximation = np.array(rgb_colors_approximation)

    # Clamp the RGB values to the range [0, 1]
    rgb_colors_tangent = np.clip(rgb_colors_tangent, 0, 1)
    rgb_colors_approximation = np.clip(rgb_colors_approximation, 0, 1)

    # Create side-by-side plots of the RGB colors with specified color range
    fig, axs = plt.subplots(1, 2, figsize=(16, 8))

    axs[0].imshow(rgb_colors_tangent.reshape(num_points_dim, num_points_dim, 3), origin='lower')
    axs[0].set_xlabel('X')
    axs[0].set_ylabel('Y')
    axs[0].set_title('Tangent from Derivative')
    axs[0].grid(True)

    axs[1].imshow(rgb_colors_approximation.reshape(num_points_dim, num_points_dim, 3), origin='lower')
    axs[1].set_xlabel('X')
    axs[1].set_ylabel('Y')
    axs[1].set_title('Tangent from finite Difference')
    axs[1].grid(True)

    plt.tight_layout()
    plt.show()


def y_plot(num_points_dim=255, correction=False, period=2):
    fig, ax = plt.subplots(figsize=(8, 8))

    # Generate a grid of vectors
    vectors_array = wsgrid_normal_vectors(num_points_dim)
    vectors_array = vectors_array.reshape(num_points_dim, num_points_dim, 3)

    def update(i):
        ax.clear()

        # flip flop animation
        index = abs(i - (num_points_dim - 1))

        # get the row of uvw coordinates
        vector_row = vectors_array[index]

        # Create arrays to store tangents for both methods
        tangents_true = []
        tangents_approximation = []
        vector_row_filtered = []

        # Iterate over uvw vectors and compute cubemap values
        for uvw_vector in vector_row:
            uv = cubemap(uvw_vector, correction=False)
            if np.all((0 <= uv) & (uv <= 1)):
                # Calculate tangent vector for cubemap using both methods
                tangent_vec_true = tangent_derivative(uvw_vector, correction)
                tangent_vec_approximation = tangent_finite_difference(uvw_vector, correction)

                tangents_true.append(tangent_vec_true)
                tangents_approximation.append(tangent_vec_approximation)

                # add the valid uvw vector to list
                vector_row_filtered.append(uvw_vector)

        # get y and z part of the coordinate vector
        x = [uvw_vector[0] for uvw_vector in vector_row_filtered]
        z = [uvw_vector[2] for uvw_vector in vector_row_filtered]

        # get x,y,z part of the tangent vector for both methods
        tangent_x_true = [tangent_vec[0] for tangent_vec in tangents_true]
        tangent_y_true = [tangent_vec[1] for tangent_vec in tangents_true]
        tangent_z_true = [tangent_vec[2] for tangent_vec in tangents_true]

        tangent_x_approximation = [tangent_vec[0] for tangent_vec in tangents_approximation]
        tangent_y_approximation = [tangent_vec[1] for tangent_vec in tangents_approximation]
        tangent_z_approximation = [tangent_vec[2] for tangent_vec in tangents_approximation]

        # Create the plot for true tangent
        ax.plot(x, tangent_x_true, label='Tangent x (Derivative)', color='orange')
        ax.plot(x, tangent_y_true, label='Tangent y (Derivative)', color='green')
        ax.plot(x, tangent_z_true, label='Tangent z (Derivative)', color='blue')

        # Create the plot for approximation tangent with dotted line style
        ax.plot(x, tangent_x_approximation, label='Tangent x (Finite Diff)', color='orange', linestyle='dotted')
        ax.plot(x, tangent_y_approximation, label='Tangent y (Finite Diff)', color='green', linestyle='dotted')
        ax.plot(x, tangent_z_approximation, label='Tangent z (Finite Diff)', color='blue', linestyle='dotted')

        # Create the plot for sphere
        ax.plot(x, z, label='Sphere', color='red')

        # Bounds of the plot
        ax.set_xlim([-0.9, 1.6])
        ax.set_ylim([-0.8, 1.4])

        # Add labels and legend
        ax.set_xlabel('X')
        ax.set_ylabel('Z')
        ax.set_title(f"Y = {vector_row[0][1]:.2f}")

        ax.grid(True)
        ax.legend()

    # animation
    frames = (num_points_dim - 1) * 2  # double for flipflop
    interval_ms = (period / num_points_dim) * 1000
    ani = animation.FuncAnimation(fig, update, frames=frames, repeat=True, interval=interval_ms)

    plt.axis('equal')
    plt.show()


def derivatives(correction=True):
    # Define the symbols
    x, y = symbols('x y')

    # Define the function q(x, y)
    q = (x / sqrt(x ** 2 + y ** 2 + 1),
         y / sqrt(x ** 2 + y ** 2 + 1),
         1 / sqrt(x ** 2 + y ** 2 + 1))

    if correction:
        # Cass Everitt's piecewise quadratic warp
        q = ((1.45109572583 - 0.451095725826 * sqrt(q[0]**2)) * q[0],
             (1.45109572583 - 0.451095725826 * sqrt(q[1]**2)) * q[1],
             (1.45109572583 - 0.451095725826 * sqrt(q[2]**2)) * q[2])

    # Compute the partial derivatives
    u = [diff(q_i, x) for q_i in q]  # ∂q/∂x
    v = [diff(q_i, y) for q_i in q]  # ∂q/∂y

    print("u_x =", u[0])
    print("u_y =", u[1])
    print("u_z =", u[2])

    print("v_x =", v[0])
    print("v_y =", v[1])
    print("v_z =", v[2])

    """
    the rescaling to 0-1 just adds a 0.5 factor for all components since we normalize this vector it cancels out.
    u = [(x²+1) /(x² + y² + 1)^(3/2),
          -xy   /(x² + y² + 1)^(3/2),
          -x    /(x² + y² + 1)^(3/2)]
    v = [-xy    /(x² + y² + 1)^(3/2),
        (x²+1)  /(x² + y² + 1)^(3/2),
          -y    /(x² + y² + 1)^(3/2)]
    can be simplified to:
    tangent = normalize(1+y^2, -xy, -x)
    bitangent = normalize(-xy, 1+x^2, -y)
    """


def symbolic_regression(num_samples=64*64, correction=True, iterations=5):
    # Generate some input data
    x = grid_normal_vectors(num_samples)

    # Initialize an array to store the output data
    y = []

    # Calculate the approximation tangent for each vector in X
    for uvw in x:
        tangent_vec = tangent_finite_difference(uvw, correction=correction)
        y.append(tangent_vec)

    # Convert the list of tangents into a NumPy array
    y = np.array(y)

    # Create a PySRRegressor object and fit the data
    model = PySRRegressor(niterations=iterations,
                          binary_operators=["plus", "sub", "mult", "div"],
                          unary_operators=["sqrt"],
                          )
    model.fit(x, y)

    # Set pandas options
    pd.set_option('display.max_columns', None)
    pd.set_option('display.expand_frame_repr', False)
    pd.set_option('max_colwidth', None)

    # Print the model
    print(model)

    # Format the timestamp as a string
    now = datetime.datetime.now()
    timestamp_str = now.strftime("%Y-%m-%d_%H%M%S.%f")
    filename = f'model_output_{timestamp_str}.txt'

    # Save the model's output to a text file
    with open(filename, 'w') as f:
        print(model, file=f)


def main():
    correction = False
    z_plot(255, correction)
    y_plot(128, correction)

    # test_cubemap_inversion(correction=correction)
    # derivatives(correction)
    # symbolic_regression(128*128, correction, 50)


if __name__ == "__main__":
    # execute only if run as a script
    main()