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 G_Cloth(roughness, LoH):
    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
    one_minus_r_sq = one_minus_r * one_minus_r
    one_minus_LoH = 1.0 - LoH

    lambda_val = 0.0
    if LoH < 0.5:
        L0 = G_Cloth_L(LoH, a0, b0, c0, d0, e0)
        L1 = G_Cloth_L(LoH, a1, b1, c1, d1, e1)
        L = lerp(L0, L1, one_minus_r_sq)
        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, one_minus_r_sq)

        L_LoH_0 = G_Cloth_L(one_minus_LoH, a0, b0, c0, d0, e0)
        L_LoH_1 = G_Cloth_L(one_minus_LoH, a1, b1, c1, d1, e1)
        L_LoH = lerp(L_LoH_0, L_LoH_1, one_minus_r_sq)

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

    # Apply terminator softening (equation 4)
    return pow(lambda_val, 1.0 + 2.0 * pow(one_minus_LoH, 8.0))


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

    A, B = 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) # NoV is V_z

        if NoL > 0:
            scale, bias = 0.0, 0.0
            # --- Standard BRDF ---
            if brdf_type == 1:
                # Note that the D term is present in the numerator and the denominator, so it cancels out.
                #D = D_GGX(roughness, NoH)
                G = G_GGXSmith(roughness, NoL, NoV_proxy)
                Fc_term = pow(1.0 - VoH, 5.0)

                # PDF of GGX Importance Sampling is D * NoH / (4 * VoH).
                # The full term is (D * G * NoL) / PDF, which simplifies to:
                # G * NoL * (4 * VoH / NoH).
                # This can be unstable when NoH is close to zero, so we clamp the denominator.
                common_term = (G * NoL * 4.0 * VoH) / max(NoH, 1e-5)

                # We are baking the two components of the split-sum approximation for IBL:
                # reflectance = f0 * scale + bias
                scale = common_term * (1.0 - Fc_term)
                bias = common_term * Fc_term
            # --- Cloth BRDF ---
            elif brdf_type == 2:
                # We are importance sampling GGX, so must account for its PDF.
                D_c = D_Cloth(roughness, NoH)
                G_c = G_Cloth(roughness, VoH)

                # PDF = D_GGX(r, NoH) * NoH / (4 * VoH)
                pdf_ggx = D_GGX(roughness, NoH) * NoH / (4.0 * VoH + 1e-6)

                # We must divide by the PDF and multiply by our target distribution and the cosine term.
                scale = (D_c * G_c * NoL) / (pdf_ggx + 1e-6)
                bias = 0.0

            A += scale
            B += bias

    return A / num_samples, B / num_samples


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

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

    r, g = 0.0, 0.0
    if brdf_type == 1: # standard
        r, g = integrate_brdf_jitted(roughness, NoV, 1, num_samples)
    elif brdf_type == 2: # cloth
        if roughness < 1e-4: return x, y, 0.0, 0.0, 0.0
        r, g = integrate_brdf_jitted(roughness, NoV, 2, num_samples)

    return x, y, r, g, 0.0


def generate_exr(resolution, output_filename, brdf_type, num_samples, num_workers):
    print(f"Generating {resolution}x{resolution} EXR '{output_filename}' ({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) }

    pixel_data = np.zeros((resolution, resolution, 3), 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, brdf_type=brdf_type, 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 = future.result()
                pixel_data[y, x] = (r, g, b)
            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()
        exr_file.writePixels({'R': r_data, 'G': g_data, 'B': b_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 DFG LUT EXR images for PBR.')
    parser.add_argument('-t', '--type', choices=['standard', 'cloth'], default='standard',
                        help='Type of DFG texture to generate (default: standard)')
    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',
                        help='Output filename (default: dfg_standard.exr or dfg_cloth.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

    brdf_type = 1 if args.type == 'standard' else 2
    output_filename = args.output if args.output else f'dfg_{args.type}.exr'

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

    return 0

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