Differentiable Ray Tracing

# Begin by importing the backend module instead of numpy or torch directly
import optiland.backend as be

# set backend
be.set_backend("torch")
be.set_device(
  "cpu"
)  # we will stick to cpu for now. "cuda" is also possible for GPU calculations.
print(be.get_precision())  # check the precision

# enable the gradient mode: to tell autograd to begin
# recording operations on a Tensor tensor
be.grad_mode.enable()

from optiland.optic import Optic
from optiland.materials import Material
from optiland.analysis import SpotDiagram


class SingletConfigurable(Optic):
  """A configurable singlet lens."""

  def __init__(self, r1, r2, t2, material_name):
      super().__init__()
      ideal_material = Material(material_name)

      self.add_surface(index=0, radius=be.inf, thickness=be.inf)
      self.add_surface(
          index=1, thickness=7.0, radius=r1, is_stop=True, material=ideal_material
      )
      self.add_surface(index=2, radius=r2, thickness=t2)
      self.add_surface(index=3)

      self.set_aperture(aperture_type="EPD", value=25)

      self.set_field_type(field_type="angle")
      self.add_field(y=0.0)

      self.add_wavelength(value=0.55, is_primary=True)

lens = SingletConfigurable(r1=70.0, r2=-70.0, t2=70.0, material_name="BK7")
lens.draw(num_rays=10)
lens.info()

# RMS Spot Size Analysis
spot_diagram = SpotDiagram(lens)
rms_value = spot_diagram.rms_spot_radius()
print(f"\nInitial RMS Spot Radius: {rms_value[0][0]:.4f} mm")

# Spot Diagram visualization
spot_diagram.view(add_airy_disk=True)

import torch.optim as optim

# --- Parameter Handling ---

# Get initial values
radius1_param = lens.surface_group.surfaces[1].geometry.radius
radius2_param = lens.surface_group.surfaces[2].geometry.radius
params_to_optimize = [radius1_param, radius2_param]
print(params_to_optimize)

# --- Optimizer ---
# We use the Adam optimizer, common for deep learning tasks.
# We pass the tensors we want to optimize.
learning_rate = 0.2  # Adjust as needed
optimizer = optim.Adam(params_to_optimize, lr=learning_rate)

# --- Loss Function ---
# The function to minimize is the RMS spot radius.
# We need a way to calculate this within the loop using the current tensor values.


def calculate_loss(lens):
  """Calculates the RMS spot radius for the current lens state."""
  rms_values_nested = SpotDiagram(lens).rms_spot_radius()
  rms_scalar = rms_values_nested[0][0]
  return rms_scalar


print("\nOptimization setup complete.")
print(f"Optimizing: R1, R2")
print(f"Initial values: (R1={radius1_param:.2f}, R2={radius2_param:.2f})")


# --- Optimization Loop ---
# We run the standard PyTorch training loop.

num_iterations = 100
print("\nStarting optimization...")
losses = []
for i in range(num_iterations):
  optimizer.zero_grad()

  # --- Calculate Loss ---
  loss = calculate_loss(lens)
  losses.append(loss.item())
  # --- Backpropagation and Optimization Step ---
  loss.backward()  # Calculate gradients
  optimizer.step()  # Update tensors based on gradients

  # --- Print Progress ---
  if (i + 1) % 25 == 0:
      print(
          f"Iteration {i + 1}/{num_iterations}, "
          f"Loss (RMS Spot Radius): {loss.item():.6f} mm"
      )
      print(f"  R1: {radius1_param.item():.4f}, R2: {radius2_param.item():.4f}")

print("Optimization finished.")

# --- Visualization of Loss per Epoch ---
import matplotlib.pyplot as plt

plt.figure()
plt.plot(range(1, num_iterations + 1), losses)
plt.xlabel("Iteration")
plt.ylabel("Loss (RMS Spot Radius) [mm]")
plt.title("Loss per Iteration")
plt.grid(True)
plt.show()

SpotDiagram(lens).view(add_airy_disk=True)