Optimization Evaluator - Batched Operands Calculation

[manuelFragata]

Checklist

• [x] I have searched the existing issues and discussions for a similar question or feature request.
• [x] I have read the documentation and tried to find an answer there.
• [x] I am using the latest version of Optiland.
• [x] I have included all necessary context.

Thanks for taking the time to go through this — it really helps us help you!

Feature Request

On the topic of performance optimization for optimization (😅 ), the current logic is doing the following: OptimizationProblem.sum_squared() iterates through the operands Each op.value property calls operand_registry.get(...), and this triggers a RayOPerand static method, which immediately calls optic.trace_generic(...) for that one ray

this creates an inefficient loop: trace -> evaluate -> trace -> evaluate -> ...

In my branch where i have been working on the custom DLS I have kind of monkey patched this issue using:

def _prepare_operand_batches(self):
  """analyze all operands to create a single batch of unique rays for tracing"""
  ray_definitions = defaultdict(list)
  
  for i, op in enumerate(self.problem.operands):
  # is the operand a ray-based operand, and if so, does it have the input data to trace
  
  if hasattr(op, "input_data") and "Hx" in op.input_data:
    input_data = op.input_data
    ray_key = (
      input_data.get("Hx", 0.0),
      input_data.get("Hy", 0.0),
      input_data.get("Px", 0.0),
      input_data.get("Py", 0.0),
      input_data.get("wavelength", optic.primary_wavelength),
      )
    ray_definitions[ray_key].append(i)
  
  self._unique_rays = list(ray_definitions.keys())
  self._trace_to_operand_map = [ray_definitions[key] for key in self._unique_rays]


def _get_residuals_and_mf_batched(self, variables_vector):
"""calculates operand residuals and MF using a single vectorized ray trace"""

  for i, var in enumerate(self.problem.variables):
    var.update(variables_vector[i])
  self.problem.update_optics()

  if not self._unique_rays:
    return be.array([]), 0.0
  
  Hx, Hy, Px, Py, wavelengths = [
    be.array([r[i] for r in self._unique_rays]) for i in range(5)
  ]
  optic.trace_generic(Hx, Hy, Px, Py, wavelengths)
  
  deltas = be.zeros(len(self.problem.operands))
  for i, op_indices in enumerate(self._trace_to_operand_map):
    # EXAMPLE: assuming operand value is the traced ycoord at the image plane
    traced_y = optic.surface_group.y[-1, i]
    for op_idx in op_indices:
      deltas[op_idx] = traced_y - self.problem.operands[op_idx].target
  
  if be.any(be.isnan(deltas)):
    return None, be.inf

  weights = be.array([op.weight for op in self.problem.operands])
  residuals = be.sqrt(weights) * deltas
  mf_value = be.sum(residuals**2)

return residuals, mf_value

Now, i would like to extend this idea to the other operands, and it should not be applied directly within the optimmizer, but rather somewhere else. I would propose creating a class, BatchedRayEvaluator, which would:

• analyze the OptimizationProblem operand list, and do the same logic as my function above, _prepare_operand_batches.
• raytrace: it will have a method maybe called evaluate(variables_vector) which would: -- 1. update the variables from the vector provided by the optimizer; -- 2. executes a single batched optic.trace_generic() with all the unique rays.
• compute all the operands values

[manuelFragata]

Take it with a grain of salt, but these are the speed ups i get with this approach: