Package: jax-xc
This library contains direct translations of exchange correlation
functionals inlibxctojax. The core calculations in libxc
are implemented inmaple. This gives us
the opportunity to translate them directly into python with the help ofCodeGeneration.UsageInstallationpipinstalljax-xcInvoking the Functionalsjax_xc’s API is functional: it receives $\rho$ a function ofCallabletype, and returns the $\varepsilon_{xc}$ as a function ofCallabletype.E_{xc} = \int \rho(r) \varepsilon_{xc}(r) drLDA and GGAUnlikelibxcwhich takes pre-computed densities and their derivative
at certain coordinates. Injax_xc, the API is designed to directly
take a density function.importjaximportjax.numpyasjnpimportjax_xcdefrho(r):"""Electron number density. We take gaussian as an example.

  A function that takes a real coordinate, and returns a scalar
  indicating the number density of electron at coordinate r.

  Args:
  r: a 3D coordinate.
  Returns:
  rho: If it is unpolarized, it is a scalar.
      If it is polarized, it is a array of shape (2,).
  """returnjnp.prod(jax.scipy.stats.norm.pdf(r,loc=0,scale=1))# create a density functionalgga_xc_pbe=jax_xc.gga_x_pbe(polarized=False)# a grid point in 3Dr=jnp.array([0.1,0.2,0.3])# pass rho and r to the functional to compute epsilon_xc (energy density) at r.# corresponding to the 'zk' in libxcepsilon_xc_r=gga_xc_pbe(rho,r)print(epsilon_xc_r)mGGAUnlike LDA and GGA that only depends on the density function, mGGA
functionals also depend on the molecular orbitals.importjaximportjax.numpyasjnpimportjax_xcdefmo(r):"""Molecular orbital. We take gaussian as an example.

  A function that takes a real coordinate, and returns the value of
  molecular orbital at this coordinate.

  Args:
    r: a 3D coordinate.
  Returns:
    mo: If it is unpolarized, it is a array of shape (N,).
        If it is polarized, it is a array of shape (N, 2).
  """# Assume we have 3 molecular orbitalsreturnjnp.array([jnp.prod(jax.scipy.stats.norm.pdf(r,loc=0,scale=1)),jnp.prod(jax.scipy.stats.norm.pdf(r,loc=0.5,scale=1)),jnp.prod(jax.scipy.stats.norm.pdf(r,loc=-0.5,scale=1))])rho=lambdar:jnp.sum(mo(r)**2,axis=0)mgga_xc_cc06=jax_xc.mgga_xc_cc06(polarized=False)# a grid point in 3Dr=jnp.array([0.1,0.2,0.3])# evaluate the exchange correlation energy per particle at this point# corresponding to the 'zk' in libxcprint(mgga_xc_cc06(rho,r,mo))Hybrid FunctionalsHybrid functionals expose the same API, with extra attributes for the
users to access parameters needed outside of libxc/jax_xc (e.g. the
fraction of exact exchange).importjaximportjax.numpyasjnpimportjax_xcdefrho(r):"""Electron number density. We take gaussian as an example.

  A function that takes a real coordinate, and returns a scalar
  indicating the number density of electron at coordinate r.

  Args:
    r: a 3D coordinate.
  Returns:
    rho: If it is unpolarized, it is a scalar.
        If it is polarized, it is a array of shape (2,).
  """returnjnp.prod(jax.scipy.stats.norm.pdf(r,loc=0,scale=1))hyb_gga_xc_pbeb0=jax_xc.hyb_gga_xc_pbeb0(polarized=False)# a grid point in 3Dr=jnp.array([0.1,0.2,0.3])# evaluate the exchange correlation energy per particle at this point# corresponding to the 'zk' in libxcprint(hyb_gga_xc_pbeb0(rho,r))# access to extra attributescam_alpha=hyb_gga_xc_pbep0.cam_alpha# fraction of full Hartree-Fock exchangeThe complete list of extra attributes can be found below:cam_alpha:floatcam_beta:floatcam_omega:floatnlc_b:floatnlc_C:floatThe meaning for each attribute is the same as libxc:cam_alpha: fraction of full Hartree-Fock exchange, used both for
usual hybrids as well as range-separated onescam_beta: fraction of short-range only(!) exchange in range-separated
hybridscam_omega: range separation constantnlc_b: non-local correlation, b parameternlc_C: non-local correlation, C parameterExperimentalWe support automatic functional derivative!importjaximportjax_xcimportautofd.operatorsasofromautofdimportfunctionimportjax.numpyasjnpfromjaxtypingimportArray,Float32@functiondefrho(r:Float32[Array,"3"])->Float32[Array,""]:"""Electron number density. We take gaussian as an example.

  A function that takes a real coordinate, and returns a scalar
  indicating the number density of electron at coordinate r.

  Args:
  r: a 3D coordinate.
  Returns:
  rho: If it is unpolarized, it is a scalar.
      If it is polarized, it is a array of shape (2,).
  """returnjnp.prod(jax.scipy.stats.norm.pdf(r,loc=0,scale=1))# create a density functionalgga_x_pbe=jax_xc.experimental.gga_x_pbeepsilon_xc=gga_x_pbe(rho)# a grid point in 3Dr=jnp.array([0.1,0.2,0.3])# pass rho and r to the functional to compute epsilon_xc (energy density) at r.# corresponding to the 'zk' in libxcprint(f"The function signature of epsilon_xc is{epsilon_xc}")energy_density=epsilon_xc(r)print(f"epsilon_xc(r) ={energy_density}")vxc=jax.grad(lambdarho:o.integrate(rho*gga_x_pbe(rho)))(rho)print(f"The function signature of vxc is{vxc}")print(vxc(r))Support FunctionalsPlease refer to thefunctionals sectioninjax_xc’s documentation
for the complete list of supported functionals.Numerical CorrectnessWe test all the functionals that are auto-generated from maple files
against the reference values inlibxc. The test is performed by
comparing the output oflibxcandjax_xcand make sure they are
within a certain tolerance, namelyatol=2e-10andrtol=2e-10.Performance BenchmarkWe report the performance benchmark ofjax_xcagainstlibxcon a
64-core machine with Intel(R) Xeon(R) Silver 4216 CPU @ 2.10GHz.We sample the points to evaluate the functionals by varying the number
of points from 1 to $10^7$. The benchmark is performed by evaluating the
runtime of the functional. Note that the runtime ofjax_xcis
measured by excluding the time of just-in-time compilation.We visualize the mean value (averaged for both polarized and unpolarized)
of the runtime ofjax_xcandlibxcin the following figure. The
y-axis is log-scale.jax_xc’s runtime is constantly belowlibxc’s
for all batch sizes. The speed up is ranging from 3x to 10x, and it is
more significant for larger batch sizes.We hypothesize that the reason
for the speed up is that Jax’s JIT compiler is able to optimize the
functionals (e.g. vectorization, parallel execution, instruction fusion,
constant folding for floating points, etc.) better than
libxc.We visualize the distribution of the runtime ratio ofjax_xcandlibxcin the following figure. The ratio is closer to 0.1 for
large batch sizes (~ 10x speed up). The ratio is constantly below 1.0.Note that, we exclude one datapointmgga_x_2d_prhg07from the
runtime ratio visualization because it is an outlier due to Jax’s lack
of support oflamberwfunction and we usetensorflow_probability.substrates.jax.math.lambertw.CaveatesThe following functionals fromlibxcare not available injax_xcbecause some functions are not available injax.gga_x_fd_lb94# Becke-Roussel not having a closed-form expressiongga_x_fd_revlb94# Becke-Roussel not having a closed-form expressiongga_x_gg99# Becke-Roussel not having a closed-form expressiongga_x_kgg99# Becke-Roussel not having a closed-form expressionhyb_gga_xc_case21# Becke-Roussel not having a closed-form expressionhyb_mgga_xc_b94_hyb# Becke-Roussel not having a closed-form expressionhyb_mgga_xc_br3p86# Becke-Roussel not having a closed-form expressionlda_x_1d_exponential# Requires explicit 1D integrationlda_x_1d_soft# Requires explicit 1D integrationmgga_c_b94# Becke-Roussel not having a closed-form expressionmgga_x_b00# Becke-Roussel not having a closed-form expressionmgga_x_bj06# Becke-Roussel not having a closed-form expressionmgga_x_br89# Becke-Roussel not having a closed-form expressionmgga_x_br89_1# Becke-Roussel not having a closed-form expressionmgga_x_mbr# Becke-Roussel not having a closed-form expressionmgga_x_mbrxc_bg# Becke-Roussel not having a closed-form expressionmgga_x_mbrxh_bg# Becke-Roussel not having a closed-form expressionmgga_x_mggac# Becke-Roussel not having a closed-form expressionmgga_x_rpp09# Becke-Roussel not having a closed-form expressionmgga_x_tb09# Becke-Roussel not having a closed-form expressiongga_x_wpbeh# jit too long for E1_scaledgga_c_ft97# jit too long for E1_scaledlda_xc_tih# vxc functionalgga_c_pbe_jrgx# vxc functionalgga_x_lb# vxc functionalBuilding from Source CodeModify the.env.exampleto fill in your envrionment variables, then
rename it to.env. Then runsource .envto load them into your
shell.OUTPUT_USER_ROOT: The path to the bazel cache. This is where the
bazel cache will be stored. This is useful if you are building on a
shared filesystem.MAPLE_PATH: The path to the maple binary.TMP_INSTALL_PATH: The path to a temporary directory where the
wheel will be installed. This is useful if you are building on a
shared filesystem.Make sure you havebazelandmapleinstalled. Your python envrionment has installed the dependencies inrequirements.txt.How to build python wheel.bazel--output_user_root=$OUTPUT_USER_ROOTbuild--action_env=PATH=$PATH:$MAPLE_PATH@jax_xc_repo//:jax_xc_wheelOnce the build finished, the python wheel could be found underbazel-bin/external/jax_xc_repo. For example, the
name for version 0.0.7 isjax_xc-0.0.7-cp310-cp310-manylinux_2_17_x86_64.whl.Install the python wheel. If needed, specify the install path bypipinstall{{wheel_name}}--target$TMP_INSTALL_PATHRunning TestThe test could be run without the command above that builds wheel from source, though it might take longer time to
build all the components needed for the test. To run all the test:bazel--output_user_root=$OUTPUT_USER_ROOTtest--action_env=PATH=$PATH:$MAPLE_PATH//tests/...To run a specific test, for exampletest_impl:bazel--output_user_root=$OUTPUT_USER_ROOTtest--action_env=PATH=$PATH:$MAPLE_PATH//tests:test_implThe test output could be found inbazel-testlogs/tests/test_impl/test.logfor thetests:test_impland similar to
the others. If you prefer output in command line, add--test_output=allto the above command.LicenseAligned withlibxc,jax_xcis licensed under the Mozilla Public License 2.0. SeeLICENSEfor the full license text.