Package: invrs-gym
invrs-gym - A collection of inverse design challengesv0.9.1OverviewTheinvrs_gympackage is an open-source gym containing a diverse set of photonic design challenges, which are relevant for a wide range of applications such as AR/VR, optical networking, LIDAR, and others.Each of the challenges consists of a high-dimensional problem in which a physical structure (the photonic device) is optimized. The structure includes typically >10,000 degrees of freedom (DoF), generally including one or more arrays representing the structure or patterning of a layer, and may also include scalar variables representing e.g. layer thickness. In general, the DoF must satisfy certain constraints to be physical: thicknesses must be positive, and layer patterns must bemanufacturable---they must not include features that are too small, or too closely spaced.In general, we seek optimization techniques thatreliablyproduce manufacturable, high-quality solutions and require reasonable compute resources. Among the techniques that could be applied are topology optimization, inverse design, and AI-guided design.invrs_gymis intended to facilitate research on such methods within the jax ecosystem. It includes several challenges that have been used in previous works, so that researchers may directly compare their results to those of the literature. While some challenges are test problems (e.g. where the structure is two-dimensional, which is unphysical but allows fast simulation), others are actual problems that are relevant e.g. for quantum computing or 3D sensing.Key conceptsThe key types of the challenge are theChallengeandComponentobjects.TheComponentrepresents the physical structure to be optimized, and has some intended excitation or operating condition (e.g. illumination with a particular wavelength from a particular direction). TheComponentincludes methods to obtain initial parameters, and to compute theresponseof a component to the excitation.EachChallengehas aComponentas an attribute, and also has a target that can be used to determine whether particular parameters "solve" the challenge. TheChallengealso provides functions to compute a scalar loss for use with gradient-based optimization, and additional metrics.Example# Select the challenge.challenge=invrs_gym.challenges.ceviche_lightweight_waveguide_bend()# Define loss function, which also returns auxilliary quantities.defloss_fn(params):response,aux=challenge.component.response(params)loss=challenge.loss(response)distance=challenge.distance_to_target(response)metrics=challenge.metrics(response,params,aux)returnloss,(response,distance,metrics,aux)value_and_grad_fn=jax.value_and_grad(loss_fn,has_aux=True)# Select an optimizer.opt=invrs_opt.density_lbfgsb(beta=4)# Generate initial parameters, and use these to initialize the optimizer state.params=challenge.component.init(jax.random.PRNGKey(0))state=opt.init(params)# Carry out the optimization.foriinrange(steps):params=opt.params(state)(value,(response,distance,metrics,aux)),grad=value_and_grad_fn(params)state=opt.update(grad=grad,value=value,params=params,state=state)With some plotting (see theexample notebook), this code will produce the following waveguide bend:ChallengesThe current list of challenges is below. Check out the notebooks for ready-to-go examples of each.Thecevichechallenges are jax-wrapped versions of theCeviche Challengesopen-sourced by Google, with defaults matchingInverse Design of Photonic Devices with Strict Foundry Fabrication Constraintsby Schubert et al. These were also studied by Ferber et al. inSurCo: Learning Linear SURrogates for COmbinatorial Nonlinear Optimization Problems.Themetagratingchallenge is a re-implementation of theMetagrating3Dproblem using thefmmaxsimulator.Thediffractive splitterchallenge involves designing a non-paraxial diffractive beamsplitter useful for 3D sensing, as discussed inLightTrans documentation.Thephoton extractorchallenge is based onInverse-designed photon extractors for optically addressable defect qubitsby Chakravarthi et al., and aims to create structures that increase photon extraction efficiency for quantum applications.Thepolarization sorterchallenge is based on thepolarization-sorting metasurface examplefrom theFMMAX simulator, and aims to create split normally-incident light into one of four subpixels in a pixel array depending upon the polarization.Installpip install invrs_gymTestingSome tests are marked as slow and are skipped by default. To run these manually, usepytest --runslow