Collimated Comparison Local

This script was run on a regular machine, not a server, and is not used in the publication.

It demonstrates the same collimated optical system comparison workflow as the server notebook, but without cloud or Colab-specific features.

All computations, optimizations, and visualizations are performed locally, and results are saved to the local filesystem.

This approach is useful for development, debugging, or running smaller experiments without the need for cloud resources

import sys
import os
import gc
import tqdm
sys.path.insert(0, os.path.abspath(".."))
device = "cuda:0"

import diffinytrace as dit
import torch
import numpy as np
torch.set_default_dtype(torch.float64)
import random

SEED = 0
np.random.seed(SEED)
torch.manual_seed(SEED)
random.seed(SEED)

from diffinytrace.source import CollimatedMonochromatic,CollimatedGaussianBeam
from diffinytrace.gaussian_smoother import GaussianSmootherSquare,make_merit_function,make_evaluation_function
import matplotlib.pyplot as plt

def run_optimiaztion(sigma,use_desired_irradiance_smoothing,return_system=False):
    in_aperture = 4.0
    in_aperture_lens = 5.0

    desired_width_square = 4.0

    out_aperture = 8.0
    source_wl = 0.589
    source_gaussian_constant = 0.035

    light_transform = dit.transforms.Identity()
    source = CollimatedGaussianBeam(light_transform,in_aperture,source_wl,source_gaussian_constant)

    lens_mat = dit.materials["NBK7"]
    env_mat = dit.materials["NONE"]

    lens1_thickness = 2.
    lens1_surf1 = dit.Bspline(in_aperture_lens,[4,4],[11,11])#dit.Legendre(in_aperture_lens,20)#
    lens1_surf2 = dit.Plane()

    det_surf = dit.Plane()

    lens1_transform = dit.transforms.Distance(5.0,parent_transform=source)
    lens1_transform.distance.requires_grad = False

    lens1 = dit.Lens(lens1_transform,lens1_thickness,lens1_surf1,lens1_surf2,lens_mat,in_aperture_lens,is_square=False)


    det_transform = dit.transforms.Distance(10.,parent_transform=lens1)
    det_transform.distance.requires_grad = False
    detector = dit.Detector(det_transform,det_surf,out_aperture)

    system = dit.SequentialOpticalSystem({"source":source,"lens1":lens1,"detector":detector},env_mat)
    system.to(device)
    sequence = ["source","lens1","detector"]

    N = 2**15
    grid_size = 128

    def desired_irradiance_fun(y):
        out = (torch.abs(y[:,0]) < desired_width_square).float() * (torch.abs(y[:,1]) < desired_width_square).float()
        return out

    smoother = GaussianSmootherSquare(out_aperture,
                                    grid_size=grid_size,
                                    sigma=sigma,
                                    desired_irradiance_fun=desired_irradiance_fun,
                                    smoothed_num_integration_points=2**21,
                                    smoothed_num_splits=10,
                                    device=device)
    #dit.plotting.quantity2D.plot(smoother.smoothed_desired_irradiance,"Smoothed [W/mm²]",out_aperture)
    #dit.plotting.quantity2D.plot(smoother.discrete_desired_irradiance,"Discrete [W/mm²]",out_aperture)
    #plt.savefig("collimated_desired_irr.png")
    """

    x,weights = source.sample(N,method="sobol")
    x = x.to("cuda:0")
    weights = weights.to("cuda:0")

    ray_multi = source.get_flux(x)*weights
    O2, D2, wl, n_func_enviroment, meta_data = system(x,sequence)
    local_pos = detector.to_local_pos(O2)[:,[0,1]]
    out = smoother.get_none_smoothed_irradiance(local_pos,ray_multi)

    dit.plotting.quantity2D.plot(out,"Irradiance Initial",out_aperture)

    x,weights = source.sample(10)
    x = x.to("cuda:0")
    weights = weights.to("cuda:0")
    O2, D2, wl, n_func_enviroment, meta_data = system(x,sequence)
    dit.plotting.system2D.plot(system,meta_data)
    """

    convergence_list = []

    merit_func = make_merit_function(optical_system=system,sequence=sequence,source=source,detector=detector,smoother=smoother,num_rays=N,use_desired_irradiance_smoothing=use_desired_irradiance_smoothing,device=device)
    evaluation_func = make_evaluation_function(optical_system=system,sequence=sequence,source=source,detector=detector,smoother=smoother,num_rays_per_split=100000,num_splits=10,device=device)

    def convergence_callback_func(convergence_list:list):
        def return_func():
            convergence_list.append(evaluation_func())

        return return_func

    result = dit.optimize.minimize(merit_func,system.parameters(),method="L-BFGS-B",save_history=True,call_before_minimize=True)
    result["sigma"] = sigma
    result["use_desired_irradiance_smoothing"] = use_desired_irradiance_smoothing
    convergence_callback_func(convergence_list)()
    result["convergence_list"] = convergence_list

    none_smooth_irr = []
    smooth_irr = []
    with torch.no_grad():
        for k in range(10):
            x,weights = source.sample(100000,method="monte_carlo")
            x = x.to("cuda:0")
            weights = weights.to("cuda:0")
            ray_multi = source.get_flux(x)*weights
            O2, D2, wl, n_func_enviroment, meta_data = system(x,sequence)
            local_pos = detector.to_local_pos(O2)[:,[0,1]]
            none_smooth_irr.append(smoother.none_smoothed_irradiance(local_pos,ray_multi).detach().cpu())
            smooth_irr.append(smoother.smoothed_irradiance(local_pos,ray_multi).detach().cpu())
    none_smooth_irr = torch.mean(torch.stack(none_smooth_irr),dim=0)
    smooth_irr = torch.mean(torch.stack(smooth_irr),dim=0)

    result["smooth_irr"] = smooth_irr.cpu().detach()
    result["none_smooth_irr"] = none_smooth_irr.cpu().detach()
    if return_system:
        return result,system

    return result

import torch

sigmas = torch.linspace(0.05,1.,20)
def run_all(use_desired_irradiance_smoothing):
    all_results = []
    for i,sigma in tqdm.tqdm(enumerate(sigmas)):
        all_results.append(run_optimiaztion(sigma,use_desired_irradiance_smoothing))
        #print(f"finished {i+1}/{len(sigmas)}")

    return all_results

baseline_results = run_all(use_desired_irradiance_smoothing=False)
our_results = run_all(use_desired_irradiance_smoothing=True)

20it [05:36, 16.84s/it]
20it [09:25, 28.28s/it]

plt.figure(figsize=(10,4))
results_folder = "results/collimated_compare/"
try:
    os.mkdir("results")
except:
    pass
try:
    os.mkdir(results_folder)
except:
    pass
baseline_conv = [result["convergence_list"][-1] for result in baseline_results]
our_conv = [result["convergence_list"][-1] for result in our_results]
sigmas = [result["sigma"] for result in our_results]
ax = plt.gca()
ax.grid(True, which='major', linestyle='-', linewidth=0.5)  # Minor grid lines (finer)
ax.grid(True, which='minor', linestyle='-', linewidth=0.5)  # Minor grid lines (finer)
ax.set_xlabel("$\\sigma$ [mm]")
ax.set_ylabel("Error")
plt.plot(sigmas,baseline_conv,label="Partially Smoothed")
plt.plot(sigmas,our_conv,label="Ours")
plt.title("Relationship Between Error and the Kernel Width")
plt.legend()
plt.savefig(results_folder+"relationshipVA.png", dpi=400, bbox_inches='tight')


idx = torch.arange(len(baseline_results)//4)*5
idx = [0,5,10,15,19]

plt.figure(figsize=(12,4))

baseline_conv = [result["convergence_list"][-1] for result in baseline_results]
our_conv = [result["convergence_list"][-1] for result in our_results]
sigmas = [result["sigma"] for result in our_results]
ax = plt.gca()
ax.grid(True, which='major', linestyle='-', linewidth=0.5)  # Minor grid lines (finer)
ax.grid(True, which='minor', linestyle='-', linewidth=0.5)  # Minor grid lines (finer)
ax.set_xlabel("$\\sigma$ [mm]")
ax.set_ylabel("Error")
plt.plot(sigmas,baseline_conv,label="Partially Smoothed")
plt.plot(sigmas,our_conv,label="Ours")
plt.title("Relationship Between Error and the Kernel Width")
plt.legend()
plt.savefig(results_folder+"relationshipVB.png", dpi=400, bbox_inches='tight')


idx = torch.arange(len(baseline_results)//4)*5
idx = [0,5,10,15,19]

xbaseline_results = [baseline_results[i] for i in idx]
xour_results = [our_results[i]for i in idx]
xbaseline_results
data_grid = [[]]*4
#result["smooth_irr"] = smooth_irr.cpu().detach()
#result["none_smooth_irr"] = none_smooth_irr.cpu().detach()

data_grid[0] = [result["smooth_irr"] for result in xbaseline_results]
data_grid[1] = [result["smooth_irr"] for result in xour_results]
data_grid[2] = [result["none_smooth_irr"] for result in xbaseline_results]
data_grid[3] = [result["none_smooth_irr"] for result in xour_results]

from PIL import Image, ImageDraw, ImageFont
from importlib import reload

from image_grid_maker import image_from_grid
import image_grid_maker
reload(image_grid_maker)
from image_grid_maker import image_from_grid
import image_grid_maker

out_aperture = 8.0
vmin = 0.0
vmax = 0.03

rows_extent = [[-out_aperture,out_aperture,-out_aperture,out_aperture]]*4
rows_vidx = ["x","x","x","x"]
rows_cmap = ["jet"]*4
cbar_titles = ["[W/mm²]"]*4
columns_title = [f'σ={result["sigma"]} mm' for result in baseline_results]
columns_title = [columns_title[i] for i in idx]
rows_title = ["(Partially Smoothed)\nSmoothed Irr.","(Ours)\nSmoothed Irr.","(Partially Smoothed)\nIrr. RC","(Ours)\nIrr. RC"]

data_grid = [data_grid[0],data_grid[2],data_grid[1],data_grid[3]]
rows_title = [rows_title[0],rows_title[2],rows_title[1],rows_title[3]]

rows_vmin = [vmin]*4
rows_vmax = [vmax]*4

kwargs = dict(
        image_grid=data_grid,
        rows_extent=rows_extent,
        rows_vidx=rows_vidx,
        rows_cmap=rows_cmap,
        rows_title=rows_title,
        cbar_titles=cbar_titles,
        columns_title=columns_title,
        rows_vmin=rows_vmin,
        rows_vmax=rows_vmax,
)
out = image_from_grid(
    **kwargs,
    max_num_column=len(columns_title),
    font_size_PIL=40,
    cbar_labelsize=20,
    cbar_title_fontsize=20,
    column_title_ratio=0.3
    )
out = out[0]



# Load an image from the file path
image = Image.open(out)
image.save(results_folder+"comparison_collimated.png")

result,system = run_optimiaztion(sigmas[0],use_desired_irradiance_smoothing=True,return_system=True)

rows vmin: [0.0, 0.0, 0.0, 0.0] rows vmax: [0.03, 0.03, 0.03, 0.03]
_image_from_grid 4 4

#([tensor(0.), tensor(0.), tensor(0.), tensor(0.)],
# [tensor(0.0288), tensor(0.4965), tensor(0.0287), tensor(0.0393)])

out_aperture = 8.0

import numpy as np
import plotly.graph_objects as go

system.cpu()
source = system.modules_dict["source"]
lens1 = system.modules_dict["lens1"]
detector = system.modules_dict["detector"]

in_aperture = 4.0
num_bins = 512

x = np.linspace(-in_aperture, in_aperture, num_bins)  # Width
y = np.linspace(-in_aperture, in_aperture, num_bins)  # Height
z = torch.zeros((num_bins, num_bins))

tmp = np.meshgrid(x,y)
tmp = torch.tensor([tmp])[0].reshape(2,-1).T

irr_source = source.get_flux(tmp)
irr_source[torch.linalg.norm(tmp,dim=1)>in_aperture]=torch.nan
irr_source = irr_source.reshape(num_bins,num_bins)
print("irr_source",irr_source.shape)

desired_width_square = 4.
def desired_irradiance_func(y):
    out = (torch.abs(y[:,0]) < desired_width_square).float() * (torch.abs(y[:,1]) < desired_width_square).float()
    return out/((desired_width_square*2)**2)
#dit.plotting.quantity2D.plot(irr_source,title="Radiant Exitance [W/mm²]",x_range=[-out_aperture,out_aperture],cmap="hot")

x = np.linspace(-out_aperture, out_aperture, num_bins)  # Width
y = np.linspace(-out_aperture, out_aperture, num_bins)  # Height
z = torch.zeros((num_bins, num_bins))

tmp = np.meshgrid(x,y)
tmp = torch.tensor([tmp])[0].reshape(2,-1).T

desired_irr = desired_irradiance_func(tmp.reshape(-1,2)).reshape(num_bins,num_bins)
print("desired_irr",desired_irr.shape,tmp.shape)

#vmax = torch.max(irr_source).item()

import matplotlib.pyplot as plt

# Assume you already have:
# - irrs: list of 2D arrays (irradiance maps)
# - rows_extent: list of [xmin, xmax, ymin, ymax] per image
# - sigmas: list of sigma values used for smoothing

cbar_labelsize=12
cbar_title_fontsize=15
# Grid dimensions
rows_extent = [[-in_aperture, in_aperture, -in_aperture, in_aperture]] +[[-out_aperture, out_aperture, -out_aperture, out_aperture]]
irrs = [irr_source,desired_irr]
irrs = [irr.cpu() for irr in irrs]

num_rows = 1
num_cols = len(irrs)

fig, axes = plt.subplots(num_rows, num_cols, figsize=(4 * num_cols, 4), constrained_layout=True)

# Titles for each column
columns_title = ["Radiant Exitance"] + [f"Desired Irradiance" for sigma in sigmas]
cmaps = ["jet","jet"]
cbar_title = "[W/mm²]"

for k in range(num_cols):
    ax = axes[k]
    img = irrs[k]
    cmap = cmaps[k]
    im = ax.imshow(img, extent=rows_extent[k],origin='lower', cmap=cmap,interpolation="nearest",vmin=vmin,vmax=vmax)
    ax.set_title(columns_title[k],fontsize=cbar_title_fontsize)
    if k != 0:
        ax.set_xticks([-4,0,4])
    ax.set_yticks([])
    ax.tick_params(labelsize=cbar_labelsize)

    cbar = plt.colorbar(im,ax=ax,shrink=0.65,aspect=9)  # Add a colorbar for reference
    cbar.ax.tick_params(labelsize=cbar_labelsize)
    cbar.ax.set_title(cbar_title, fontsize=cbar_title_fontsize, pad=10,loc="left")  # Set label above
    offset_text = cbar.ax.yaxis.offsetText
    offset_text.set_size(cbar_labelsize)  # Set the font size
    offset_text.set_ha('left')  # Align the text to the left

#plt.suptitle("Irradiance Maps from Ray Counting and Smoothing", fontsize=16)
plt.savefig(results_folder+"radiant_exitance_desired_irr.png", dpi=400, bbox_inches='tight')

plt.show()

C:\Users\marti\AppData\Local\Temp\ipykernel_21216\1628606591.py:19: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\torch\csrc\utils\tensor_new.cpp:281.)
  tmp = torch.tensor([tmp])[0].reshape(2,-1).T

irr_source torch.Size([512, 512])
desired_irr torch.Size([512, 512]) torch.Size([262144, 2])

font_multi = 1.3
for k in range(num_cols):
    ax = plt.gca()
    img = irrs[k]
    cmap = "jet"
    im = ax.imshow(img, extent=rows_extent[k],origin='lower', cmap=cmap,interpolation="nearest",vmin=vmin,vmax=vmax)
    ax.set_title(columns_title[k],fontsize=cbar_title_fontsize*font_multi)
    if k != 0:
        ax.set_xticks([-4,0,4])
    ax.set_yticks([])
    ax.tick_params(labelsize=cbar_labelsize*font_multi)

    cbar = plt.colorbar(im,ax=ax,shrink=0.65,aspect=9)  # Add a colorbar for reference
    cbar.ax.tick_params(labelsize=cbar_labelsize*font_multi)
    cbar.ax.set_title(cbar_title, fontsize=cbar_title_fontsize*font_multi, pad=10,loc="left")  # Set label above
    offset_text = cbar.ax.yaxis.offsetText
    offset_text.set_size(cbar_labelsize*font_multi)  # Set the font size
    offset_text.set_ha('left')  # Align the text to the left

    plt.savefig(results_folder+f"radiant_exitance_desired_irr_sep{k}.png", dpi=400, bbox_inches='tight')
    plt.show()

xr,_ = source.sample(15)
sequence = ["source","lens1","detector"]

O,D,wave_len,_,RayPaths = system(xr,sequence)

if not RayPaths is None:
    if isinstance(RayPaths,dict):
        rays = RayPaths["ray_paths"]

show_grid=True
xlabel="x [mm]"
ylabel="y [mm]"
zlabel="z [mm]"
xticks=None
yticks=None
zticks=None
axislabel_font_size=10
tick_font_size=10
ray_color="#9673A6"
ray_linewidth=3.

data = []
resolution = 32
data += dit.plotting.system3D._plot_surface_recursively(lens1,"",resolution)
data += dit.plotting.system3D._plot_surface_recursively(source,"",resolution)

data += dit.plotting.system3D.ray_paths(rays,ray_color,ray_linewidth)
layout = dit.plotting.system3D.get_optical_system_layout(False,xlabel,ylabel,zlabel,xticks,yticks,zticks,axislabel_font_size,tick_font_size)


import plotly.graph_objects as go
none_smooth_irr = result["none_smooth_irr"]

num_bins = none_smooth_irr.shape[0]

out_aperture = 8.0

x = np.linspace(-out_aperture, out_aperture, num_bins)  # Width
y = np.linspace(-out_aperture, out_aperture, num_bins)  # Height
z = torch.ones((num_bins, num_bins))*detector.get_transformation_matrix()[2,3]        # Flat surface
z = z.detach()
gosurface2 = go.Surface(
        x=x,
        y=y,
        z=z,
        surfacecolor=none_smooth_irr,
        cmin=vmin,
        cmax=vmax,
        colorbar=dict(
        title=dict(text='[W/mm²]', font=dict(size=16)),  # Increase title font size
        tickfont=dict(size=14)  # Increase tick font size
        ),
        colorscale="jet",
        showscale=True,
        showlegend=False
)

data += [gosurface2]
fig = go.Figure(data=data,layout=layout)
fig.write_html(results_folder+"collimated2.html")
#fig.show()

c:\Users\marti\anaconda3\Lib\site-packages\torch\functional.py:513: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\aten\src\ATen\native\TensorShape.cpp:3610.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]