Torch Lens Maker – Differentiable Geometric Optics in PyTorch

Hacker News

Share card

Torch Lens Maker – Differentiable Geometric Optics in PyTorch

Hello HN! For the past 6 months I've been working on an open source python library that implements differentiable geometric optics in PyTorch. It's very experimental still, but eventually the goal is to use it to design optical systems with a state of the art optimization framework and a beautiful code based API. Think OpenSCAD, but for optical systems. 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!