Chromatix: a differentiable, GPU-accelerated wave-optics library

Dit paper introduceert Chromatix, een open-source, GPU-versnelde en differentieerbare bibliotheek voor golf-optische simulaties die gebaseerd is op JAX en standaardisatie, snelheid en herbruikbaarheid biedt voor computationele optica-toepassingen zoals microscopie en holografie.

De kern

Stel je voor dat je een fotograaf bent, maar in plaats van gewone foto's maakt je foto's van dingen die je met het blote oog niet kunt zien: levende cellen, zenuwcellen in een hersen, of zelfs de structuur van een heel embryo. Om dit te doen, gebruiken wetenschappers geavanceerde microscopen. Maar hier zit de twist: deze microscopen zijn niet alleen maar lenzen en spiegels; ze zijn ook een soort computerprogramma.

Vroeger was het bouwen van zo'n "slimme" microscoop als het bouwen van een auto in een garage waar elke monteur zijn eigen gereedschap en taal gebruikte. Als je een nieuwe auto wilde bouwen, moest je alles opnieuw uitvinden: de motor, de wielen, de remmen. Het was traag, vol fouten en niemand kon makkelijk samenwerken.

Chromatix is de oplossing voor dit probleem. Het is een nieuwe, open-source "gereedschapskist" voor wetenschappers die licht en optica bestuderen. Hier is hoe het werkt, vertaald in alledaagse taal:

1. De "LEGO-blokken" voor licht

Stel je voor dat je een optisch systeem (zoals een microscoop) wilt ontwerpen. In het verleden moest je de wiskunde voor elk stukje licht dat door de lens gaat, handmatig uitrekenen. Dat is als proberen een huis te bouwen door elke steen zelf te bakken.

Chromatix werkt als een LEGO-set. Het heeft standaard blokken voor alles wat licht doet: lenzen, spiegels, sensoren en zelfs hoe licht door troebel water (zoals weefsel) reist. Je hoeft niet meer de wiskunde te kennen; je klikt gewoon de blokken aan elkaar om je microscoop te bouwen.

2. De "Super-kracht" van de computer (GPU)

Lichtsimulaties zijn extreem rekenintensief. Het is alsof je probeert de beweging van elke druppel regen in een storm te voorspellen. Normaal gesproken duurt dit dagen op een gewone computer.

Chromatix gebruikt de kracht van GPU's (de krachtige grafische kaarten die gamers gebruiken voor snelle video's). Het is alsof je in plaats van één persoon die de regen druppel voor druppel meet, een heel leger van duizenden mensen hebt die het tegelijk doen.

• Het resultaat: Wat voorheen uren duurde, duurt nu seconden. Soms zelfs 22 keer sneller.

3. De "Zelflerende" simulator

Dit is misschien wel het coolste deel. Chromatix is differentieerbaar. Dat klinkt ingewikkeld, maar het betekent simpelweg: de computer kan niet alleen simuleren wat er gebeurt, maar ook leren van zijn fouten.

Stel je voor dat je een blinddoek op hebt en probeert een doolhof te vinden. Normaal moet je blindelings lopen en hopen dat je eruit komt. Met Chromatix heeft de computer een "gevoel" voor welke kant hij moet draaien om sneller bij de uitgang te komen.

• In de praktijk: Als je een microscoop wilt bouwen die de beste foto's maakt, kan de computer automatisch duizenden variaties van lenzen en spiegels testen, zien welke het beste werkt, en zichzelf verbeteren. Het is alsof je een robot hebt die ontwerpt en bouwt terwijl je toekijkt.

Waarom is dit belangrijk?

Voorheen moesten onderzoekers hun eigen simulators bouwen voor elk nieuw project. Dat was als elke kok die een taart wilde bakken, zijn eigen oven en bloem moest uitvinden.

Met Chromatix:

• Snelheid: Onderzoekers kunnen nu in plaats van weken, dagen of uren besteden aan het ontwerpen van hun experimenten.
• Samenwerking: Iedereen gebruikt dezelfde "LEGO-blokken". Als iemand een nieuw blok bedenkt, kan iedereen het direct gebruiken.
• Nieuwe ontdekkingen: Omdat het zo snel en makkelijk is, kunnen wetenschappers nu dingen proberen die ze vroeger te bang of te traag vonden om te proberen. Denk aan het maken van 3D-foto's van hele hersenen in één klap, of het besturen van zenuwcellen met licht (optogenetica) zonder dat het licht verstrooit in het weefsel.

Kortom: Chromatix is de "iPhone" voor computergestuurde microscopie. Het maakt complexe optica toegankelijk, snel en makkelijk te gebruiken, zodat wetenschappers zich kunnen focussen op het ontdekken van nieuwe dingen in het leven, in plaats van vast te zitten aan de wiskunde.

Technische samenvatting

Probleemstelling

Moderne microscopie-methoden zijn steeds meer afhankelijk van computationele modellering, waarbij optische systeemontwerp en algoritmen voor reconstructie of optimalisatie geïntegreerd worden. Deze "model-based" benaderingen vereisen vaak differentieerbare simulaties van golfoptica (wave optics) om inverse problemen op te lossen of het ontwerp van het optische systeem te optimaliseren.

Huidige uitdagingen in dit veld zijn:

• Gebrek aan standaardisatie: Onderzoekers moeten voor elk project opnieuw simulaties van nul opbouwen, wat leidt tot beperkte herbruikbaarheid en inconsistenties in conventies.
• Rekenkracht en differentieerbaarheid: Bestaande tools (zoals Zemax of CODE V) zijn niet efficiënt geïmplementeerd voor de specifieke golf-optica modellen die nodig zijn voor moderne computationele optica, en zijn vaak niet differentieerbaar of compatibel met diepe neurale netwerken.
• Schaalbaarheid: Het programmeren van snelle golf-optica simulaties op moderne hardware (zoals GPU's) is complex. Bestaande open-source bibliotheken zijn vaak gebaseerd op straaloptica (ray-based) in plaats van golfoptica, of missen ondersteuning voor complexe scenario's zoals meervoudige verstrooiing in 3D-monsters en polarisatie.

Methodologie: Chromatix

De auteurs presenteren Chromatix, een open-source, GPU-versnelde, differentieerbare bibliotheek voor golf-optica simulaties. De kern van de methode is gebaseerd op de principes van moderne deep learning-frameworks (zoals PyTorch en TensorFlow), maar toegepast op optica.

Kernarchitectuur en principes:

1. Differentieerbaarheid (Differentiability): Chromatix is gebouwd op JAX (Just After eXecution), een numerieke rekenbibliotheek voor Python. Dit zorgt voor automatische differentiatie (AD), waardoor gradiënten automatisch worden berekend voor elke parameter in een simulatie. Dit maakt gradient-based optimalisatie mogelijk voor complexe fysieke modellen (bijv. verstrooiing of sample-deformatie) zonder handmatig afgeleiden te moeten berekenen.
2. Composabiliteit (Composability): Net als neurale netwerken zijn optische systemen in Chromatix opgebouwd uit een reeks fundamentele elementen (linsen, sensoren, SLM's, vrije ruimte-propagatie). Onderzoekers kunnen deze elementen als bouwstenen combineren om complexe systemen te modelleren. De "hidden state" is het complexe lichtveld (amplitude en fase), inclusief golflengte, polarisatie en ruimtelijke sampling.
3. Schaalbaarheid (Scalability): Dankzij JAX en de onderliggende XLA-compiler (Accelerated Linear Algebra) kan Chromatix code naadloos versnellen op CPU's, GPU's en TPUs. Het ondersteunt automatisch vectoriseren (batch-verwerking op één GPU) en parallelisatie over meerdere GPU's met minimale code-aanpassingen.

Fysieke modellen:
De bibliotheek implementeert diverse modellen, waaronder:

• Scalar en vectoriële golfpropagatie (vrije ruimte, verstrooiende media).
• Meervoudige verstrooiing in 3D-monsters (via Multislice Beam Propagation).
• Diffractieve elementen zoals Spatial Light Modulators (SLM) en DMD's.
• Zowel coherente als incoherente lichtbronnen.

Belangrijkste Bijdragen

1. Eerste standaardbibliotheek: Chromatix biedt de eerste open-source, differentieerbare bibliotheek die specifiek is ontworpen voor golf-optica simulaties met ondersteuning voor GPU-parallelisatie en complexe 3D-verstrooiende monsters.
2. Democratisering van complexe simulaties: Door een gestandaardiseerde interface te bieden, kunnen onderzoekers zonder diepgaande kennis van lage-niveau GPU-programmering complexe optische systemen modelleren en optimaliseren.
3. Integratie met Deep Learning: De bibliotheek maakt naadloze integratie toe tussen fysieke optische modellen en neurale netwerken, wat end-to-end optimalisatie van zowel het hardware-ontwerp als de reconstructie-algoritmen mogelijk maakt.

Resultaten en Experimenten

De auteurs demonstreren de kracht van Chromatix via zes experimenten die vier hoofdcategorieën bestrijken:

1. Oplossen van inverse problemen voor monsterreconstructie:

   • Ring deconvolutie: Voor een Miniscope met GRIN-lens werden ruimtelijk variërende aberraties gecorrigeerd. Chromatix kon het volledige beeldveld (2048x2048 pixels) verwerken door parallelisatie over meerdere GPU's, terwijl de originele PyTorch-implementatie vastliep op het geheugen van één GPU. Snelheidswinst: 4,5x (1 GPU) tot 19x (8 GPU's).
   • Computational Aberration Correction (COCOA): Gebruikmakend van Implicit Neural Representations (INR) voor 3D-reconstructie. Chromatix leverde betere resultaten op (gladdere dendrieten) en was 2x sneller (1 GPU) tot 9x sneller (8 GPU's).
   • 3D Refractive Index Microscopy: Reconstructie van de brekingsindex van een sterk verstrooiend D. rerio embryo. De Chromatix-implementatie was 3-13x sneller dan de originele MATLAB-code, wat leidde tot minder rasterartefacten en een snellere convergentie.

2. Programmeerbare optica en Deep Learning:

   • 3D Snapshot Microscopie (Holoscope): End-to-end optimalisatie van een fase-masker (SLM) en een reconstructie-netwerk (FourierNet) voor volumetrische imaging. Chromatix reduceerde de trainingstijd van weken naar dagen (~7x sneller) met vergelijkbare reconstructiekwaliteit.
   • DeepCGH (Computer Generated Holography): Een neurale netwerkbased aanpak om hologrammen in één doorloop te genereren. Chromatix bood een 2,5x snelheidswinst (1 GPU) tot >10x (8 GPU's) ten opzichte van de originele TensorFlow-implementatie.

3. Flexibele modellering:

   • Spectroscopische Single-Molecule Localization: Optimalisatie van een PSF voor multikleurige imaging vanuit één kanaal.
   • Holografie door verstrooiend weefsel: Combinatie van holografie en verstrooiingsmodellen om optogenetische stimulatiepatronen te corrigeren voor verstrooiing in biologisch weefsel, wat resulteerde in uniforme stimulatie.

4. Prestaties:

   • Over alle geteste methoden toonde Chromatix een snelheidswinst van 2-6x op één GPU en tot 22x op 8 GPU's in vergelijking met bestaande implementaties (MATLAB, PyTorch, TensorFlow).

Betekenis en Impact

Chromatix vult een cruciale lacune in het veld van computationele optica. Door een gestandaardiseerde, hoogpresterende en differentieerbare infrastructuur te bieden, verandert het de manier waarop onderzoekers optische systemen ontwerpen en analyseren:

• Versnelling van innovatie: Onderzoekers kunnen nu sneller itereren op nieuwe ideeën en complexe modellen testen die voorheen te rekenintensief waren.
• End-to-end optimalisatie: Het maakt het mogelijk om het ontwerp van het optische systeem (hardware) en de beeldverwerking (software) gelijktijdig te optimaliseren, wat leidt tot superieure prestaties.
• Toekomstperspectief: De auteurs zien potentie voor toepassingen zoals "optical architecture search" (automatisch ontwerpen van lenzen en optica) en hardware-in-the-loop kalibratie.

Kortom, Chromatix democratiseert hoogwaardige golf-optica simulaties en stelt de gemeenschap in staat om een veel rijker ontwerpruimte te verkennen voor biologische imaging en andere toepassingen.