{"id":13833072,"url":"https://github.com/sail-sg/jax_xc","last_synced_at":"2025-07-09T17:09:26.979Z","repository":{"id":65010109,"uuid":"578105369","full_name":"sail-sg/jax_xc","owner":"sail-sg","description":"Exchange correlation functionals translated from libxc to jax","archived":false,"fork":false,"pushed_at":"2025-03-24T16:59:46.000Z","size":681,"stargazers_count":45,"open_issues_count":4,"forks_count":2,"subscribers_count":6,"default_branch":"main","last_synced_at":"2025-06-09T15:52:58.741Z","etag":null,"topics":["density-functional-theory","jax","libxc","python"],"latest_commit_sha":null,"homepage":"https://jax-xc.readthedocs.io/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mpl-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/sail-sg.png","metadata":{"files":{"readme":"README.rst","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2022-12-14T09:16:19.000Z","updated_at":"2025-03-26T17:14:21.000Z","dependencies_parsed_at":"2024-01-15T17:39:27.111Z","dependency_job_id":null,"html_url":"https://github.com/sail-sg/jax_xc","commit_stats":null,"previous_names":[],"tags_count":10,"template":false,"template_full_name":null,"purl":"pkg:github/sail-sg/jax_xc","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sail-sg%2Fjax_xc","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sail-sg%2Fjax_xc/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sail-sg%2Fjax_xc/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sail-sg%2Fjax_xc/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/sail-sg","download_url":"https://codeload.github.com/sail-sg/jax_xc/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sail-sg%2Fjax_xc/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":264502181,"owners_count":23618558,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["density-functional-theory","jax","libxc","python"],"created_at":"2024-08-04T11:00:38.812Z","updated_at":"2025-07-09T17:09:26.341Z","avatar_url":"https://github.com/sail-sg.png","language":"Python","funding_links":[],"categories":["programs"],"sub_categories":[],"readme":"JAX Exchange Correlation Library\n================================\n\n.. image:: https://github.com/sail-sg/jax_xc/raw/main/figures/logo.png\n   :width: 200\n   :align: center\n   \n\n.. image:: https://img.shields.io/pypi/v/jax-xc.svg\n   :target: https://pypi.org/project/jax-xc/\n\n.. image:: https://readthedocs.org/projects/ansicolortags/badge/?version=latest\n   :target: https://jax-xc.readthedocs.io/en/latest/\n\n\n\nThis library contains direct translations of exchange correlation\nfunctionals in `libxc \u003chttps://tddft.org/programs/libxc/\u003e`__ to\n`jax \u003chttps://github.com/google/jax\u003e`__. The core calculations in libxc\nare implemented in `maple \u003chttps://www.maplesoft.com/\u003e`__. This gives us\nthe opportunity to translate them directly into python with the help of\n`CodeGeneration \u003chttps://www.maplesoft.com/support/help/maple/view.aspx?path=CodeGeneration%2fPython\u003e`__.\n\nUsage\n-----\n\nInstallation\n~~~~~~~~~~~~\n\n.. code:: sh\n\n   pip install jax-xc\n\nInvoking the Functionals\n~~~~~~~~~~~~~~~~~~~~~~~~\n\njax_xc's API is functional: it receives $\\\\rho$ a function of ``Callable``\ntype, and returns the $\\\\varepsilon_{xc}$ as a function of ``Callable``\ntype.\n\n.. code:: math\n\n   E_{xc} = \\int \\rho(r) \\varepsilon_{xc}(r) dr\n\nLDA and GGA\n^^^^^^^^^^^\n\nUnlike ``libxc`` which takes pre-computed densities and their derivative\nat certain coordinates. In ``jax_xc``, the API is designed to directly\ntake a density function.\n\n.. code:: python\n\n  import jax\n  import jax.numpy as jnp\n  import jax_xc\n\n\n  def rho(r):\n    \"\"\"Electron number density. We take gaussian as an example.\n\n    A function that takes a real coordinate, and returns a scalar\n    indicating the number density of electron at coordinate r.\n\n    Args:\n    r: a 3D coordinate.\n    Returns:\n    rho: If it is unpolarized, it is a scalar.\n        If it is polarized, it is a array of shape (2,).\n    \"\"\"\n    return jnp.prod(jax.scipy.stats.norm.pdf(r, loc=0, scale=1))\n\n  # create a density functional\n  gga_xc_pbe = jax_xc.gga_x_pbe(polarized=False)\n\n  # a grid point in 3D\n  r = jnp.array([0.1, 0.2, 0.3])\n\n  # pass rho and r to the functional to compute epsilon_xc (energy density) at r.\n  # corresponding to the 'zk' in libxc\n  epsilon_xc_r = gga_xc_pbe(rho, r)\n  print(epsilon_xc_r)\n\nmGGA\n^^^^\n\nUnlike LDA and GGA that only depends on the density function, mGGA\nfunctionals also depend on the molecular orbitals.\n\n.. code:: python\n\n  import jax\n  import jax.numpy as jnp\n  import jax_xc\n  \n\n  def mo(r):\n    \"\"\"Molecular orbital. We take gaussian as an example.\n\n    A function that takes a real coordinate, and returns the value of\n    molecular orbital at this coordinate.\n\n    Args:\n      r: a 3D coordinate.\n    Returns:\n      mo: If it is unpolarized, it is a array of shape (N,).\n          If it is polarized, it is a array of shape (N, 2).\n    \"\"\"\n    # Assume we have 3 molecular orbitals\n    return jnp.array([\n        jnp.prod(jax.scipy.stats.norm.pdf(r, loc=0, scale=1)),\n        jnp.prod(jax.scipy.stats.norm.pdf(r, loc=0.5, scale=1)),\n        jnp.prod(jax.scipy.stats.norm.pdf(r, loc=-0.5, scale=1))\n    ])\n\n\n  rho = lambda r: jnp.sum(mo(r)**2, axis=0)\n  mgga_xc_cc06 = jax_xc.mgga_xc_cc06(polarized=False)\n\n  # a grid point in 3D\n  r = jnp.array([0.1, 0.2, 0.3])\n\n  # evaluate the exchange correlation energy per particle at this point\n  # corresponding to the 'zk' in libxc\n  print(mgga_xc_cc06(rho, r, mo))\n\nHybrid Functionals\n^^^^^^^^^^^^^^^^^^\n\nHybrid functionals expose the same API, with extra attributes for the\nusers to access parameters needed outside of libxc/jax_xc (e.g. the\nfraction of exact exchange).\n\n.. code:: python\n\n  import jax\n  import jax.numpy as jnp\n  import jax_xc\n\n\n  def rho(r):\n    \"\"\"Electron number density. We take gaussian as an example.\n\n    A function that takes a real coordinate, and returns a scalar\n    indicating the number density of electron at coordinate r.\n\n    Args:\n      r: a 3D coordinate.\n    Returns:\n      rho: If it is unpolarized, it is a scalar.\n          If it is polarized, it is a array of shape (2,).\n    \"\"\"\n    return jnp.prod(jax.scipy.stats.norm.pdf(r, loc=0, scale=1))\n\n\n  hyb_gga_xc_pbeb0 = jax_xc.hyb_gga_xc_pbeb0(polarized=False)\n\n  # a grid point in 3D\n  r = jnp.array([0.1, 0.2, 0.3])\n\n  # evaluate the exchange correlation energy per particle at this point\n  # corresponding to the 'zk' in libxc\n  print(hyb_gga_xc_pbeb0(rho, r))\n\n  # access to extra attributes\n  cam_alpha = hyb_gga_xc_pbep0.cam_alpha  # fraction of full Hartree-Fock exchange\n\nThe complete list of extra attributes can be found below:\n\n.. code:: python\n\n     cam_alpha: float\n     cam_beta: float\n     cam_omega: float\n     nlc_b: float\n     nlc_C: float\n\nThe meaning for each attribute is the same as libxc:\n\n-  cam_alpha: fraction of full Hartree-Fock exchange, used both for\n   usual hybrids as well as range-separated ones\n-  cam_beta: fraction of short-range only(!) exchange in range-separated\n   hybrids\n-  cam_omega: range separation constant\n-  nlc_b: non-local correlation, b parameter\n-  nlc_C: non-local correlation, C parameter\n\nExperimental\n------------\n\nWe support automatic functional derivative!\n\n.. code:: python\n\n  import jax\n  import jax_xc\n  import autofd.operators as o\n  from autofd import function\n  import jax.numpy as jnp\n  from jaxtyping import Array, Float32\n\n  @function\n  def rho(r: Float32[Array, \"3\"]) -\u003e Float32[Array, \"\"]:\n    \"\"\"Electron number density. We take gaussian as an example.\n\n    A function that takes a real coordinate, and returns a scalar\n    indicating the number density of electron at coordinate r.\n\n    Args:\n    r: a 3D coordinate.\n    Returns:\n    rho: If it is unpolarized, it is a scalar.\n        If it is polarized, it is a array of shape (2,).\n    \"\"\"\n    return jnp.prod(jax.scipy.stats.norm.pdf(r, loc=0, scale=1))\n\n  # create a density functional\n  gga_x_pbe = jax_xc.experimental.gga_x_pbe\n  epsilon_xc = gga_x_pbe(rho)\n\n  # a grid point in 3D\n  r = jnp.array([0.1, 0.2, 0.3])\n\n  # pass rho and r to the functional to compute epsilon_xc (energy density) at r.\n  # corresponding to the 'zk' in libxc\n  print(f\"The function signature of epsilon_xc is {epsilon_xc}\")\n\n  energy_density = epsilon_xc(r)\n  print(f\"epsilon_xc(r) = {energy_density}\")\n\n  vxc = jax.grad(lambda rho: o.integrate(rho * gga_x_pbe(rho)))(rho)\n  print(f\"The function signature of vxc is {vxc}\")\n  print(vxc(r))\n\n\nSupport Functionals\n-------------------\n\nPlease refer to the `functionals section \u003chttps://jax-xc.readthedocs.io/en/latest/sources/jax_xc.html#module-jax_xc.functionals\u003e`_ \nin ``jax_xc``'s documentation\nfor the complete list of supported functionals.\n\nNumerical Correctness\n---------------------\n\nWe test all the functionals that are auto-generated from maple files\nagainst the reference values in ``libxc``. The test is performed by\ncomparing the output of ``libxc`` and ``jax_xc`` and make sure they are\nwithin a certain tolerance, namely ``atol=2e-10`` and ``rtol=2e-10``.\n\nPerformance Benchmark\n---------------------\n\nWe report the performance benchmark of ``jax_xc`` against ``libxc`` on a\n64-core machine with Intel(R) Xeon(R) Silver 4216 CPU @ 2.10GHz.\n\nWe sample the points to evaluate the functionals by varying the number\nof points from 1 to $10^7$. The benchmark is performed by evaluating the\nruntime of the functional. Note that the runtime of ``jax_xc`` is\nmeasured by excluding the time of just-in-time compilation.\n\nWe visualize the mean value (averaged for both polarized and unpolarized)\nof the runtime of ``jax_xc`` and ``libxc`` in the following figure. The\ny-axis is log-scale. \n\n``jax_xc``'s runtime is constantly below ``libxc``'s\nfor all batch sizes. The speed up is ranging from 3x to 10x, and it is\nmore significant for larger batch sizes. \n\nWe hypothesize that the reason\nfor the speed up is that Jax's JIT compiler is able to optimize the\nfunctionals (e.g. vectorization, parallel execution, instruction fusion, \nconstant folding for floating points, etc.) better than\nlibxc.\n\n.. image:: https://raw.githubusercontent.com/sail-sg/jax_xc/main/figures/jax_xc_speed.svg\n\nWe visualize the distribution of the runtime ratio of ``jax_xc`` and\n``libxc`` in the following figure. The ratio is closer to 0.1 for\nlarge batch sizes (~ 10x speed up). The ratio is constantly below 1.0.\n\n.. image:: https://raw.githubusercontent.com/sail-sg/jax_xc/main/figures/jax_xc_ratio.svg\n\nNote that, we exclude one datapoint ``mgga_x_2d_prhg07`` from the\nruntime ratio visualization because it is an outlier due to Jax's lack\nof support of\\ ``lamberw`` function and we use\n``tensorflow_probability.substrates.jax.math.lambertw``.\n\nCaveates\n--------\n\nThe following functionals from ``libxc`` are not available in ``jax_xc``\nbecause some functions are not available in ``jax``.\n\n.. code:: python\n\n   gga_x_fd_lb94          # Becke-Roussel not having a closed-form expression\n   gga_x_fd_revlb94       # Becke-Roussel not having a closed-form expression\n   gga_x_gg99             # Becke-Roussel not having a closed-form expression\n   gga_x_kgg99            # Becke-Roussel not having a closed-form expression\n   hyb_gga_xc_case21      # Becke-Roussel not having a closed-form expression\n   hyb_mgga_xc_b94_hyb    # Becke-Roussel not having a closed-form expression\n   hyb_mgga_xc_br3p86     # Becke-Roussel not having a closed-form expression\n   lda_x_1d_exponential   # Requires explicit 1D integration\n   lda_x_1d_soft          # Requires explicit 1D integration\n   mgga_c_b94             # Becke-Roussel not having a closed-form expression\n   mgga_x_b00             # Becke-Roussel not having a closed-form expression\n   mgga_x_bj06            # Becke-Roussel not having a closed-form expression\n   mgga_x_br89            # Becke-Roussel not having a closed-form expression\n   mgga_x_br89_1          # Becke-Roussel not having a closed-form expression\n   mgga_x_mbr             # Becke-Roussel not having a closed-form expression\n   mgga_x_mbrxc_bg        # Becke-Roussel not having a closed-form expression\n   mgga_x_mbrxh_bg        # Becke-Roussel not having a closed-form expression\n   mgga_x_mggac           # Becke-Roussel not having a closed-form expression\n   mgga_x_rpp09           # Becke-Roussel not having a closed-form expression\n   mgga_x_tb09            # Becke-Roussel not having a closed-form expression\n   gga_x_wpbeh            # jit too long for E1_scaled\n   gga_c_ft97             # jit too long for E1_scaled\n   lda_xc_tih             # vxc functional\n   gga_c_pbe_jrgx         # vxc functional\n   gga_x_lb               # vxc functional\n\nBuilding from Source Code\n-------------------------\n\nModify the ``.env.example`` to fill in your envrionment variables, then\nrename it to ``.env``. Then run ``source .env`` to load them into your\nshell.\n\n-  ``OUTPUT_USER_ROOT``: The path to the bazel cache. This is where the\n   bazel cache will be stored. This is useful if you are building on a\n   shared filesystem.\n\n-  ``MAPLE_PATH``: The path to the maple binary.\n\n-  ``TMP_INSTALL_PATH``: The path to a temporary directory where the\n   wheel will be installed. This is useful if you are building on a\n   shared filesystem.\n\nMake sure you have ``bazel`` and ``maple`` installed. Your python envrionment has installed the dependencies in\n``requirements.txt``.\n\nHow to build python wheel.\n\n.. code:: sh\n\n   bazel --output_user_root=$OUTPUT_USER_ROOT build --action_env=PATH=$PATH:$MAPLE_PATH @jax_xc_repo//:jax_xc_wheel\n\nOnce the build finished, the python wheel could be found under ``bazel-bin/external/jax_xc_repo``. For example, the\nname for version 0.0.7 is ``jax_xc-0.0.7-cp310-cp310-manylinux_2_17_x86_64.whl``.\n\nInstall the python wheel. If needed, specify the install path by\n\n.. code:: sh\n\n   pip install {{wheel_name}} --target $TMP_INSTALL_PATH\n\nRunning Test\n------------\n\nThe test could be run without the command above that builds wheel from source, though it might take longer time to\nbuild all the components needed for the test. To run all the test:\n\n.. code:: sh\n\n   bazel --output_user_root=$OUTPUT_USER_ROOT test --action_env=PATH=$PATH:$MAPLE_PATH //tests/...\n\nTo run a specific test, for example ``test_impl``:\n\n.. code:: sh\n\n   bazel --output_user_root=$OUTPUT_USER_ROOT test --action_env=PATH=$PATH:$MAPLE_PATH //tests:test_impl\n\nThe test output could be found in ``bazel-testlogs/tests/test_impl/test.log`` for the ``tests:test_impl`` and similar to\nthe others. If you prefer output in command line, add ``--test_output=all`` to the above command.\n\nLicense\n-------\n\nAligned with ``libxc``, ``jax_xc`` is licensed under the Mozilla Public License 2.0. See\n``LICENSE`` for the full license text.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsail-sg%2Fjax_xc","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsail-sg%2Fjax_xc","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsail-sg%2Fjax_xc/lists"}