Create Rose Lens surface in Blender polar and cartesian coordinates

blender_create_rose_lens_surface.py

import bpy
import math
import numpy as np

import sys
sys.path.append(r"C:\Users\daguiam264\AppData\Roaming\Python\Python311\site-packages")



from scipy.constants import micro, nano, milli
from scipy import ndimage

from scipy import interpolate




def digitize_array_to_bins(array, levels):
    """Digitizes the given array to within the number of levels provided 
    
    Args:
        :array:  input array of values
        :levels: integer number of levels to consider or array of levels
        
    Returns:
        :bins:      bins corresponding to the levels
        :digitized: digitized array
        
    To do:
        Consider the midpoint selection in the future
    """    
    assert isinstance(levels, (np.ndarray, int)), "levels must be a scalar or numpy array"
    if isinstance(levels, int):
        bins = np.linspace(np.nanmin(array), np.nanmax(array) , levels, endpoint=False) 
    else:
        bins = levels
    
    dig = np.digitize(array, bins, )
    
    # Everything below the minimum bin level is changed to the minimum level
    dig[dig==0] = 1
    dig = dig-1
    # dig[np.isnan(array)] = np.nan
    return bins, dig
def phase2height(phase, wavelength, n1, n0=1):
        """Converts the phase to height
        Args:
            :wavelength:    Wavelength of the light
            :n1:            Refractive index of the medium where the light is propagating
            :n0:            Refractive index of the medium background"""


        height = phase * wavelength/(2*np.pi*(n1-n0))

        return height


def height2phase(height, wavelength, n1, n0=1):
    """Converts the height to phase
    Args:
        :wavelength:    Wavelength of the light
        :n1:            Refractive index of the medium where the light is propagating
        :n0:            Refractive index of the medium background"""

    phase = height * 2*np.pi*(n1-n0)/wavelength
    return phase



def fresnel_lens_phase(XX,YY,focal_length,wavelength, phase_offset=np.pi): 
    """
    returns the COMPLEX PHASE of a fresnel lens with input meshgrid (x,y) with center at (x0,y0)
    
    Args:
        :XX:            x array from meshgrid 
        :YY:            y array from meshgrid 
        :focal_length:  focal distance 
        :wavelength:    wavelength of design
    
    Note: for angle (in rad), call numpy.angle(...)
    """

    rc = np.sqrt((XX)**2 + (YY)**2)
    fresn = np.exp(1.0j*((focal_length-np.sqrt(focal_length**2 + rc**2))*(2*np.pi)/(wavelength) + phase_offset))
    fresn = np.angle(fresn)
    # fresn = fresn-np.min(fresn)
    
    return fresn     


def fresnel_lens_phase_unwrap(XX,YY,focal_length,wavelength, phase_offset=np.pi): 
    """
    returns the COMPLEX PHASE of a fresnel lens with input meshgrid (x,y) with center at (x0,y0)
    
    Args:
        :XX:            x array from meshgrid 
        :YY:            y array from meshgrid 
        :focal_length:  focal distance 
        :wavelength:    wavelength of design
    
    Note: for angle (in rad), call numpy.angle(...)
    """

    rc = np.sqrt((XX)**2 + (YY)**2)
    fresn = ((focal_length-np.sqrt(focal_length**2 + rc**2))*(2*np.pi)/(wavelength) + phase_offset)
    # fresn = fresn-np.min(fresn)
    
    return fresn     


def theta_to_wavelength(theta, wavelengths, L_repetitions=1, ramp_up_down=True): 
    """
    Converts a theta map to the respetive wavelengths for axisimmetric fresnel phase

    """
    wavelength_range = wavelengths

    if ramp_up_down: 
        wavelength_range = np.concatenate((wavelengths, np.flip(wavelengths)))
    else:
        wavelength_range = wavelengths
        wavelength_range = np.concatenate((wavelengths, []))


    # theta_range = np.linspace(-np.pi, np.pi, len(wavelength_range))
    # theta_range = np.linspace(0, 1, len(wavelength_range), endpoint=True)
    theta_range = np.arange(0, 1, 1/len(wavelength_range))


    theta_norm = np.mod(theta, 2*np.pi)/(2*np.pi)
    theta_norm *= L_repetitions
    theta_norm = np.mod(theta_norm, 1)

    wavelength_interp = interpolate.interp1d(theta_range, wavelength_range, kind='previous',fill_value="extrapolate")

    ramp_wavelength = wavelength_interp(theta_norm)

    unique, counts = np.unique(ramp_wavelength, return_counts=True)
    
#    print("Unique counts: ", unique, counts)
    
    return ramp_wavelength 



def get_material_n(wavelength, filename=r"C:\Users\daguiam264\INL\Sensitive Industry (Internal INL) - General\2. Development\32. Rose lens grayscale optimized\blender\maP1225_av_nk.txt"): 

    
    """
    Interpolates the refractive index of a material given a wavelength
    """
    material_wl, material_n, _,_,_,_,_,_,_ = np.loadtxt(filename, skiprows=3, unpack=True) 

    material_wl *= 1e-9
    n_interp = interpolate.interp1d(material_wl, material_n, kind='cubic', fill_value="extrapolate")
    return n_interp(wavelength)






def best_orders_for_target_height(wavelengths,target_height, orders, n_material, n0=1, VERBOSE_PLOT=False):

    array_heights = []

    # orders = np.arange(1,20)
    for order in orders:
        heights = wavelengths/(n_material-n0)* order

        array_heights.append(heights)
 

    array_heights = np.array(array_heights)


    # target_height = 6*micro

    merit_function_order = np.abs(array_heights-target_height)

    order_min_idx = np.argmin(merit_function_order, axis=0)
    best_orders = orders[order_min_idx]
 
    # find the heights for the best orders

    best_heights = array_heights[order_min_idx, np.arange(len(wavelengths))]

    return (best_orders, best_heights)






def create_math_surface():
    # Configuration
    name = "MathSurface"
    resolution = 50  # Number of segments
    size = 1        # Physical size of the grid
    
    # Create mesh and object
    mesh = bpy.data.meshes.new(name)
    obj = bpy.data.objects.new(name, mesh)
    bpy.context.collection.objects.link(obj)
    
    verts = []
    faces = []
    
    # Generate Vertices
    for i in range(resolution):
        for j in range(resolution):
            # Map i, j to x, y coordinates
            x = (i / (resolution - 1) - 0.5) * size
            y = (j / (resolution - 1) - 0.5) * size
            
            # --- YOUR MATH FUNCTION HERE ---
            # Example: A ripple effect
            d = math.sqrt(x**2 + y**2)
            z = math.sin(d * 4) * 0.5
            # -------------------------------
            
            verts.append((x, y, z))
            
    # Generate Faces
    for i in range(resolution - 1):
        for j in range(resolution - 1):
            # Indexing the grid vertices
            v1 = i * resolution + j
            v2 = (i + 1) * resolution + j
            v3 = (i + 1) * resolution + (j + 1)
            v4 = i * resolution + (j + 1)
            faces.append((v1, v2, v3, v4))
            
    # Load geometry into the mesh
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    
    
    
def modulos(aperture, mod, normalize_to_max=True,mod_tolerance=1e-64, align_max=True,  no_align=True):

        aux = aperture

        func_align = np.min
        if align_max:
            func_align = np.max
        aperture = (aux-func_align(aux)-mod_tolerance) % (mod)
        if no_align:
            aperture = (aux-mod_tolerance) % (mod)
        return aperture


def create_rose_lens_cartesian(XX,YY,N, focal_length, L_repetitions, wavelengths, target_height=10):
    # Configuration
    name = "RoseLens"
    resolution = XX.shape[0]  # Number of segments
    
    scale = micro/milli

    
    # Create mesh and object
    mesh = bpy.data.meshes.new(name)
    obj = bpy.data.objects.new(name, mesh)
    bpy.context.collection.objects.link(obj)
    
    verts = []
    faces = []
    
    
    n_material = get_material_n(wavelengths)
    n0 = 1
    orders = np.arange(1,20)
#    print("orders", orders)
    best_orders, best_heights = best_orders_for_target_height(wavelengths,target_height, orders, n_material,  VERBOSE_PLOT=False)
    orders = best_orders
    
    theta = np.arctan2(YY,XX)
    
    phase_offset = np.pi
    ramp_up_down = True
    
    wavelength =     theta_to_wavelength(theta, wavelengths, L_repetitions=L_repetitions, ramp_up_down=ramp_up_down)
    

    
    wavelength_order =     theta_to_wavelength(theta, orders, L_repetitions=L_repetitions, ramp_up_down=ramp_up_down)
    
    aperture = fresnel_lens_phase_unwrap(XX,YY, focal_length, wavelength, phase_offset=phase_offset)

    aperture -= np.min(aperture)
    

    MOD_order = wavelength_order
    aperture = modulos(aperture, 2*np.pi*MOD_order, align_max=True)
    
    refractive_index = get_material_n(wavelength)
    aperture = phase2height(aperture, wavelength, refractive_index)
    
    RR = np.sqrt(XX**2+YY**2)
    max_R = np.ptp(XX)/2
    aperture /= 20
    scale_space = 10/milli
    XX *= scale_space
    YY *= scale_space
    ZZ = aperture/micro

    offset = N * N

  
# 4. Flatten data for Blender (Top + Bottom layers)
    # We stack the top coordinates and bottom coordinates
    top_verts = np.stack((XX.flatten(), YY.flatten(), ZZ.flatten()), axis=-1)
    bot_verts = np.stack((XX.flatten(), YY.flatten(), np.full(N*N, floor_z)), axis=-1)
    all_verts = np.vstack((top_verts, bot_verts))
    print(all_verts)
# 6. Optimized Face Generation
    # Standard grid indexing for Top and Bottom
    for i in range(N - 1):
        row_start = i * N
        next_row_start = (i + 1) * N
        for j in range(N - 1):
            v1, v2 = row_start + j, row_start + j + 1
            v3, v4 = next_row_start + j + 1, next_row_start + j
            
            # Top Face
            faces.append((v1, v2, v3, v4))
            # Bottom Face (flipped)
            faces.append((v4 + offset, v3 + offset, v2 + offset, v1 + offset))

    # 7. Side Walls (Boundary Stitching)
    for i in range(N - 1):
        # Left/Right edges
        faces.append((i*N, (i+1)*N, (i+1)*N + offset, i*N + offset))
        faces.append(((i+1)*N - 1, (i+2)*N - 1, (i+2)*N - 1 + offset, (i+1)*N - 1 + offset))
        # Front/Back edges
        faces.append((i, i + 1, i + 1 + offset, i + offset))
        faces.append((offset - N + i, offset - N + i + 1, 2*offset - N + i + 1, 2*offset - N + i))

    # 8. Load into Blender
    mesh.from_pydata(all_verts.tolist(), [], faces)
    mesh.update()



def create_rose_lens_polar(aperture_radius, rings, segments, floor_z, focal_length, L_repetitions, wavelengths, target_height=10*micro):
    # Configuration
    name = "RoseLensPolar"
    # Create mesh and object
    mesh = bpy.data.meshes.new(name)
    obj = bpy.data.objects.new(name, mesh)
    bpy.context.collection.objects.link(obj)
    
    verts = []
    faces = []
    print("Creating vertices")

    n_material = get_material_n(wavelengths)
    n0 = 1
    orders = np.arange(1,20)
#    print("orders", orders)
    best_orders, best_heights = best_orders_for_target_height(wavelengths,target_height, orders, n_material,  VERBOSE_PLOT=False)
    orders = best_orders
    
    print(wavelengths, orders, target_height, n_material)

    refractive_indexes = get_material_n(wavelengths)
    # 2. Generate Vertices (Top and Bottom)
    # We create two layers: Top (0) and Bottom (1)
    for layer_z in [None, floor_z]:
        for r_idx in range(rings):
            r = (r_idx / (rings - 1)) * aperture_radius
            
            for s_idx in range(segments):
                theta = (s_idx / segments) * 2 * np.pi
                
 
                
                x = r * np.cos(theta)
                y = r * np.sin(theta)
                
                # If layer_z is None, calculate the Top surface math
                if layer_z is None:

                    phase_offset = np.pi
                    wavelength =     theta_to_wavelength(theta, wavelengths, L_repetitions=L_repetitions, ramp_up_down=ramp_up_down)
#                    wavelength_order =     theta_to_wavelength(theta, orders, L_repetitions=L_repetitions, ramp_up_down=ramp_up_down)

                    aperture = fresnel_lens_phase_unwrap(x,y, focal_length, wavelength, phase_offset=phase_offset)
                    
                    max_aperture = fresnel_lens_phase_unwrap(0,0, focal_length, wavelength, phase_offset=phase_offset)
                    aperture -= max_aperture          
                    MOD_order = best_orders[np.where(wavelengths==wavelength)]
                    aperture = modulos(aperture, 2*np.pi*MOD_order, align_max=True)            
                    n = refractive_indexes[np.where(wavelengths==wavelength)]
#                    print("iteration ", r_idx, s_idx, wavelength, n, MOD_order, aperture, 2*np.pi*MOD_order)                    
                    aperture = phase2height(aperture, wavelength, n)            
                    aperture /= 20
                    z =  aperture
                    

                else:
                    z = layer_z

                scale_space = 10/milli
                x *= scale_space
                y *= scale_space
                z = z/micro
                verts.append((x, y, z))
            print("iteration ", r_idx, s_idx, wavelength, n, MOD_order, aperture, 2*np.pi*MOD_order)                    
    # 3. Generate Faces
    offset = rings * segments # Offset between Top and Bottom layers
    print(" creating faces")
    for r in range(rings - 1):
        for s in range(segments):
            # Modular arithmetic to wrap the circle (connect last segment to first)
            s_next = (s + 1) % segments
            
            # Indices for Top
            v1 = r * segments + s
            v2 = r * segments + s_next
            v3 = (r + 1) * segments + s_next
            v4 = (r + 1) * segments + s
            
            # Top Layer Face
            faces.append((v1, v2, v3, v4))
            
            # Bottom Layer Face (Flipped)
            faces.append((v4 + offset, v3 + offset, v2 + offset, v1 + offset))
            
            # 4. Side Walls (Only on the outermost ring)
            if r == rings - 2:
                faces.append((v4, v3, v3 + offset, v4 + offset))

    # 5. Build Mesh
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    
    # Set smooth shading
    for poly in mesh.polygons:
        poly.use_smooth = True
# Run the function
#create_math_surface()


wavelengths = np.array([450,550,650])*nano

focal_length = 25*milli
L_repetitions = 10
target_height = 10*micro
ramp_up_down = False
ramp_up_down = True

# Configuration
name = "RoseLens"
aperture_radius = 1000 * micro
rings = 200    # Resolution along the radius (N)
segments = 360 # Resolution around the circle (Theta)
floor_z = -0.1 * micro

create_rose_lens_polar(aperture_radius, rings, segments, floor_z, focal_length, L_repetitions, wavelengths, target_height=target_height)

#aperture_width = 2000*micro
#N = 500

#x =y=np.linspace(-aperture_width/2,aperture_width/2,N)


#aperture_radius = 1000 * micro
#rings = N    # Resolution along the radius (N)
#segments = 200 # Resolution around the circle (Theta)
#floor_z = -100 * micro
#XX,YY = np.meshgrid(x,y)


#print("creating rose lens")
#create_rose_lens(XX,YY,N, focal_length, L_repetitions, wavelengths, target_height=target_height)