Added density_functional_approximation_dm21 to deepmind-research repo

PiperOrigin-RevId: 415267993
2026-05-26 01:15:26 +08:00 · 2021-12-08 12:01:54 +00:00
parent d55816beb3
commit 72c72d530f
39 changed files with 2934 additions and 0 deletions
@@ -24,6 +24,7 @@ https://deepmind.com/research/publications/

 ## Projects

+*   [Pushing the Frontiers of Density Functionals by Solving the Fractional Electron Problem](density_functional_approximation_dm21), Science 2021
 *   [Mind the Gap: Assessing Temporal Generalization in Neural Language Models](pitfalls_static_language_models), NeurIPS 2021
 *   [The Difficulty of Passive Learning in Deep Reinforcement Learning](tandem_dqn), NeurIPS 2021
 *   [Skilful precipitation nowcasting using deep generative models of radar](nowcasting), Nature 2021
@@ -0,0 +1,125 @@
+"""DM21 functionals."""
+
+load("@org_tensorflow//tensorflow/python/tools:tools.bzl", "saved_model_compile_aot")
+load("@rules_python//python:defs.bzl", "py_binary", "py_library")
+load("@external_py_deps//:requirements.bzl", "requirement")
+
+licenses(["notice"])
+
+py_library(
+    name = "compute_hfx_density",
+    srcs = ["density_functional_approximation_dm21/compute_hfx_density.py"],
+    srcs_version = "PY3",
+    deps = [
+        requirement("numpy"),
+        requirement("pyscf"),
+    ],
+)
+
+py_test(
+    name = "compute_hfx_density_test",
+    srcs = ["density_functional_approximation_dm21/compute_hfx_density_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        ":compute_hfx_density",
+        requirement("attrs"),
+        requirement("numpy"),
+        requirement("scipy"),
+        requirement("pyscf"),
+        "@io_abseil_py//absl/testing:absltest",
+        "@io_abseil_py//absl/testing:parameterized",
+    ],
+)
+
+filegroup(
+    name = "dm21_checkpoints",
+    srcs = glob(["density_functional_approximation_dm21/checkpoints/**"]),
+)
+
+py_library(
+    name = "neural_numint",
+    srcs = ["density_functional_approximation_dm21/neural_numint.py"],
+    data = [":dm21_checkpoints"],
+    srcs_version = "PY3",
+    deps = [
+        ":compute_hfx_density",
+        requirement("attrs"),
+        requirement("keras"),
+        requirement("numpy"),
+        requirement("pyscf"),
+        "@org_tensorflow//tensorflow:tensorflow_py",
+        requirement("tensorflow-hub"),
+    ],
+)
+
+py_test(
+    name = "neural_numint_test",
+    srcs = ["density_functional_approximation_dm21/neural_numint_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        ":neural_numint",
+        requirement("attrs"),
+        requirement("pyscf"),
+        "@io_abseil_py//absl/testing:parameterized",
+        "@org_tensorflow//tensorflow:tensorflow_py",
+    ],
+)
+
+py_binary(
+    name = "export_saved_model",
+    srcs = ["density_functional_approximation_dm21/export_saved_model.py"],
+    data = [":dm21_checkpoints"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        ":neural_numint",
+    ],
+)
+
+EXPORTED_SAVED_MODEL_OBJECTS = [
+    "DM21/saved_model.pb",
+    "DM21/variables/variables.index",
+    "DM21/variables/variables.data-00000-of-00001",
+]
+
+genrule(
+    name = "create_model_for_aot_compile",
+    outs = EXPORTED_SAVED_MODEL_OBJECTS,
+    # The functional used in dm21_aot_compiled_example can be changed by
+    # setting the --functional flag to the desired functional.
+    cmd = "$(location :export_saved_model) --functional DM21 --batch_size 1000 --out_dir $(@D)/DM21",
+    tools = [":export_saved_model"],
+)
+
+filegroup(
+    name = "dm21_exported_model",
+    srcs = EXPORTED_SAVED_MODEL_OBJECTS,
+)
+
+saved_model_compile_aot(
+    name = "aot_compiled_dm21",
+    cpp_class = "dm21::functional",
+    directory = ":DM21",
+    filegroups = [":dm21_exported_model"],
+    force_without_xla_support_flag = False,
+    multithreading = False,
+    signature_def = "default",
+    tag_set = "''",
+    target_triple = "x86_64-pc-linux",
+)
+
+cc_library(
+    name = "dm21_aot_compiled_example",
+    srcs = ["cc/dm21_aot_compiled_example.cc"],
+    hdrs = ["cc/dm21_aot_compiled_example.h"],
+    deps = [":aot_compiled_dm21"],
+)
+
+cc_binary(
+    name = "run_dm21_aot_compiled_example",
+    srcs = ["cc/run_dm21_aot_compiled_example.cc"],
+    copts = ["-DXLA_AVAILABLE"],
+    deps = [":dm21_aot_compiled_example"],
+)
@@ -0,0 +1,248 @@
+# Pushing the Frontiers of Density Functionals by Solving the Fractional Electron Problem
+
+This package provides a PySCF interface to the DM21 (DeepMind 21) family of
+exchange-correlation functionals described in the paper ["Pushing the Frontiers
+of Density Functionals by Solving the Fractional Electron
+Problem"](https://doi.org/10.1126/science.abj6511).
+
+## Installation
+
+`pip install .` installs the interfaces to the DM21 functionals and required
+dependencies. This is best done inside a
+[virtual environment](https://docs.python-guide.org/dev/virtualenvs/).
+
+Note: using PySCF 2.0 (or later) enables substantially more efficient
+calculation of the load Hartree-Fock features, resulting in a large speed
+increase.
+
+### Installing directly
+
+To install without cloning or downloading the deepmind_research repository,
+execute:
+
+```shell
+python3 -m venv ~/venv/DM21
+source ~/venv/DM21/bin/activate
+pip install git+git://github.com/deepmind/deepmind-research.git#subdirectory=density_functional_approximation_dm21
+```
+
+### Downloading and installing from a local git repository
+
+Alternatively, clone the deepmind_research repository and install from a local
+repository:
+
+```shell
+git clone https://github.com/deepmind/deepmind-research.git
+cd deepmind-research/density_functional_approximation_dm21
+python3 -m venv ~/venv/DM21
+source ~/venv/DM21/bin/activate
+pip install .
+```
+
+The tests can be run either by running the test files directly or using
+`py.test`, again from the `density_functional_approximation_dm21` subdirectory:
+
+```shell
+pip install '.[testing]'
+py.test
+```
+
+Alternatively, conda can be used to manage the installation of the dependencies
+instead. Please see the [conda documentation](https://docs.conda.io/en/latest/)
+for install instructions. Both miniconda and anaconda work. Given a working
+conda installation, and having cloned the deepmind_research repository :
+
+```shell
+conda create -n dm21
+conda activate dm21
+conda install absl-py "h5py<3.3.0" numpy pytest scipy tensorflow
+conda install -c pyscf pyscf
+conda install -c conda-forge tensorflow-hub
+git clone https://github.com/deepmind/deepmind-research.git
+cd deepmind-research/density_functional_approximation_dm21
+pip install .
+py.test
+```
+
+## PySCF interface
+
+A typical DFT calculation with PySCF is set up and run using:
+
+```python
+from pyscf import gto
+from pyscf import dft
+
+# Create the molecule of interest and select the basis set.
+mol = gto.Mole()
+mol.atom = 'Ne 0.0 0.0 0.0'
+mol.basis = 'cc-pVDZ'
+mol.build()
+
+# Create a DFT solver and select the exchange-correlation functional.
+mf = dft.RKS(mol)
+mf.xc = 'b3lyp'
+
+# Run the DFT calculation.
+mf.kernel()
+```
+
+The DM21 functionals can be used in a very similar way, except we need to
+compute local Hartree-Fock features, which does not fit in with the interface
+used by conventional functionals. Instead, this package supplies a lightweight
+wrapper around PySCF's numerical integration class which evaluates the
+exchange-correlation energy and potential over a real-space grid. To use the
+DM21 functional with PySCF, the above code needs to be changed to:
+
+```python
+import density_functional_approximation_dm21 as dm21
+from pyscf import gto
+from pyscf import dft
+
+# Create the molecule of interest and select the basis set.
+mol = gto.Mole()
+mol.atom = 'Ne 0.0 0.0 0.0'
+mol.basis = 'cc-pVDZ'
+mol.build()
+
+# Create a DFT solver and insert the DM21 functional into the solver.
+mf = dft.RKS(mol)
+mf._numint = dm21.NeuralNumInt(dm21.Functional.DM21)
+
+# Run the DFT calculation.
+mf.kernel()
+```
+
+i.e. instead of specifying the functional with `mf.xc = <functional name>`, the
+functional is specified using `mf._numint = dm21.NeuralNumInt(<DM21
+functional>)`, where `<DM21 functional>` is the corresponding member of the
+`Functional` enum.
+
+Available functionals are:
+
+*   `DM21` - trained on molecules dataset, and fractional charge, and fractional
+    spin constraints.
+*   `DM21m` - trained on molecules dataset.
+*   `DM21mc` - trained on molecules dataset, and fractional charge constraints.
+*   `DM21mu` - trained on molecules dataset, and electron gas constraints.
+
+Full details of the network architecture, training method and datasets used can
+be found in the paper (reference below). Note that the results in our paper also
+include D3 corrections, which must be
+[included separately](https://pyscf.org/user/dft.html#dispersion-corrections).
+
+## Using DM21 from C++
+
+There are two options for using the DM21 from C++.
+
+1.  Load the SavedModel using
+    [TensorFlow's C++ API](https://www.tensorflow.org/guide/saved_model#load_a_savedmodel_in_c).
+    This requires a run-time dependency on the TensorFlow library.
+2.  Compile the model ahead-of-time into a standalone library. This requires all
+    array dimensions to be fixed at compile time, which imposes a minor
+    limitation on the interface for using the functional. As the DM21
+    functionals act on grid points independently, this does not restrict the
+    system size.
+
+The first option is more flexible but trickier to setup. Consequently we
+demonstrate the second option: compiling the functional into a standalone
+library. An example of running the DM21 functional using a standalone compiled
+library is provided in `cc/dm21_aot_compiled_example.cc`. This requires a
+link-time dependency on parts of the `xla_compiled_cpu_runtime_standalone`
+library, which are not included in the compiled functional library. The easiest
+way to build this is to use [Bazel](https://bazel.build). The first step is to
+[install Bazel](https://docs.bazel.build/versions/4.2.0/install.html).
+[Bazelisk](https://docs.bazel.build/versions/main/install-bazelisk.html) is
+another way to install Bazel if a native installer is not available. The
+following has been tested with Bazel 4.2.0. It is best to continue working
+inside a virtual environment.
+
+Assuming the above installation steps using `git clone` have been followed, and
+`Bazel` has been installed, the example can be built and run using:
+
+```
+pip install -r requirements_aot_compilation.txt
+bazel run --experimental_repo_remote_exec :run_dm21_aot_compiled_example
+```
+
+where the `pip install` command is only required if a fresh virtual environment
+is used, and installs required prerequisites for TensorFlow. See the
+[TensorFlow documentation](https://www.tensorflow.org/install/source) for more
+details.
+
+A static library, using the `cc_library` rule, can similarly be built and then
+linked against from a separate program with no additional dependencies required,
+shown in `cc/dm21_aot_compiled_example.cc` and the `dm21_aot_compiled_example`
+build rule.
+
+For calling from C, we recommend wrapping a C++ interface in `extern C { ... }`
+to create a C API. For calling from Fortran, the C API can be used via the
+Fortran 2003 ISO_C_BINDING feature.
+
+### Detailed explanation
+
+The supplied functionals in the `checkpoints` subdirectory are not easy to use
+from C++, as they only contain operations for the forward pass through the
+functional. This means they can only be used to evaluate the energy on a fixed
+density. Self-consistent calculations require various gradients of the
+functional, and it is easiest to create these using TensorFlow's Python API.
+`NeuralNumInt` contains a method for exporting the functional and functional
+derivatives. Assuming the above installation steps have been followed, a
+functional and its derivatives can be exported by:
+
+```shell
+export_saved_model.py --functional=DM21 \
+                      --out_dir=/path/to/export/DM21 \
+                      --batch_size=1000
+```
+
+where `--out_dir` specifies the directory to save the model containing the
+functional and functional derivatives to, and `--batch_size` the number of grid
+points the functional will be evaluated on at a time. The functional must be the
+name of a functional in the `neural_numint.Functional` enum. Note that the batch
+size needs only be fixed for exporting models to be used with ahead-of-time
+compilation.
+
+The functional can now be compiled using the `saved_model_cli` tool supplied
+with TensorFlow:
+
+```shell
+$ saved_model_cli aot_compile_cpu \
+                  --dir /path/to/export/DM21 \
+                  --output_prefix /path/to/compiled/DM21/dm21 \
+                  --cpp_class dm21::functional \
+                  --tag_set '' \
+                  --signature_def_key default
+```
+
+The output prefix and C++ class can be arbitrarily chosen. The `tag_set` and
+`signature_def_key` arguments must be as given above. This creates the following
+files in the output directory:
+
+-   dm21.h: header file defining the C++ interface to the DM21 functional.
+-   dm21_makefile.inc: a snippet to be included in a Makefile for setting
+    include, library and compilation flags.
+-   dm21_metadata.o, dm21.o: object files for running the DM21 functional.
+
+## Reference
+
+If this repository is helpful for your research please cite the following
+publication:
+
+Pushing the Frontiers of Density Functionals by Solving the Fractional Electron
+Problem, James Kirkpatrick, Brendan McMorrow, David H. P. Turban, Alexander L.
+Gaunt, James S. Spencer, Alexander G. D. G. Matthews, Annette Obika, Louis
+Thiry, Meire Fortunato, David Pfau, Lara Román Castellanos, Stig Petersen,
+Alexander W. R. Nelson, Pushmeet Kohli, Paula Mori-Sánchez, Demis Hassabis, Aron
+J. Cohen, Science, DOI: https://doi.org/10.1126/science.abj6511
+
+## License
+
+All code is made available under the Apache 2.0 License. Model parameters
+(contained in the `density_functional_approximation_dm21/checkpoints/`
+subdirectory) are made available under the Creative Commons Attribution 4.0
+International (CC BY 4.0) License. See
+https://creativecommons.org/licenses/by/4.0/legalcode for more details.
+
+## Disclaimer
+
+This is not an official Google product.
@@ -0,0 +1,62 @@
+workspace(name = "org_density_functional_approximation_dm21")
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+http_archive(
+    name = "rules_python",
+    sha256 = "934c9ceb552e84577b0faf1e5a2f0450314985b4d8712b2b70717dc679fdc01b",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.3.0/rules_python-0.3.0.tar.gz",
+)
+
+load("@rules_python//python:pip.bzl", "pip_install")
+
+# Create a central external repo, @external_py_deps, that contains Bazel
+# targets for all the third-party packages specified in the requirements.txt
+# file.
+pip_install(
+    name = "external_py_deps",
+    requirements = "//:requirements.txt",
+)
+
+http_archive(
+    name = "io_abseil_py",
+    strip_prefix = "abseil-py-pypi-v0.15.0",
+    urls = ["https://github.com/abseil/abseil-py/archive/pypi-v0.15.0.tar.gz"],
+)
+
+http_archive(
+    name = "six_archive",
+    build_file = "@io_abseil_py//third_party:six.BUILD",
+    strip_prefix = "six-1.12.0",
+    urls = [
+        "https://pypi.python.org/packages/source/s/six/six-1.12.0.tar.gz",
+    ],
+)
+
+http_archive(
+    name = "org_tensorflow",
+    patch_args = ["-p1"],
+    patches = ["tf_bazel.patch"],
+    strip_prefix = "tensorflow-2.7.0",
+    urls = ["https://github.com/tensorflow/tensorflow/archive/refs/tags/v2.7.0.tar.gz"],
+)
+
+# The cascade of load() statements and tf_workspace?() calls works around the
+# restriction that load() statements need to be at the top of .bzl files.
+# E.g. we can not retrieve a new repository with http_archive and then load()
+# a macro from that repository in the same file.
+load("@org_tensorflow//tensorflow:workspace3.bzl", "tf_workspace3")
+
+tf_workspace3()
+
+load("@org_tensorflow//tensorflow:workspace2.bzl", "tf_workspace2")
+
+tf_workspace2()
+
+load("@org_tensorflow//tensorflow:workspace1.bzl", "tf_workspace1")
+
+tf_workspace1()
+
+load("@org_tensorflow//tensorflow:workspace0.bzl", "tf_workspace0")
+
+tf_workspace0()
@@ -0,0 +1,97 @@
+// Copyright 2021 DeepMind Technologies Limited.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cstdio>
+
+// Users should adjust header path to match generated location. The generated
+// location is determined by the -output_prefix flags passed to saved_model_cli
+// aot_compile_cpu.
+#include "aot_compiled_dm21.h"  // generated
+
+void run_dm21_compiled_functional() {
+  // The functional class name is set by a flag passed to saved_model_cli
+  // aot_compile_cpu.
+  dm21::functional dm21_xc;
+
+  // This assumes the functional was compiled using instructions given in the
+  // README.md, in particular that the model was exported with a batch size of
+  // 1000.  We will compute the functional just at a single point and pad the
+  // rest of the batch using grid_weight = 0.
+
+  constexpr int batch_dim = 1000;
+  // See docstring for neural_numint.FunctionalInputs for descriptions on each
+  // input feature.  Note that, due how the model is compiled, we don't need to
+  // pass any values for the grid coordinates.
+  float rho_a[6][batch_dim] = {
+      {1.238148615359934e-11},   {-5.4047671667795604e-11},
+      {-5.4047671667795604e-11}, {-7.2613887530595865e-12},
+      {4.436179569857956e-10},   {5.972594378309204e-11},
+  };
+  float hfx_a[batch_dim][2] = {{-4.33218591e-13, -4.32842821e-13}};
+  float grid_weights[batch_dim] = {1.7594189642339968};
+
+  // Use same values for both alpha and beta electrons (restricted calculation)
+  dm21_xc.set_arg_feed_rho_a_data(rho_a);
+  dm21_xc.set_arg_feed_rho_b_data(rho_a);
+  dm21_xc.set_arg_feed_hfx_a_data(hfx_a);
+  dm21_xc.set_arg_feed_hfx_b_data(hfx_a);
+  dm21_xc.set_arg_feed_grid_weights_data(grid_weights);
+
+  std::puts("Running functional...");
+  bool status = dm21_xc.Run();
+
+  if (status) {
+    std::puts("Successfully ran functional.");
+    // Fetch results.
+    // Other methods for fetching results exist which may be more convenient.
+    // Please see the generated header.
+    // See neural_numint.NeuralNumint._build_graph and
+    // See neural_numint.NeuralNumint.eval_xc for more details.
+    // XC potential at each grid point, shape (batch_dim).
+    const float* vxc = dm21_xc.result_fetch_vxc_data();
+    // Derivative of the energy wtih respect to the density.
+    // In python, this has shape (2, batch_dim), where the zeroth component is
+    // with respect to the alpha density and the first component with respect to
+    // the beta density. In C++, a flat 1D array is returned.
+    const float* vrho = dm21_xc.result_fetch_vrho_data();
+    // Derivative of the energy wtih respect to sigma.
+    // In python, this has shape (3, batch_dim), where the zeroth component is
+    // with respect to the alpha spin channel, the first component with respect
+    // to the spin channel and the third component with respect to the total. In
+    // C++, a flat 1D array is returned.
+    const float* vsigma = dm21_xc.result_fetch_vsigma_data();
+    // Derivative of the energy wtih respect to tau, the kinetic energy density.
+    // In python, this has shape (2, batch_dim), where the zeroth component is
+    // with respect to the alpha spin channel, and the first component with
+    // respect to the spin channel. In C++, a flat 1D array is returned.
+    const float* vtau = dm21_xc.result_fetch_vtau_data();
+    // Intermediates required for evaluating the contribution of local
+    // Hartree-Fock features to the derivative of the Fock matrix. See
+    // docstrings and comments in compute_hfx_density.py and neural_numint.py.
+    // In python, this has shape (2, batch_dim, nomega), where nomega is the
+    // number of omega values used for the Hartree-Fock kernels. The zeroth
+    // component is with respect to the Hartree-Fock energy density at each grid
+    // point for the alpha-spin density and the first component with respect to
+    // the beta-spin density. In C++, a flat 1D array is returned.
+    const float* vhf = dm21_xc.result_fetch_vhf_data();
+    std::printf("vxc[0] = %.6g\n", vxc[0]);
+    std::printf("vrho[0] = %.6g, %.6g\n", vrho[0], vrho[batch_dim]);
+    std::printf("vsigma[0] = %.6g, %.6g %.6g\n", vsigma[0], vsigma[batch_dim],
+                vsigma[2 * batch_dim]);
+    std::printf("vtau[0] = %.6g, %.6g\n", vtau[0], vtau[batch_dim]);
+    std::printf("vhf[0] = %.6g, %.6g\n", vhf[0], vhf[1]);
+  } else {
+    std::puts("Failed to run functional.");
+  }
+}
@@ -0,0 +1,21 @@
+// Copyright 2021 DeepMind Technologies Limited.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef THIRD_PARTY_DEEPMIND_DEEPMIND_RESEARCH_DENSITY_FUNCTIONAL_APPROXIMATION_DM21_CC_DM21_AOT_COMPILED_EXAMPLE_H_
+#define THIRD_PARTY_DEEPMIND_DEEPMIND_RESEARCH_DENSITY_FUNCTIONAL_APPROXIMATION_DM21_CC_DM21_AOT_COMPILED_EXAMPLE_H_
+
+void run_dm21_compiled_functional();
+
+
+#endif  // THIRD_PARTY_DEEPMIND_DEEPMIND_RESEARCH_DENSITY_FUNCTIONAL_APPROXIMATION_DM21_CC_DM21_AOT_COMPILED_EXAMPLE_H_
@@ -0,0 +1,30 @@
+// Copyright 2021 DeepMind Technologies Limited.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cstdio>
+
+#ifdef XLA_AVAILABLE
+#include "cc/dm21_aot_compiled_example.h"
+#endif
+
+int main(int argc, char** argv) {
+#ifdef XLA_AVAILABLE
+  run_dm21_compiled_functional();
+  return 0;
+#else
+  std::puts("Built without XLA support. Cannot run functional!");
+  return 1;
+#endif
+}
+
@@ -0,0 +1,18 @@
+# Copyright 2021 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""An interface to DM21 family of exchange-correlation functionals for PySCF."""
+
+from density_functional_approximation_dm21.neural_numint import Functional
+from density_functional_approximation_dm21.neural_numint import NeuralNumInt
@@ -0,0 +1 @@
+
@@ -0,0 +1 @@
+
@@ -0,0 +1 @@
+
@@ -0,0 +1 @@
+
@@ -0,0 +1,335 @@
+# Copyright 2021 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+r"""Computation of the Hartree-Fock exchange density.
+
+We consider two types of potential:
+
+  1. Coulomb potential v(r,r') = 1/|r-r'|, which results in the full HF exchange
+     density and energy.
+  2. Screened (long-range) Coulomb potential v(r,r') = erf(\omega|r-r'|)/|r-r'|,
+     which results in the screened HF exchange density energy.
+
+Note that PySCF and libcint treat a value of omega=0 to correspond to the
+Coulomb potential. In the following, HF refers to full HF exchange if the
+Coulomb potential is used and to screened HF exchange if the screened Coulomb
+potential is used.
+
+The Hartree-Fock (HF) exchange energy can be written as:
+
+-2 HF_x = \sum_{a,b,c,d} D_{ab} D_{cd} \int dr \int dr'
+   [ \chi_a(r) \chi_c(r) v(r, r') \chi_b(r') \chi_d(r') ]
+
+where D is the density matrix, \chi_a the atomic basis functions and r, r' are
+coordinates. For clarity we have dropped the spin-channel label of the density
+matrix.
+
+Defining the following intermediates:
+
+\nu_{bd}(r) = \int dr' (\chi_b(r') v(r, r') \chi_d(r'))
+E_b(r) = \sum_a D_{ab} \chi_a(r)
+
+we get the following expression for HF:
+
+-2 HF_x = \int dr \sum_{bd} E_b(r) E_d(r) \nu_{bd}(r)
+
+Therefore the quantity
+
+exx(r) = -0.5 sum_{bd} E_b(r) E_d(r) \nu_{bd}(r)
+
+represents an energy density at location r which integrates to the HF exchange
+energy.
+
+The Fock matrix, F, is the derivative of the energy with respect to the density
+matrix. If the energy depends upon the set of features {x}, then the Fock matrix
+can be evaluated as \sum_x dE/dx dx/dD_{ab}. The derivatives with respect to the
+features can be easily evaluated using automatic differentiation. We hence
+require the derivative of exx with respect to the density matrix:
+
+dexx(r)/dD_{ab} = -D_{cd} \chi_a(r) \chi_c(r) \nu_{bd}(r)
+
+This is too large to store, so we instead compute the following intermediate,
+and evaluate the derivative as required on the fly:
+
+fxx_a(r) = D_{bd} \chi_a(r) \nu_{bd}(r)
+
+Note: we compute exx and fxx for each spin channel for both restricted and
+unrestricted calculations.
+"""
+
+from typing import Generator, Optional, Tuple, Union
+
+import attr
+import numpy as np
+from pyscf.dft import numint
+from pyscf.gto import mole
+from pyscf.lib import logger
+from pyscf.lib import numpy_helper
+
+
+def _evaluate_nu_slow(mol: mole.Mole,
+                      coords: np.ndarray,
+                      omega: float,
+                      hermi: int = 1) -> np.ndarray:
+  """Computes nu integrals for given coordinates using a slow loop."""
+  nu = []
+  # Use the Gaussian nuclear model in int1e_rinv_sph to evaluate the screened
+  # integrals.
+  with mol.with_rinv_zeta(zeta=omega * omega):
+    # This is going to be slow...
+    for coord in coords:
+      with mol.with_rinv_origin(coord):
+        nu.append(mol.intor('int1e_rinv_sph', hermi=hermi))
+  return np.asarray(nu)
+
+
+def _evaluate_nu(mol: mole.Mole,
+                 coords: np.ndarray,
+                 omega: float,
+                 hermi: int = 1) -> np.ndarray:
+  """Computes nu integrals for given coordinates."""
+  try:
+    with mol.with_range_coulomb(omega=omega):
+      # grids keyword argument supported in pyscf 2.0.0-alpha.
+      nu = mol.intor('int1e_grids_sph', hermi=hermi, grids=coords)  # pytype: disable=wrong-keyword-args
+  except TypeError:
+    logger.info(
+        mol, 'Support for int1e_grids not found (requires libcint 4.4.1 and '
+        'pyscf 2.0.0a or later. Falling back to slow loop over individual grid '
+        'points.')
+    nu = _evaluate_nu_slow(mol, coords, omega)
+  return nu
+
+
+def _nu_chunk(mol: mole.Mole,
+              coords: np.ndarray,
+              omega: float,
+              chunk_size: int = 1000
+             ) -> Generator[Tuple[int, int, np.ndarray], None, None]:
+  r"""Yields chunks of nu integrals over the grid.
+
+  Args:
+    mol: pyscf Mole object.
+    coords: coordinates, r', at which to evaluate the nu integrals, shape (N,3).
+    omega: range separation parameter. A value of 0 disables range-separation
+      (i.e. uses the kernel v(r,r') = 1/|r-r'| instead of
+      v(r,r') = erf(\omega |r-r'|) / |r-r'|)
+    chunk_size: number of coordinates to evaluate the integrals at a time.
+
+  Yields:
+    start_index, end_index, nu_{ab}(r) where
+      start_index, end_index are indices into coords,
+      nu is an array of shape (end_index-start_index, nao, nao), where nao is
+      the number of atomic orbitals and contains
+      nu_{ab}(r) = <a(r')|v(r,r')| b(r')>, where a,b are atomic
+      orbitals and r' are the grid coordinates in coords[start_index:end_index].
+
+  Raises:
+    ValueError: if omega is negative.
+  """
+  if omega < 0:
+    raise ValueError('Range-separated parameter omega must be non-negative!')
+  ncoords = len(coords)
+  for chunk_index in range(0, ncoords, chunk_size):
+    end_index = min(chunk_index + chunk_size, ncoords)
+    coords_chunk = coords[chunk_index:end_index]
+    nu_chunk = _evaluate_nu(mol, coords_chunk, omega=omega)
+    yield chunk_index, end_index, nu_chunk
+
+
+def _compute_exx_block(nu: np.ndarray,
+                       e: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+  r"""Computes exx and fxx.
+
+  Args:
+    nu: batch of <i|v(r,r_k)|j> integrals, in format (k,i,j) where r_k is the
+      position of the k-th grid point, i and j label atomic orbitals.
+    e: density matrix in the AO basis at each grid point.
+
+  Returns:
+    exx and fxx, where
+    fxx_{gb} =\sum_c nu_{gbc} e_{gc} and
+    exx_{g} = -0.5 \sum_b e_{gb} fxx_{gb}.
+  """
+  fxx = np.einsum('gbc,gc->gb', nu, e)
+  exx = -0.5 * np.einsum('gb,gb->g', e, fxx)
+  return exx, fxx
+
+
+def _compute_jk_block(nu: np.ndarray, fxx: np.ndarray, dm: np.ndarray,
+                      ao_value: np.ndarray,
+                      weights: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+  """Computes J and K contributions from the given block of nu integrals."""
+  batch_size = nu.shape[0]
+  vj = numpy_helper.dot(nu.reshape(batch_size, -1), dm.reshape(-1, 1))
+  vj = np.squeeze(vj)
+  vj_ao = np.einsum('g,gb->gb', vj * weights, ao_value)
+  j = numpy_helper.dot(ao_value.T, vj_ao)
+  w_ao = np.einsum('g,gb->gb', weights, ao_value)
+  k = numpy_helper.dot(fxx.T, w_ao)
+  return j, k
+
+
+@attr.s(auto_attribs=True)
+class HFDensityResult:
+  r"""Container for results returned by get_hf_density.
+
+  Note that the kernel used in all integrals is defined by the omega input
+  argument.
+
+  Attributes:
+    exx: exchange energy density at position r on the grid for the alpha, beta
+      spin channels.  Each array is shape (N), where N is the number of grid
+      points.
+    fxx: intermediate for evaluating dexx/dD^{\sigma}_{ab}, where D is the
+      density matrix and \sigma is the spin coordinate. See top-level docstring
+      for details.  Each array is shape (N, nao), where nao is the number of
+      atomic orbitals.
+    coulomb: coulomb matrix (restricted calculations) or matrices (unrestricted
+      calculations). Each array is shape (nao, nao).
+      Restricted calculations: \sum_{} D_{cd} (ab|cd)
+      Unrestricted calculations: \sum_{} D^{\sigma}_{cd} (ab|cd)
+    exchange: exchange matrix (restricted calculations) or matrices
+      (unrestricted calculations). Each array is shape (nao, nao).
+      Restricted calculations: \sum_{} D_{cd} (ab|cd)
+      Unrestricted calculations: \sum_{} D^{\sigma}_{cd} (ac|bd).
+  """
+  exx: Tuple[np.ndarray, np.ndarray]
+  fxx: Optional[Tuple[np.ndarray, np.ndarray]] = None
+  coulomb: Optional[Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]] = None
+  exchange: Optional[Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]] = None
+
+
+def get_hf_density(
+    mol: mole.Mole,
+    dm: Union[Tuple[np.ndarray, np.ndarray], np.ndarray],
+    coords: np.ndarray,
+    omega: float = 0.,
+    deriv: int = 0,
+    ao: Optional[np.ndarray] = None,
+    chunk_size: int = 1000,
+    weights: Optional[np.ndarray] = None,
+) -> HFDensityResult:
+  r"""Computes the (range-separated) HF energy density.
+
+  Args:
+    mol: PySCF molecule.
+    dm: The density matrix. For restricted calculations, an array of shape
+      (M, M), where M is the number of atomic orbitals. For unrestricted
+      calculations, either an array of shape (2, M, M) or a tuple of arrays,
+      each of shape (M, M), where dm[0] is the density matrix for the alpha
+      electrons and dm[1] the density matrix for the beta electrons.
+    coords: The coordinates to compute the HF density at, shape (N, 3), where N
+      is the number of grid points.
+    omega: The inverse width of the error function. An omega of 0. means range
+      separation and a 1/|r-R| kernel is used in the nu integrals. Otherwise,
+      the kernel erf(\omega|r-R|)/|r-R|) is used. Must be non-negative.
+    deriv: The derivative order. Only first derivatives (deriv=1) are currently
+      implemented. deriv=0 indicates no derivatives are required.
+    ao: The atomic orbitals evaluated on the grid, shape (N, M). These are
+      computed if not supplied.
+    chunk_size: The number of coordinates to compute the HF density for at once.
+      Reducing this saves memory since we don't have to keep as many Nus (nbasis
+      x nbasis) in memory at once.
+    weights: weight of each grid point, shape (N). If present, the Coulomb and
+      exchange matrices are also computed semi-numerically, otherwise only the
+      HF density and (if deriv=1) its first derivative are computed.
+
+  Returns:
+    HFDensityResult object with the HF density (exx), the derivative of the HF
+    density with respect to the density (fxx) if deriv is 1, and the Coulomb and
+    exchange matrices if the weights argument is provided.
+
+  Raises:
+    NotImplementedError: if a Cartesian basis set is used or if deriv is greater
+    than 1.
+    ValueError: if omega or deriv are negative.
+  """
+  if mol.cart:
+    raise NotImplementedError('Local HF exchange is not implmented for basis '
+                              'sets with Cartesian functions!')
+  if deriv < 0:
+    raise ValueError(f'`deriv` must be non-negative, got {deriv}')
+  if omega < 0:
+    raise ValueError(f'`omega` must be non-negative, got {omega}')
+  if deriv > 1:
+    raise NotImplementedError('Higher order derivatives are not implemented.')
+
+  if isinstance(dm, tuple) or dm.ndim == 3:
+    dma, dmb = dm
+    restricted = False
+  else:
+    dma = dm / 2
+    dmb = dm / 2
+    restricted = True
+
+  logger.info(mol, 'Computing contracted density matrix ...')
+  if ao is None:
+    ao = numint.eval_ao(mol, coords, deriv=0)
+  e_a = np.dot(ao, dma)
+  e_b = np.dot(ao, dmb)
+
+  exxa = []
+  exxb = []
+  fxxa = []
+  fxxb = []
+  ja = np.zeros_like(dma)
+  jb = np.zeros_like(dmb)
+  ka = np.zeros_like(dma)
+  kb = np.zeros_like(dmb)
+
+  for start, end, nu in _nu_chunk(mol, coords, omega, chunk_size=chunk_size):
+    logger.info(mol, 'Computing exx %s / %s ...', end, len(e_a))
+    exxa_block, fxxa_block = _compute_exx_block(nu, e_a[start:end])
+    exxa.extend(exxa_block)
+    if not restricted:
+      exxb_block, fxxb_block = _compute_exx_block(nu, e_b[start:end])
+      exxb.extend(exxb_block)
+    if deriv == 1:
+      fxxa.extend(fxxa_block)
+      if not restricted:
+        fxxb.extend(fxxb_block)
+
+    if weights is not None:
+      ja_block, ka_block = _compute_jk_block(nu, fxxa_block, dma, ao[start:end],
+                                             weights[start:end])
+      ja += ja_block
+      ka += ka_block
+      if not restricted:
+        jb_block, kb_block = _compute_jk_block(nu, fxxb_block, dmb,
+                                               ao[start:end],
+                                               weights[start:end])
+        jb += jb_block
+        kb += kb_block
+
+  exxa = np.asarray(exxa)
+  fxxa = np.asarray(fxxa)
+  if restricted:
+    exxb = exxa
+    fxxb = fxxa
+  else:
+    exxb = np.asarray(exxb)
+    fxxb = np.asarray(fxxb)
+
+  result = HFDensityResult(exx=(exxa, exxb))
+  if deriv == 1:
+    result.fxx = (fxxa, fxxb)
+  if weights is not None:
+    if restricted:
+      result.coulomb = 2 * ja
+      result.exchange = 2 * ka
+    else:
+      result.coulomb = (ja, jb)
+      result.exchange = (ka, kb)
+  return result
@@ -0,0 +1,226 @@
+# Copyright 2021 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for compute_hfx_density."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import numpy as np
+from pyscf import dft
+from pyscf import gto
+from pyscf import lib
+from pyscf import scf
+import scipy
+from density_functional_approximation_dm21 import compute_hfx_density
+
+
+class ComputeHfxDensityTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    lib.param.TMPDIR = None
+    lib.num_threads(1)
+
+  @parameterized.named_parameters(
+      {'testcase_name': 'local_hf', 'omega': 0.},
+      {'testcase_name': 'range_separated_local_hf_0.5', 'omega': 0.5},
+      {'testcase_name': 'range_separated_local_hf_1.0', 'omega': 1.0},
+      {'testcase_name': 'range_separated_local_hf_2.0', 'omega': 2.0},
+  )
+  def test_closed_shell(self, omega):
+    mol = gto.M(atom='He 0. 0. 0.', basis='3-21g')
+    solver = dft.RKS(mol)
+    solver.grids.level = 2
+    solver.grids.build()
+    solver.kernel()
+    dm = solver.make_rdm1()
+    with mol.with_range_coulomb(omega=omega):
+      target_j, target_k = scf.hf.get_jk(mol, dm)
+    target_hf = -0.25 * np.einsum('ij,ji', dm, target_k)
+    target_coulomb = np.einsum('ij,ji', dm, target_j)
+
+    coords = solver.grids.coords
+    weights = solver.grids.weights
+
+    results = compute_hfx_density.get_hf_density(
+        mol, dm, coords, omega=omega, weights=weights)
+    coulomb = np.einsum('ij,ji', dm, results.coulomb)
+    hf = -0.25 * np.einsum('ij,ji', dm, results.exchange)
+    predicted_hf = np.sum((results.exx[0] + results.exx[1]) * weights)
+
+    with self.subTest('test_hf_density'):
+      self.assertAlmostEqual(target_hf, predicted_hf)
+
+    with self.subTest('test_get_jk'):
+      np.testing.assert_allclose(results.coulomb, target_j)
+      np.testing.assert_allclose(results.exchange, target_k)
+      self.assertAlmostEqual(coulomb, target_coulomb)
+      self.assertAlmostEqual(hf, target_hf)
+
+  @parameterized.named_parameters(
+      {'testcase_name': 'local_hf', 'omega': 0.},
+      {'testcase_name': 'range_separated_local_hf_0.5', 'omega': 0.5},
+      {'testcase_name': 'range_separated_local_hf_1.0', 'omega': 1.0},
+      {'testcase_name': 'range_separated_local_hf_2.0', 'omega': 2.0},
+  )
+  def test_hf_density_on_open_shell(self, omega):
+    mol = gto.M(atom='He 0. 0. 0.', basis='3-21g', charge=1, spin=1)
+    solver = dft.UKS(mol)
+    solver.grids.level = 2
+    solver.grids.build()
+    solver.kernel()
+    dm = solver.make_rdm1()
+    with mol.with_range_coulomb(omega=omega):
+      target_j, target_k = scf.hf.get_jk(mol, dm)
+    target_hf = -0.5 * (
+        np.einsum('ij,ji', dm[0], target_k[0]) +
+        np.einsum('ij,ji', dm[1], target_k[1]))
+    target_coulomb = np.einsum('ij,ji', dm[0], target_j[0]) + np.einsum(
+        'ij,ji', dm[1], target_j[1])
+
+    coords = solver.grids.coords
+    weights = solver.grids.weights
+
+    results = compute_hfx_density.get_hf_density(
+        mol, dm, coords, omega=omega, weights=weights)
+
+    predicted_hf = np.sum((results.exx[0] + results.exx[1]) * weights)
+    coulomb = (
+        np.einsum('ij,ji', dm[0], results.coulomb[0]) +
+        np.einsum('ij,ji', dm[1], results.coulomb[1]))
+    hf = -0.5 * (
+        np.einsum('ij,ji', dm[0], results.exchange[0]) +
+        np.einsum('ij,ji', dm[1], results.exchange[1]))
+
+    with self.subTest('test_hf_density'):
+      self.assertAlmostEqual(target_hf, predicted_hf, places=3)
+
+    with self.subTest('test_get_jk'):
+      np.testing.assert_allclose(results.coulomb[0], target_j[0])
+      np.testing.assert_allclose(results.coulomb[1], target_j[1])
+      np.testing.assert_allclose(results.exchange[0], target_k[0])
+      np.testing.assert_allclose(results.exchange[1], target_k[1])
+      self.assertAlmostEqual(coulomb, target_coulomb)
+      self.assertAlmostEqual(hf, target_hf)
+
+
+def _nu_test_systems():
+  systems = [
+      {
+          'atom': 'N 0 0 0; N 0 0 2.4',
+          'charge': 0,
+          'spin': 0,
+          'basis': 'cc-pVDZ',
+          'num_grids': -1
+      },
+      {
+          'atom': 'N 0 0 0; N 0 0 2.4',
+          'charge': 0,
+          'spin': 0,
+          'basis': 'cc-pVDZ',
+          'num_grids': 1
+      },
+      {
+          'atom': 'N 0 0 0; N 0 0 2.4',
+          'charge': 0,
+          'spin': 0,
+          'basis': 'cc-pVDZ',
+          'num_grids': 2
+      },
+      {
+          'atom': 'N 0 0 0; N 0 0 2.4',
+          'charge': 0,
+          'spin': 0,
+          'basis': 'cc-pVDZ',
+          'num_grids': 10
+      },
+      {
+          'atom': 'N 0 0 0; N 0 0 2.4',
+          'charge': 0,
+          'spin': 0,
+          'basis': 'cc-pVDZ',
+          'num_grids': 32
+      },
+      {
+          'atom': 'N 0 0 0; N 0 0 2.4',
+          'charge': 0,
+          'spin': 0,
+          'basis': 'cc-pVDZ',
+          'num_grids': 33
+      },
+      {
+          'atom': 'Li 0 0 0',
+          'charge': 0,
+          'spin': 1,
+          'basis': 'cc-pVTZ',
+          'num_grids': -1
+      },
+      {
+          'atom': 'H 0 0 0',
+          'charge': 0,
+          'spin': 1,
+          'basis': 'cc-pVQZ',
+          'num_grids': -1
+      },
+  ]
+  system_names = ['N2', 'N2_1', 'N2_2', 'N2_10', 'N2_32', 'N2_33', 'Li', 'H']
+  for name, system in zip(system_names, systems):
+    yield {'testcase_name': f'{name}_hermitian', 'hermi': 0, **system}
+    yield {'testcase_name': f'{name}_non_hermitian', 'hermi': 1, **system}
+
+
+class NuTest(parameterized.TestCase):
+
+  def setUp(self):
+    super(NuTest, self).setUp()
+    lib.param.TMPDIR = None
+    lib.num_threads(1)
+
+  @parameterized.named_parameters(_nu_test_systems())
+  def test_nu_integrals(self, atom, charge, spin, basis, num_grids, hermi):
+    mol = gto.M(atom=atom, charge=charge, spin=spin, basis=basis)
+    mf = dft.UKS(mol)
+    mf.grids.build()
+    if num_grids == -1:
+      test_coords = mf.grids.coords
+    else:
+      test_coords = mf.grids.coords[0:num_grids]
+    nu_slow = compute_hfx_density._evaluate_nu_slow(
+        mol, test_coords, omega=0.0, hermi=hermi)
+    nu_fast = compute_hfx_density._evaluate_nu(
+        mol, test_coords, omega=0.0, hermi=hermi)
+    np.testing.assert_allclose(nu_slow, nu_fast, atol=1E-13)
+
+  def test_range_separated_nu(self):
+    mol = gto.M(atom='He 0 0 0', basis='cc-pVDZ')
+    r0 = np.array([[0.1, 0.2, 1.]])
+    omega = 1.
+    result = np.squeeze(compute_hfx_density._evaluate_nu(mol, r0, omega=omega))
+
+    solver = dft.RKS(mol)
+    solver.grids.level = 2
+    solver.grids.build()
+    coords = solver.grids.coords
+    weights = solver.grids.weights
+    ao_value = dft.numint.eval_ao(mol, coords, deriv=0)
+    dist = np.linalg.norm(coords - r0, axis=1)
+    erf = scipy.special.erf(omega * dist) / dist
+    expected_result = np.squeeze(
+        np.einsum('g,ga,gb->ab', weights * erf, ao_value, ao_value))
+
+    np.testing.assert_allclose(result, expected_result)
+
+
+if __name__ == '__main__':
+  absltest.main()
@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+# Copyright 2021 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helper for exporting a functional and its derivatives to a saved_model."""
+
+from typing import Sequence
+
+from absl import app
+from absl import flags
+
+from density_functional_approximation_dm21 import neural_numint
+
+_OUT_DIR = flags.DEFINE_string(
+    'out_dir', None, 'Output directory.', required=True)
+_BATCH_SIZE = flags.DEFINE_integer(
+    'batch_size',
+    1000,
+    'Number of grid points exported functional will process in a single call.',
+    lower_bound=0)
+_FUNCTIONAL = flags.DEFINE_enum_class('functional',
+                                      neural_numint.Functional.DM21,
+                                      neural_numint.Functional,
+                                      'Functional to export.')
+
+
+def export(
+    functional: neural_numint.Functional,
+    export_path: str,
+    batch_dim: int,
+) -> None:
+  """Export a functional and its derivatives to a single saved_model.
+
+  Args:
+    functional: functional to export.
+    export_path: path to saved the model to.
+    batch_dim: number of grid points to process in a single call.
+  """
+  ni = neural_numint.NeuralNumInt(functional)
+  ni.export_functional_and_derivatives(
+      export_path=export_path, batch_dim=batch_dim)
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+  export(_FUNCTIONAL.value, _OUT_DIR.value, _BATCH_SIZE.value)
+
+
+if __name__ == '__main__':
+  app.run(main)
@@ -0,0 +1,158 @@
+# Copyright 2021 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for neural_numint."""
+
+import os
+
+
+from absl.testing import parameterized
+import attr
+from pyscf import dft
+from pyscf import gto
+from pyscf import lib
+import tensorflow.compat.v1 as tf
+
+from density_functional_approximation_dm21 import neural_numint
+
+
+class NeuralNumintTest(tf.test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    lib.param.TMPDIR = None
+    lib.num_threads(1)
+
+  # Golden values were obtained using the version of PySCF (including integral
+  # generation) reported in the DM21 paper.
+  @parameterized.parameters(
+      {
+          'functional': neural_numint.Functional.DM21,
+          'expected_energy': -126.898521
+      },
+      {
+          'functional': neural_numint.Functional.DM21m,
+          'expected_energy': -126.907332
+      },
+      {
+          'functional': neural_numint.Functional.DM21mc,
+          'expected_energy': -126.922127
+      },
+      {
+          'functional': neural_numint.Functional.DM21mu,
+          'expected_energy': -126.898178
+      },
+  )
+  def test_rks(self, functional, expected_energy):
+    ni = neural_numint.NeuralNumInt(functional)
+
+    mol = gto.Mole()
+    mol.atom = [['Ne', 0., 0., 0.]]
+    mol.basis = 'sto-3g'
+    mol.build()
+
+    mf = dft.RKS(mol)
+    mf._numint = ni
+    mf.run()
+    self.assertAlmostEqual(mf.e_tot, expected_energy, delta=2.e-4)
+
+  @parameterized.parameters(
+      {
+          'functional': neural_numint.Functional.DM21,
+          'expected_energy': -37.34184876
+      },
+      {
+          'functional': neural_numint.Functional.DM21m,
+          'expected_energy': -37.3377766
+      },
+      {
+          'functional': neural_numint.Functional.DM21mc,
+          'expected_energy': -37.33489173
+      },
+      {
+          'functional': neural_numint.Functional.DM21mu,
+          'expected_energy': -37.34015315
+      },
+  )
+  def test_uks(self, functional, expected_energy):
+    ni = neural_numint.NeuralNumInt(functional)
+
+    mol = gto.Mole()
+    mol.atom = [['C', 0., 0., 0.]]
+    mol.spin = 2
+    mol.basis = 'sto-3g'
+    mol.build()
+
+    mf = dft.UKS(mol)
+    mf._numint = ni
+    mf.run()
+    self.assertAlmostEqual(mf.e_tot, expected_energy, delta=2.e-4)
+
+  def test_exported_model(self):
+
+    mol = gto.Mole()
+    mol.atom = [['C', 0., 0., 0.]]
+    mol.spin = 2
+    mol.basis = 'sto-3g'
+    mol.build()
+
+    ni = neural_numint.NeuralNumInt(neural_numint.Functional.DM21)
+    mf = dft.UKS(mol)
+    mf._numint = ni
+    mf.run()
+
+    dms = mf.make_rdm1()
+    ao = ni.eval_ao(mol, mf.grids.coords, deriv=2)
+    rho_a = ni.eval_rho(mol, ao, dms[0], xctype='MGGA')
+    rho_b = ni.eval_rho(mol, ao, dms[1], xctype='MGGA')
+    inputs, _ = ni.construct_functional_inputs(
+        mol=mol,
+        dms=dms,
+        spin=1,
+        coords=mf.grids.coords,
+        weights=mf.grids.weights,
+        rho=(rho_a, rho_b),
+        ao=ao[0])
+
+    feed_dict = dict(
+        zip(
+            attr.asdict(ni._placeholders).values(),
+            attr.asdict(inputs).values(),
+        ))
+    with ni._graph.as_default():
+      outputs = ni._session.run(
+          {
+              'vxc': ni._vxc,
+              'vrho': ni._vrho,
+              'vsigma': ni._vsigma,
+              'vtau': ni._vtau,
+              'vhf': ni._vhf
+          },
+          feed_dict=feed_dict)
+
+    export_path = os.path.join(self.get_temp_dir(), 'export')
+    ni.export_functional_and_derivatives(export_path)
+    model = tf.saved_model.load_v2(export_path)
+    tensor_inputs = {
+        k: tf.constant(v, dtype=tf.float32)
+        for k, v in attr.asdict(inputs).items()
+    }
+    exported_output_tensors = model.signatures['default'](**tensor_inputs)
+    with tf.Session() as session:
+      session.run(tf.global_variables_initializer())
+      exported_outputs = session.run(exported_output_tensors)
+    self.assertAllClose(outputs, exported_outputs, atol=5.e-5, rtol=1.e-5)
+
+
+if __name__ == '__main__':
+  tf.test.main()
@@ -0,0 +1,78 @@
+From 76edcc2da8a754cb2b283d7de9a7b37f51f7f380 Mon Sep 17 00:00:00 2001
+From: James Spencer <jamessspencer@google.com>
+Date: Thu, 26 Aug 2021 10:00:53 +0100
+Subject: [PATCH] Set dependencies for saved_model_compile_aot rule,
+ visibility for xla_compiled_cpu_runtime_standalone to work outside of
+ tensorflow, and absl version.
+
+---
+ tensorflow/compiler/tf2xla/BUILD  | 2 +-
+ tensorflow/python/tools/tools.bzl | 4 ++--
+ tensorflow/workspace2.bzl         | 9 +++++----
+ 3 files changed, 8 insertions(+), 7 deletions(-)
+
+diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
+index cfe63b16675..cd273aaf29e 100644
+--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
+@@ -235,7 +235,7 @@ cc_library(
+     copts = runtime_copts() + tf_openmp_copts(),
+     features = ["fully_static_link"],
+     linkstatic = 1,
+-    visibility = [":friends"],
+    visibility = ["//visibility:public"],
+     # Note, we specifically removed MKL and multithreaded dependencies so the
+     # standalone does not require the MKL binary blob or threading libraries.
+     #
+diff --git a/tensorflow/python/tools/tools.bzl b/tensorflow/python/tools/tools.bzl
+index db886746006..bc597e29de9 100644
+--- a/tensorflow/python/tools/tools.bzl
+++ b/tensorflow/python/tools/tools.bzl
+@@ -148,7 +148,7 @@ def saved_model_compile_aot(
+         ),
+         tags = tags,
+         tools = [
+-            "//tensorflow/python/tools:saved_model_cli",
+            "@org_tensorflow//tensorflow/python/tools:saved_model_cli",
+         ],
+     )
+ 
+@@ -170,7 +170,7 @@ def saved_model_compile_aot(
+         tags = tags,
+         deps = _maybe_force_compile(
+             [
+-                "//tensorflow/compiler/tf2xla:xla_compiled_cpu_runtime_standalone",
+                "@org_tensorflow//tensorflow/compiler/tf2xla:xla_compiled_cpu_runtime_standalone",
+             ],
+             force_compile = force_without_xla_support_flag,
+         ),
+diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl
+index df4b687d720..f1173507e72 100644
+--- a/tensorflow/workspace2.bzl
+++ b/tensorflow/workspace2.bzl
+@@ -488,8 +488,8 @@ def _tf_repositories():
+ 
+     tf_http_archive(
+         name = "absl_py",
+-        sha256 = "588a23406b2e28ea368897dbebc1210165414e87212d4fdd4b2ee968f0a772c6",
+-        strip_prefix = "abseil-py-pypi-v0.10.0",
+        sha256 = "0be59b82d65dfa1f995365dcfea2cc57989297b065fda696ef13f30fcc6c8e5b",
+        strip_prefix = "abseil-py-pypi-v0.15.0",
+         system_build_file = "//third_party/systemlibs:absl_py.BUILD",
+         system_link_files = {
+             "//third_party/systemlibs:absl_py.absl.BUILD": "absl/BUILD",
+@@ -498,8 +498,9 @@ def _tf_repositories():
+             "//third_party/systemlibs:absl_py.absl.logging.BUILD": "absl/logging/BUILD",
+         },
+         urls = [
+-            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/abseil/abseil-py/archive/pypi-v0.10.0.tar.gz",
+-            "https://github.com/abseil/abseil-py/archive/pypi-v0.10.0.tar.gz",
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/abseil/abseil-py/archive/pypi-v0.15.0.tar.gz",
+            "https://github.com/abseil/abseil-py/archive/pypi-v0.15.0.tar.gz",
+            "https://files.pythonhosted.org/packages/75/c6/ea1b86d2e7068e77f5204f8280edb2434596c1bad59fe03564f3d11d5caf/absl-py-0.15.0.tar.gz",
+         ],
+     )
+ 
+-- 
+2.33.0.259.gc128427fd7-goog
+
@@ -0,0 +1,9 @@
+absl-py==0.13.0
+attrs==21.2.0
+h5py==3.1.0
+numpy==1.19.5
+pyscf==1.7.6.post1
+pytest==6.2.4
+scipy==1.7.1
+tensorflow==2.6.0
+tensorflow-hub==0.12.0
@@ -0,0 +1,3 @@
+Keras-Preprocessing==1.1.2
+numpy==1.19.5
+tensorflow-estimator==2.7.0
@@ -0,0 +1,27 @@
+#!/bin/bash
+# Copyright 2021 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Install density_functional_approximation_dm21 in a clean virtualenv and run
+# the tests.  This assumes the working directory is the top-level directory of
+# the deepmind-research repository, i.e.:
+# git clone git@github.com:deepmind/deepmind-research.git
+# cd deepmind_research
+# density_functional_approximation_dm21/run.sh
+
+python3 -m venv /tmp/DM21
+source /tmp/DM21/bin/activate
+pip3 install -r density_functional_approximation_dm21/requirements.txt
+pip install 'density_functional_approximation_dm21/[testing]'
+py.test density_functional_approximation_dm21/
@@ -0,0 +1,50 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Setup for DM21 functionals interface to PySCF."""
+
+from setuptools import setup
+
+REQUIRED_PACKAGES = [
+    'absl-py',
+    'attrs',
+    # PySCF 1.7.6 and older do not support h5py 3.3.0:
+    # https://github.com/pyscf/pyscf/issues/1016
+    # If using PySCF 2.0 or later, this restriction can be lifted.
+    'h5py<3.3.0',
+    'numpy',
+    'pyscf',
+    'tensorflow',
+    'tensorflow_hub',
+]
+CHECKPOINT_DATA = ['checkpoints/DM21*/*.pb', 'checkpoints/DM21*/variables/*']
+
+setup(
+    name='density_functional_approximation_dm21',
+    version='0.1',
+    description='An interface to PySCF for the DM21 functionals.',
+    url='https://github.com/deepmind/deepmind-research/density_functional_approximation_dm21',
+    author='DeepMind',
+    author_email='no-reply@google.com',
+    # Contained modules and scripts.
+    packages=['density_functional_approximation_dm21'],
+    package_data={
+        'density_functional_approximation_dm21': CHECKPOINT_DATA,
+    },
+    scripts=['density_functional_approximation_dm21/export_saved_model.py'],
+    install_requires=REQUIRED_PACKAGES,
+    platforms=['any'],
+    license='Apache 2.0',
+    extras_require={'testing': ['pytest', 'scipy']},
+)