Package: jaxdf
jaxdf - JAX-based Discretization FrameworkOverview|Example|Installation|Documentation|SupportOverviewJaxdf is a package based onJAXthat provides a coding framework for creating differentiable numerical simulators with arbitrary discretizations.The primary objective of Jaxdf is to aid in the construction of numerical models for physical systems, like wave propagation, or the numerical resolution of partial differential equations, in a manner that is easily tailored to the user's research requirements. These models are pure functions that can be seamlessly integrated into arbitrary differentiable programs written inJAX. For instance, they can be employed as layers within neural networks, or utilized in constructing a physics loss function.ExampleThe script below constructs the non-linear operator(∇2+ sin), applying a Fourier spectral discretization on a square 2D domain. It then utilizes this operator to define a loss function. The gradient of this loss function is calculated using JAX's Automatic Differentiation.fromjaxdfimportoperatorsasjopsfromjaxdfimportFourierSeries,operatorfromjaxdf.geometryimportDomainfromjaximportnumpyasjnpfromjaximportjit,grad# Defining operator@operatordefcustom_op(u,*,params=None):grad_u=jops.gradient(u)diag_jacobian=jops.diag_jacobian(grad_u)laplacian=jops.sum_over_dims(diag_jacobian)sin_u=jops.compose(u)(jnp.sin)returnlaplacian+sin_u# Defining discretizationsdomain=Domain((128,128),(1.,1.))parameters=jnp.ones((128,128,1))u=FourierSeries(parameters,domain)# Define a differentiable loss function@jitdefloss(u):v=custom_op(u)returnjnp.mean(jnp.abs(v.on_grid)**2)gradient=grad(loss)(u)# gradient is a FourierSeriesInstallationBefore proceeding with the installation ofjaxdf, ensure thatJAX is already installedon your system. If you intend to utilizejaxdfwith NVidia GPU support, follow the instructions to install JAX accordingly.To installjaxdffrom PyPI, use thepipcommand:pipinstalljaxdfFor development purposes, installjaxdfby either cloning the repository or downloading and extracting the compressed archive. Afterward, navigate to the root folder in a terminal, and execute the following command:pipinstall--upgradepoetry
poetryinstallThis will install the dependencies and the package itself (in editable mode).SupportIf you encounter any issues with the code or wish to suggest new features, please feel free to open an issue. If you seek guidance, wish to discuss something, or simply want to say hi, don't hesitate to write a message in ourDiscord channel.ContributingContributions are absolutely welcome! Most contributions start with an issue. Please don't hesitate to create issues in which you ask for features, give feedback on performances, or simply want to reach out.To make a pull request, please look at the detailedContributing guidefor how to do it, but fundamentally keep in mind the following main guidelines:If you add a new feature or fix a bug:Make sure it is covered by testsAdd a line in the changelog usingkacl-cliIf you changed something in the documentation, make sure that the documentation site can be correctly build usingmkdocs serveCitationAn initial version of this package was presented at theDifferentiable Programming workshopat NeurIPS 2021.@article{stanziola2021jaxdf,author={Stanziola, Antonio and Arridge, Simon and Cox, Ben T. and Treeby, Bradley E.},title={A research framework for writing differentiable PDE discretizations in JAX},year={2021},journal={Differentiable Programming workshop at Neural Information Processing Systems 2021}}AcknowledgementsSome of the packaging of this repository is done by editingthis templace from @rochacbrunoThe multiple-dispatch method employed is based onplum, check out this amazing project:https://github.com/wesselb/plumRelated projectsodlOperator Discretization Library (ODL) is a python library for fast prototyping focusing on (but not restricted to) inverse problems.deepXDE: a TensorFlow and PyTorch library for scientific machine learning.SciML: SciML is a NumFOCUS sponsored open source software organization created to unify the packages for scientific machine learning.