Package: galilei
galileigalileiis a python package that makes emulating a numerical functions easier and more composable. It supports multiple backends such as pytorch-based neural networks, GPy-based gaussian process regression, etc. As of now, it defaults to a jax+flax+optax backend which supports automatic differenciation of the emulated function and easy composibility with the rest of the jax-based eco-system.The motivation of emulating a function is that sometimes computing a function could be a time consuming task, so one may need to find fast approximations of a function that's better than basic interpolation techniques. Machine learning techniques such as neural networks offer a solution to this problem which is generic enough to emulate any arbitrarily shaped function. In contrast to the original function, a neural-network based emulator runs blazingly fast and even more so with GPU, often achieveing over many orders of magnitude speed-up.This idea of emulating function is not new. In the field of cosmology we have powerful tools such ascosmopowerand its derived works such asaxionEmu, whose idea inspired this work. My aim in this work differs from the previous approach, as I intend to make a both generic and easily-composible function emulator that can take any arbitrary parametrized numerical function as an input, and return a function with the exact same signature as a drop-in replacement in existing code base, with no additional code changes needed. In addition, I also focus on making the emulated function automatically differenciable regardless of its original implementation.FeaturesSupport multiple backends:torch,sklearn,gpy(for gaussian process regression),jax.Flexible: Able to emulate generic numerical functions.Automatic differenciable (supported by selected backends): emulated function is automatically differenciable and easily composible with jax-based tools.Easy to use: just add a decorator@emulateand use your emulated function as a drop-in replacement of your existing function in code-base without additional modification.Allow arbitrary transformation of function output before training through the use ofPreconditioner.Installationpip install galileiBasic usageSuppose that we have an expensive function that we want to emulatedefmyfun(a,b):# assume this is a slow functionx=np.linspace(0,10,100)returnnp.sin(a*x)+np.sin(b*x)If you want to emulate this function, you can simply add a decorator@emulateand supply the parameters that you want to evaluate this function at to build up the training data set.fromgalileiimportemulate@emulate(samples={'a':np.random.rand(1000),'b':np.random.rand(1000)})defmyfun(a,b):x=np.linspace(0,10,100)returnnp.sin(a*x)+np.sin(b*x)Here we are just making 1000 pairs of random numbers from 0 to 1 to train our function. When executing these lines, the emulator will start training, and once it is done, the originalmyfunfunction will be automatically replaced with the emulated version and should behave in the same way, except much faster!Training emulator...
100%|██████████| 500/500 [00:09<00:00, 50.50it/s, loss=0.023]
Ave Test loss: 0.025With the defaultjaxbackend, the emulated function is automatically jax compatible, which means one can easily compose them withjaxmachinary, such as in example below where I have compiled the emulated function withjitand then vectorized it over its first argument withvmap.fromjaximportjit,vmapvmap(jit(myfun),in_axes=(0,None))(np.linspace(0,1,10),0.5)Output:Array([[-2.39813775e-02,  2.16133818e-02,  8.05920288e-02,
         1.66035295e-01,  2.01425016e-01,  2.42054626e-01,
         2.74079561e-01,  3.50277930e-01,  4.12616253e-01,
         4.33193207e-01,  4.82740909e-01,  5.66871405e-01,
         5.73131263e-01,  6.51429832e-01,  6.55564785e-01,
         ...The emulated function will also be automatically differenciable regardless of the original implementation details. For example, we could easily evaluate its jacobian (without finite differencing) withfromjaximportjacfwdjacfwd(myfun,argnums=(0,1))(0.2,0.8)Output:(Array([ 0.05104188,  0.18436229,  0.08595917,  0.06582363,  0.17270228, ...],      dtype=float32),
 Array([-3.3511031e-01,  1.2647966e-01,  4.3209594e-02,  2.4325712e-01, ...],      dtype=float32))You can also easily save your trained model with thesaveoption@emulate(samples={'a':np.random.rand(100),'b':np.random.rand(100)},backend='sklearn',save="test.pkl")defmyfun(a,b):x=np.linspace(0,10,100)returnnp.sin(a*x)+np.sin(b*x)and when you use it in production, simply load a pretrained model with@emulate(backend='sklearn',load="test.pkl")defmyfun(a,b):x=np.linspace(0,10,100)returnnp.sin(a*x)+np.sin(b*x)and your function will be replaced with a fast emulated version.It's also possible to sample training points based on latin hypercube using thebuild_samplesfunction. For example, here I build a 100 sample latin hypercube for a given range ofaandbfromgalilei.samplingimportbuild_samples@emulate(samples=build_samples({"a":[0,2],"b":[0,2]},100),backend='sklearn')defmyfun(a,b):x=np.linspace(0,10,100)returnnp.sin(a*x)+np.sin(b*x)Sometimes one might want to collect training data only instead of training the emulator. This could
be done byfromgalilei.experimentalimportcollect@collect(samples=build_samples({"a":[0,2],"b":[0,2]},100),save="collection.pkl",mpi=True)defmyfun(a,b):x=np.linspace(0,10,100)returnnp.sin(a*x)+np.sin(b*x)which will save a precomputed collection tocollection.pklfor future loading. Note that the option to usempidepends on the user having a workingmpi4pywhich needs to be installed by the user.
The collection could be loaded for training emulator using@emulate(collection="collection.pkl",backend='sklearn')defmyfunc(a,b):raiseException()myfunc(1,1)since the function will not be evaluated in this case, we note that the implementation ofmyfuncmakes no difference (otherwise it would have given an error).For more usage examples, see this notebook:RoadmapTODO add prebuild preconditionersTODO support downloading files from webTODO auto infer backendTODO chains of preconditionersCreditsThis package was created with theppwtool. For more information, please visit theproject page.If this package is helpful in your work, please consider citing:@article{yguan_2023,title={galilei: a generic function emulator},DOI={10.5281/zenodo.7651315},publisher={Zenodo},author={Yilun Guan},year={2023},month={Feb}}Free software: MIT