Show HN: Torch Lens Maker – Differentiable Geometric Optics in PyTorch (victorpoughon.github.io)
Not only is PyTorch's autograd an amazing general purpose optimizer, but torch.nn (the neural network building blocks) can be used pretty much out of the box to model an optical system. This is because there is a strong analogy to be made between layers of a neural network, and optical elements in a so-called sequential optical system. So the magic is that we can stack lenses as if we were stacking Conv2D and ReLu layers and everything works out. Instead of Conv2D you have ray-surface collision detection, instead of ReLu you have the law of refraction. Designing lenses is surprisingly like training a neural network.
Check out the docs for examples of using the API. My favorite one is the rainbow :) https://victorpoughon.github.io/torchlensmaker/examples/rain...
You should be able to `pip install torchlensmaker` to try it out, but I just set it up so let me know if there's any trouble.
I was part of the Winter 1'24 batch at the Recurse Center (https://www.recurse.com/) working on this project pretty much full time. I'm happy to talk about that experience too!
46 comments
[ 2.7 ms ] story [ 93.4 ms ] thread> The key idea is that there is a strong analogy to be made between layers of a neural network, and optical elements in a so-called sequential optical system. If we have a compound optical system made of a series of lenses, mirrors, etc., we can treat each optical element as the layer of a neural network. The data flowing through this network are not images, sounds, or text, but rays of light. Each layer affects light rays depending on its internal parameters (surface shape, refractive material...) and following the very much non‑linear Snell's law. Inference, or the forward model, is the optical simulation where given some input light, we compute the system's output light. Training, or optimization, is finding the best shapes for lenses to focus light where we want it.
Any plans to include diffractive optics as well? (A totally self-serving question, given that refractive optics is much more common.) In a past life I taught holography and wrote interactive programs to visualize the image forming properties of holograms.
I've been working off and on on a similar hobby project, working through the book _Computational Fourier Optics: A MATLAB Tutorial_, and implementing it in Jax.
My main interest is adaptive optics, but I'm only a hobbyist (limited physics background) and honestly haven't had much time to put into it.
A few notes, though paraxial approximations are "dumb", they are very useful tools for lens designers and understanding/constraining the design space - calculating the F/#, aperture stop, principal planes and is critical in some approaches. This pushes what autodiff tools are capable of because you need to get Hessians of your surface. There's also a rich history in objective function definition and quadrature integration techniques thereof which you can work to implement, and you may like to have users be able to specify explicit parametric constraints.
> There's also a rich history in objective function definition and quadrature integration techniques thereof which you can work to implement, and you may like to have users be able to specify explicit parametric constraints.
Yes, this is definitely the direction I want to take the project in. If you have any reference material to share I'd be interested!
> Forbes, G. W. (1989). Optical system assessment for design: numerical ray tracing in the Gaussian pupil. Journal of the Optical Society of America A, 6(8), 1123. https://doi.org/10.1364/josaa.6.001123
In general, you'll want to look at MTF calculation (look at Zemax's manual for explanation/how-to). There is also a technique to target optimization at particular spatial frequencies:
> K. E. Moore, E. Elliott, et. al. "Digital Contrast Optimization - A faster and better method for optimizing system MTF," in Optical Design and Fabrication 2017 (Freeform, IODC, OFT), OSA Technical Digest (online) (Optical Society of America, 2017), paper IW1A.3
Do you have any plans to add stock lens catalogue matching? To make it easier for hobbyists to manufacture lens assemblies.
You could get rid of the layer lines with post-processing like grinding/polishing, but grinding/polishing is how lenses are traditionally made anyway. Maybe you get to skip the first few steps (rough shaping of the glass blank), but you're still probably left with an optically inferior product.
It seems that dipping the lens in resin is the best solution.
How close is something like this to being competitive with ray-tracing (as featured in video game engines, or as featured in something like Blender)? I guess, since it is using Torch it should be… surprisingly performant, right? You get some hardware acceleration at least.
Ray tracing for rendering typically needs to figure out which surface a ray is hitting as part of collision detection. This is typically done with something called Bouding Volume Hierachies. Optical design (at least in sequential mode) side steps that issue completely because the order of surface collisions is known in advance.
Another big difference is that ray tracing for optical design needs to be differentiable. This is why I made this project in PyTorch, so that the entire collision detection code and physics implementation (refraction, reflection) can be differentiated with respect to parameters that describe the shape of surfaces. Then you can gradient descent the entire optical system to find optimal parameters.
Finally, rendering raytracing typically implements a lot of realism functions like diffuse or partial reflection which makes the code acutally more complex in some ways. But optical design will care more precisely about things like precise modeling of dispersion, which is not a huge focus for rendering. And there can be real-time performance constraints if you're doing a video game also. Here the implementation really doesn't care about any real time stuff.
My layman perception of NNs is they are a formalism that defines a parametric function family (for instance the family of affine functions has m + mxn parameters for n input and m output space dimensions) using some base primitives and composition rules. By tacking on a cost function, an optimizer, and a bunch of (input,output) pairs (training set) to this, one obtains optimal parameters such that the cost function is minimized over the training set (in some norm, I imagine). The NN can then be used to map never seen before inputs to outputs in a manner, hopefully, that leads to a small value of the loss function (i.e. adequately).
Even if this is wrong to some extent, could you confirm that the optimizer is but one component of a larger system, and in fact one that exists independently of NNs as well? Such as a stochastic gradient descent. In that case, what is the role of NNs in what you mentioned, would it not be simpler to yank the optimizer out and apply it to your application directly? It seems to me trying to recast a given problem to a NN to make use of a Python library's included optimizer is a sort of "XY problem", if one could just write their cost function and pass it to the optimizer directly (which presumably is no less open source than the library that includes it).
I may be misinterpreting, because this is not my field, however interpolation or projection are things very familiar to me, so I may have a bias to interpret things to resemble this. In that case, I'd welcome corrections.
Check out this project[1] which kinda does that, although it's 2D only as far as I know. But it's fully interactive, which is super neat.
[1] https://phydemo.app/ray-optics/
I can only hope that projects like these help build better lenses for the future.
By the way, the roadmap link on the Github repo page is broken (though I did find it in the project's website).
I saw you're interested in handling Bezier splines. This might be premature generalization, but would you want to support whole BREPs? Is that something that is used in applications? I imagine lenses are generally simple geometries, so perhaps not.
Honestly, the main challenge at this point feels like finding the time / resources to just work on everything I've got on the roadmap. I think I've got enough ideas for at least 1y of full time work, perhaps more.
Bezier splines are interesting. I actually had a working implementation a few weeks ago, so I know it will work. It's just that I refactored the internals of how surfaces work and didn't port the old bezier spline code yet, but it shoudln't be too difficult.
I haven't given too much thought about volume representation so far. Currently it's just a list of unconnected surfaces which seems sufficient for now.