Package: jax-dips
JAX-DIPSJAX implementation of a differentiable PDE solver with jump conditions across irregular interfaces in 3D.JAX-DIPS implements the neural bootstrapping method (NBM):Streamlines of solution gradients (left), and jump in solution (right) calculated by thedragonexample.TestingDopytest tests/test_*.pyof each of the available tests from the parent directory:test_advection: a sphere is rotated 360 degrees around the box to replicate initial configuration. The L2 error in level-set function should be less than 1e-4 to pass. The advection is performed using semi-Lagrangian scheme with Sussman reinitialization.test_reinitialization: starting from a sphere level-set function with -1 inside sphere and +1 outside, we repeatedly perform Sussman reinitialization until the signed-distance property of the level-set is achieved. Center of the box should have level-set value equal to radius of the sphere, and corner of the box should be at a pre-specified distance to pass.test_geometric_integrations: integrating surface area of a sphere along with its volume. Small differences with associated theoretical values are expected to pass.test_poisson: tests for both the pointwise and the grid-based Poisson solvers over a star and a sphere interfaces. Note that in the current implementation the grid-based solver does not support batching and is therefore faster. Fixing this issue will be done in the future versions.InstallationTo install the latest released version from PyPI dopip install jax-dips.Development & Usage Environment1. Virtual EnvironmentCreate a virtual environment by running the following command./create_virtualenv.shand theenv_jax_dipsvirtual environment will be created. Then you can launch into this environment bysource env_jax_dips/bin/activate. After you are done,deactivate.2. DockerPrerequisitesFirst you will need to install the nvidia driver and docker engines:Nvidia Driverapt install $(nvidia-detector)Docker Engine
Please referhttps://docs.docker.com/engine/install/ubuntu/nvidia-docker2
Please referhttps://nvidia.github.io/nvidia-docker/to setup apt repoapt-get install nvidia-docker2InstructionsDocker images provide an isolated and consistent runtime environment, ensuring that the application behaves the same regardless of the host system. We recommend using the docker image provided here as it is fully loaded with libraries for datacenter scale simulatiopns and optimized for NVIDIA GPUs. For a full list of the supported software and specific versions that come packaged with this container image see the Frameworks Support Matrixhttps://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.htmlTo personalize your development environment you need to set up a.envfile that containsIMAGE_NAME=docker.io/pourion/jax_dips:latest       # default docker image available for download! Change this if you want to build new docker images and push to your preferred docker registry
DATA_PATH=/data/                                   # default data path inside your docker container, will mirror your DATA_MOUNT_PATH directory
DATA_MOUNT_PATH=/data                              # default data mount path inside your machine, will mirror into your docker container's DATA_PATH directory
RESULT_PATH=/results/                              # default result path inside your docker container
RESULT_MOUNT_PATH=/results/                        # default result mount path inside your machine
REGISTRY=<your-preferred-registry-name>            # (optional) choices are docker.io (default), nvcr.io, ghcr.io, etc.
REGISTRY_USER=<your-registry-username>             # (optional) your username to connect to docker registry
REGISTRY_ACCESS_TOKEN=<your-registry-access-token> # (optional) your access token to connect to docker registry
WANDB_API_KEY=NotSpecified                         # (optional) your API key to connect to Weights and Biases service
JUPYTER_PORT=8888                                  # (optional) port to connect to jupyter serverPull development containerCurrently the latest docker image available on Docker Hub is available atdocker.io/pourion/jax_dips. Instead of building the container, you can only pull the latest docker image by running./launch.sh pullwhich pulls from docker hub; i.e., equivalent to$ docker pull pourion/jax_dips:latest.Build development containerAlternatively you can build the container by running the following command./launch.sh buildIn case you want to add additional libraries to your container this is the recommended way.Start developement containerThis will create a container and places the user in the container with source code mounted../launch.sh devAdditionally, you can run the container in background without having your terminal jump into the container. This can be done by passing the-dflag for daemon:./launch.sh dev -dYou can always attach your teminal to the runningjax_dipscontainer by./launch.sh attachDevelopment in VS CodeOnce the container is created and is running on your machine, user can attach to this container from VS code; i.e., you need to install Microsoft'sDev Containersextension in your VS Code, thenCtrl+Shift+Pand chooseDev Containers: Attach to Running Container..., then choose thejax_dipscontainer from the list of running containers on your machine.Cite JAX-DIPSIf you use JAX-DIPS in your research please use the following citations:@article{mistani2022jax,title={JAX-DIPS: Neural bootstrapping of finite discretization methods and application to elliptic problems with discontinuities},author={Mistani, Pouria and Pakravan, Samira and Ilango, Rajesh and Gibou, Frederic},journal={arXiv preprint arXiv:2210.14312},year={2022}}@article{mistani2022neuro,title={Neuro-symbolic partial differential equation solver},author={Mistani, Pouria and Pakravan, Samira and Ilango, Rajesh and Choudhry, Sanjay and Gibou, Frederic},journal={arXiv preprint arXiv:2210.14907},year={2022}}Contributing to JAX-DIPSReporting bugs.To report a bug please open an issue in the GitHubIssues.Suggesting enhancements.To submit an enhancement suggestion, including new features or improvements to existing functionality, let us know by opening an issue in the GitHubIssues.Pull requests.If you made improvements to JAX-DIPS, fixed a bug, or added a new example, feel free to send us a pull-request.The TeamJAX-DIPS was developed byPouria MistaniandSamira PakravanandRajesh Ilangounder the supervision of Prof.Frederic Gibouduring 2019-2022 atUniversity of California Santa Barbara. This project was partially funded by theUS Office of Naval Research.LicenseLGPL-2.1 License