path: root/Scripts/make_dfg_lut.py

#!/usr/bin/env python3

# To run with uv:
# uv run -w OpenEXR -w numba ./make_dfg_lut.py

import argparse
import math
import numpy as np
import OpenEXR
import Imath
import numba
import random
import concurrent.futures
import os
from functools import partial


@numba.njit(cache=True)
def rcp(a):
    return 1.0 / a


@numba.njit(cache=True)
def lerp(a, b, t):
    return a + (b - a) * t


@numba.njit(cache=True)
def saturate(a):
    if a < 0.0: return 0.0
    if a > 1.0: return 1.0
    return a


# Standard BRDF components.
@numba.njit(cache=True)
def F_Schlick(LoH, f0, f90=1.0):
    term = 1.0 - LoH
    term2 = term * term
    term5 = term2 * term2 * term
    return f0 + (f90 - f0) * term5


@numba.njit(cache=True)
def D_GGX(roughness, NoH):
    r2 = roughness * roughness
    NoH2 = NoH * NoH
    NoH4 = NoH2 * NoH2
    k = rcp(NoH2) - 1.0
    r2_plus_k = r2 + k
    denom = NoH4 * r2_plus_k * r2_plus_k
    return r2 / (denom + 1e-6)


@numba.njit(cache=True)
def G_GGXSmith(roughness, NoL, NoV):
    denom = 2.0 * lerp(2.0 * NoL * NoV, NoL + NoV, roughness)
    return rcp(denom + 1e-6)


# Cloth BRDF components.
@numba.njit(cache=True)
def D_Cloth(roughness, NoH):
    if roughness < 1e-4: return 0.0
    r_rcp = rcp(roughness)
    sin2H = 1.0 - NoH * NoH
    return (2.0 + r_rcp) * pow(sin2H, r_rcp * 0.5) / (2.0 * math.pi)


@numba.njit(cache=True)
def G_Cloth_L(x, a, b, c, d, e):
    return a / (1.0 + b * pow(x, c)) + d * x + e


@numba.njit(cache=True)
def Lambda_Cloth_Raw(roughness, cos_theta):
    a0, a1 = 25.3245, 21.5473
    b0, b1 = 3.32435, 3.82987
    c0, c1 = 0.16801, 0.19823
    d0, d1 = -1.27393, -1.97760
    e0, e1 = -4.85967, -4.32054

    one_minus_r = 1.0 - roughness
    interp = one_minus_r * one_minus_r
    rough_weight = 1.0 - interp

    lambda_val = 0.0
    if cos_theta < 0.5:
        L0 = G_Cloth_L(cos_theta, a0, b0, c0, d0, e0)
        L1 = G_Cloth_L(cos_theta, a1, b1, c1, d1, e1)
        L = lerp(L0, L1, rough_weight)
        lambda_val = math.exp(L)
    else:
        L_05_0 = G_Cloth_L(0.5, a0, b0, c0, d0, e0)
        L_05_1 = G_Cloth_L(0.5, a1, b1, c1, d1, e1)
        L_05 = lerp(L_05_0, L_05_1, rough_weight)

        one_minus_cos = 1.0 - cos_theta
        L_c_0 = G_Cloth_L(one_minus_cos, a0, b0, c0, d0, e0)
        L_c_1 = G_Cloth_L(one_minus_cos, a1, b1, c1, d1, e1)
        L_c = lerp(L_c_0, L_c_1, rough_weight)

        lambda_val = math.exp(2.0 * L_05 - L_c)

    return lambda_val


@numba.njit(cache=True)
def Lambda_Cloth_Softened(roughness, cos_theta):
    lambda_val = Lambda_Cloth_Raw(roughness, cos_theta)
    return pow(lambda_val, 1.0 + 2.0 * pow(1.0 - cos_theta, 8.0))


@numba.njit(cache=True)
def V_Cloth_Outgoing(roughness, NoL, NoV):
    # Height-correlated Smith: G2 / (4 * NoL * NoV)
    lambda_l = Lambda_Cloth_Softened(roughness, NoL)
    lambda_v = Lambda_Cloth_Raw(roughness, NoV)
    return 1.0 / ((1.0 + lambda_l + lambda_v) * 4.0 * NoL * NoV + 1e-6)


@numba.njit(cache=True)
def V_Cloth_Incoming(roughness, NoL, NoV):
    lambda_l = Lambda_Cloth_Softened(roughness, NoL)
    lambda_v = Lambda_Cloth_Raw(roughness, NoV)
    return 1.0 / ((1.0 + lambda_l + lambda_v) * 4.0 * NoL * NoV + 1e-6)


@numba.njit(cache=True)
def integrate_brdf_jitted(roughness, NoV, num_samples):
    V_x = math.sqrt(1.0 - NoV * NoV)
    V_y = 0.0
    V_z = NoV

    # R: GGX scale, G: GGX bias, B: cloth outgoing albedo, A: cloth incoming albedo
    std_scale, std_bias, cloth_out, cloth_in = 0.0, 0.0, 0.0, 0.0

    for i in range(num_samples):
        e1, e2 = random.random(), random.random()

        # Importance sample GGX
        a = roughness
        a2 = a * a

        phi = 2.0 * math.pi * e1
        cos_theta = math.sqrt((1.0 - e2) / (1.0 + (a2 - 1.0) * e2))
        sin_theta = math.sqrt(1.0 - cos_theta * cos_theta)

        H_x = math.cos(phi) * sin_theta
        H_y = math.sin(phi) * sin_theta
        H_z = cos_theta

        VoH = H_x * V_x + H_y * V_y + H_z * V_z
        if VoH <= 0: continue

        L_x = 2.0 * VoH * H_x - V_x
        L_y = 2.0 * VoH * H_y - V_y
        L_z = 2.0 * VoH * H_z - V_z

        NoL = saturate(L_z)
        NoH = saturate(H_z)
        NoV_proxy = saturate(V_z)

        if NoL > 0:
            # --- Standard BRDF ---
            # D cancels between numerator and PDF.
            G = G_GGXSmith(roughness, NoL, NoV_proxy)
            Fc_term = pow(1.0 - VoH, 5.0)

            # PDF = D_GGX * NoH / (4 * VoH), so (D * G * NoL) / PDF simplifies to:
            common_term = (G * NoL * 4.0 * VoH) / max(NoH, 1e-5)

            std_scale += common_term * (1.0 - Fc_term)
            std_bias += common_term * Fc_term

            # --- Cloth BRDF ---
            # Same GGX importance samples, reweighted for cloth D and V.
            if roughness >= 1e-4:
                D_c = D_Cloth(roughness, NoH)
                pdf_ggx = D_GGX(roughness, NoH) * NoH / (4.0 * VoH + 1e-6)
                V_out = V_Cloth_Outgoing(roughness, NoL, NoV_proxy)
                V_in = V_Cloth_Incoming(roughness, NoV_proxy, NoL)
                cloth_out += (D_c * V_out * NoL) / (pdf_ggx + 1e-6)
                cloth_in += (D_c * V_in * NoL) / (pdf_ggx + 1e-6)

    inv_n = 1.0 / num_samples
    return std_scale * inv_n, std_bias * inv_n, cloth_out * inv_n, cloth_in * inv_n


def calculate_pixel(coords, resolution, num_samples):
    x, y = coords
    u = (x + 0.5) / resolution
    v = (y + 0.5) / resolution

    NoV = saturate(u)
    perceptual_roughness = saturate(v)
    roughness = max(perceptual_roughness * perceptual_roughness, 1e-4)
    if NoV < 1e-4: return x, y, 0.0, 0.0, 0.0, 0.0

    std_scale, std_bias, cloth_out, cloth_in = integrate_brdf_jitted(roughness, NoV, num_samples)

    # R: GGX scale, G: GGX bias, B: cloth outgoing albedo, A: cloth incoming albedo
    return x, y, std_scale, std_bias, cloth_out, cloth_in


def generate_exr(resolution, output_filename, num_samples, num_workers):
    print(f"Generating {resolution}x{resolution} EXR '{output_filename}' (R=GGX scale, G=GGX bias, B=cloth out, A=cloth in) ({num_samples} samples/pixel) using {num_workers} workers.")
    header = OpenEXR.Header(resolution, resolution)
    pt = Imath.PixelType(Imath.PixelType.FLOAT)
    header['channels'] = {
        'R': Imath.Channel(pt),
        'G': Imath.Channel(pt),
        'B': Imath.Channel(pt),
        'A': Imath.Channel(pt),
    }

    pixel_data = np.zeros((resolution, resolution, 4), dtype=np.float32)

    coords_to_process = [(x, y) for y in range(resolution) for x in range(resolution)]
    worker_func = partial(calculate_pixel, resolution=resolution, num_samples=num_samples)

    processed_count = 0
    total_pixels = len(coords_to_process)
    print(f"Starting pixel processing...")

    with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:
        futures = {executor.submit(worker_func, coord): coord for coord in coords_to_process}

        for future in concurrent.futures.as_completed(futures):
            try:
                x, y, r, g, b, a = future.result()
                pixel_data[y, x] = (r, g, b, a)
            except Exception as exc:
                coord = futures[future]
                print(f'\nPixel at {coord} generated an exception: {exc}')

            processed_count += 1
            print(f"  ...processed {processed_count}/{total_pixels} pixels ({processed_count/total_pixels:.1%})", end='\r')

    print(f"\nProcessing complete. Writing to {output_filename}...")
    try:
        # Vertically flip to match UV coordinates (0,0 at bottom-left).
        pixel_data = np.flipud(pixel_data)

        exr_file = OpenEXR.OutputFile(output_filename, header)
        r_data = pixel_data[:, :, 0].ravel().tobytes()
        g_data = pixel_data[:, :, 1].ravel().tobytes()
        b_data = pixel_data[:, :, 2].ravel().tobytes()
        a_data = pixel_data[:, :, 3].ravel().tobytes()
        exr_file.writePixels({'R': r_data, 'G': g_data, 'B': b_data, 'A': a_data})
        exr_file.close()
        print(f"Successfully generated {output_filename}")
    except Exception as e:
        raise RuntimeError(f"Failed to write EXR file '{output_filename}': {e}")

def main():
    parser = argparse.ArgumentParser(description='Generate packed DFG LUT (R=GGX scale, G=GGX bias, B=cloth out, A=cloth in).')
    parser.add_argument('-r', '--resolution', type=int, default=512,
                        help='Resolution of the square EXR image (default: 512)')
    parser.add_argument('-s', '--samples', type=int, default=8192,
                        help='Number of samples per pixel for integration (default: 8192)')
    parser.add_argument('-o', '--output', default='dfg.exr',
                        help='Output filename (default: dfg.exr)')
    parser.add_argument('-j', '--workers', type=int, default=os.cpu_count(),
                        help=f'Number of worker processes (default: {os.cpu_count()})')

    args = parser.parse_args()

    if args.resolution <= 0:
        print("Error: Resolution must be a positive integer")
        return 1

    try:
        generate_exr(args.resolution, args.output, args.samples, args.workers)
    except Exception as e:
        print(f"Error: {e}")
        return 1

    return 0

if __name__ == '__main__':
    exit(main())