HarrisonKramer
Support for stereographic coordinates based computations
Hello there friends from optiland,
Thanks for making this project available to everyone.
I am trying to use optiland to: I need to obtain the partial derivatives of different intermediate ray directions (expressed as stereographic coordinates) with respect to the input ray directions (also expressed as stereographic coordinates).
I have already added my own function to the RayGenerator class which returns an instance of RealRays.
Then I can perform the ray-tracing, compute the stereographic coordinate functions that I need. So far, until this point all works nicely.
The issue that I am having is found once I move to calculating the partial derivatives of the output stereographic coordinates with respect to the input ones.
As far as I undestand it, I cannot just use backward() in this case for the entire returned stereographic output tensors. I tried using a small loop to go through each entry and compute the derivatives. However, when I observe the returned gradients, I obtain tensors which include the derivatives with respect to all input stereographic coordinates.
In my specific case, there is a one to one mapping relationship betwen each input and output pair of coordinates, so I dont need to compute the full Jacobians. (When I try to do this, this is really slow; I am just using a CPU).
I was trying to use vmap to vectorize the operation (and avoid computing the full jacobians) using the following:
def stereo_coords_s_single(u_i,v_i):
u = u_i.unsqueeze(0)
v = v_i.unsqueeze(0)
my_stereo_local = StereoCoords(my_surfs,u,v,wavelengths=0.650)
stereo_coords = my_stereo_local.stereo_coords(1)
return stereo_coords[0,0], stereo_coords[1,0]
def grad_fn(u_i,v_i):
u_i = u_i.clone()#.detach().requires_grad_(True)
v_i = v_i.clone()#.detach().requires_grad_(True)
out1, out2 = stereo_coords_s_single(u_i,v_i)
grads = be.autograd.jacrev(outputs=(out1, out2), inputs=(u_i, v_i), grad_outputs=(be.tensor(1., dtype=be.float64), be.tensor(1., dtype=be.float64)), create_graph=False)
return grads
and then calling:
stereo_params = np.loadtxt("ellip_b_to_circ_ff")
u = be.tensor(stereo_params[:,0],dtype=be.float64,requires_grad=True)
v = be.tensor(stereo_params[:,1],dtype=be.float64,requires_grad=True)
grad_fn_vmap = be.vmap(grad_fn)
grad_fn_vmap(u,v)
However, this results in:
File "/home/kalosu/fenics_gen_func/optiland_stereo_grads.py", line 51, in _generate_ray_dirs_full
data.append(self._generate_ray_dirs(u,v,wavelengths))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/kalosu/fenics_gen_func/optiland_stereo_grads.py", line 58, in _generate_ray_dirs
self.optic.myowm_trace(u,v,wavelength)
File "/usr/local/lib/python3.12/dist-packages/optiland/optic/optic.py", line 427, in myowm_trace
return self.ray_tracer.myown_stereo_tracer(u,v,wavelength)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/optiland/raytrace/real_ray_tracer.py", line 66, in myown_stereo_tracer
self.optic.surface_group.trace(rays)##Here we do not touch anything. trace from surface_groups just needs a RealRays instance
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/optiland/surfaces/surface_group.py", line 205, in trace
surface.trace(rays)
File "/usr/local/lib/python3.12/dist-packages/optiland/surfaces/standard_surface.py", line 86, in trace
return self._trace_real(rays)
^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/optiland/surfaces/standard_surface.py", line 233, in _trace_real
t = self.geometry.distance(rays)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/optiland/geometries/standard.py", line 112, in distance
t[a == 0] = -c[a == 0] / b[a == 0]
~^^^^^^^^
RuntimeError: vmap: We do not support batching operators that can support dynamic shape. Attempting to batch over indexing with a boolean mask.
So as you can see, the error happens even before calling jacrev from autograd. As far as I can see, it is related to some issue related to how vmap tries to vectorize some operations that are part of the internals of the ray-tracing engine.
Before I try anything like using a different mask which is compatible with vmap, I wanted to ask if there was any suggestion/comment or example on how could I achieve what I want without having to go and modify baseline code.
Any thoughts??
Hi @kalosu ,
Thanks for your message and for taking the time to explore Optiland so deeply. It sounds like you've already made it quite far and are converging.
Regarding the issue with using vmap: you're right that this stems from internal boolean indexing (e.g., t[a == 0] = ...) which is not compatible with PyTorch's current vectorization backend. Optiland wasn’t originally designed with vmap support in mind (although we could change this), so this kind of error can appear if you try to batch across internal operations that use dynamic shapes or masks.
Since you're already modifying the execution flow (e.g., adding custom ray generation and tracing steps), it might make sense to also consider modifying or bypassing the indexing step causing the failure. For example, in your case, removing or replacing the line involving the mask (t[a == 0] = ...) could work if you're confident that a == 0 won't occur or that such cases are easy enough to catch (e.g., by detecting a divide-by-zero later). That should remove the need for boolean masks and make vectorization feasible. Just be aware that this changes how certain edge cases are handled, so I'd recommend making this explicit in your workflow.
In the meantime, as a workaround, you might find that computing the gradients per-ray using a simple loop (e.g., with autograd.grad on each scalar pair) is both easier to implement and fast enough for most practical CPU-scale problems. That way, you avoid vmap altogether and keep Optiland's internals largely untouched.
I'll also point out that the PyTorch backend is still relatively new, and we're actively expanding examples and improving common workflows. If you have interesting use cases or feedback, we’d love to hear more.
Let me know how it goes here.
Regards, Kramer
Hi there @HarrisonKramer,
Thanks a lot for the feedback. At the end I managed to replace that specific line with an operation which is supported by vmap.
Specifically, I just used be.where instead of the index mask. The good thing is that this was the only part in the computational graph which was not compatible with vmap. I did not have to change anything else.
After this, I could use vectorization as I needed it.
Thanks a lot for the feedback and help with this!
Hi @kalosu ,
Great! Glad it was a simple fix. I'll plan to update that call to be.where, as well as look for other instances of boolean indexing throughout the package that might also need to be replaced. I'll keep this issue open to track that update.
Kramer