diff --git a/README.md b/README.md index 5d56dcd..25860f2 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/density_functional_approximation_dm21/.bazelrc b/density_functional_approximation_dm21/.bazelrc new file mode 100644 index 0000000..68d0fcf --- /dev/null +++ b/density_functional_approximation_dm21/.bazelrc @@ -0,0 +1,609 @@ +# The following configuration has been taken from the TensorFlow 2.5.1 release. +# +# TensorFlow Bazel configuration file. +# This file tries to group and simplify build options for TensorFlow +# +# ----CONFIG OPTIONS---- +# Android options: +# android: +# android_arm: +# android_arm64: +# android_x86: +# android_x86_64: +# +# iOS options: +# ios: +# ios_armv7: +# ios_arm64: +# ios_i386: +# ios_x86_64: +# ios_fat: +# +# Macosx options +# darwin_arm64: +# +# Compiler options: +# cuda_clang: Use clang when building CUDA code. +# c++17: Build with C++17 options (links with libc++) +# c++1z: Build with C++17 options (links with libc++) +# c++17_gcc: Build with C++17 options (links with stdlibc++) +# c++1z_gcc: Build with C++17 options (links with stdlibc++) +# avx_linux: Build with avx instruction set on linux. +# avx2_linux: Build with avx2 instruction set on linux. +# native_arch_linux: Build with instruction sets available to the host machine on linux +# avx_win: Build with avx instruction set on windows +# avx2_win: Build with avx2 instruction set on windows +# +# Other build options: +# short_logs: Only log errors during build, skip warnings. +# verbose_logs: Show all compiler warnings during build. +# monolithic: Build all TF C++ code into a single shared object. +# dynamic_kernels: Try to link all kernels dynamically (experimental). +# libc++: Link against libc++ instead of stdlibc++ +# asan: Build with the clang address sanitizer +# msan: Build with the clang memory sanitizer +# ubsan: Build with the clang undefined behavior sanitizer +# +# +# TF version options; +# v1: Build TF V1 (without contrib) +# v2: Build TF v2 +# +# Feature and Third party library support options: +# xla: Build TF with XLA +# tpu: Build TF with TPU support +# cuda: Build with full cuda support. +# rocm: Build with AMD GPU support (rocm). +# mkl: Enable full mkl support. +# tensorrt: Enable Tensorrt support. +# numa: Enable numa using hwloc. +# noaws: Disable AWS S3 storage support +# nogcp: Disable GCS support. +# nohdfs: Disable hadoop hdfs support. +# nonccl: Disable nccl support. +# +# +# Remote build execution options (only configured to work with TF team projects for now.) +# rbe: General RBE options shared by all flavors. +# rbe_linux: General RBE options used on all linux builds. +# rbe_win: General RBE options used on all windows builds. +# +# rbe_cpu_linux: RBE options to build with only CPU support. +# rbe_linux_cuda_nvcc_py*: RBE options to build with GPU support using nvcc. +# +# rbe_linux_py2: Linux Python 2 RBE config. +# rbe_linux_py3: Linux Python 3 RBE config +# +# rbe_win_py37: Windows Python 3.7 RBE config +# rbe_win_py38: Windows Python 3.8 RBE config +# +# tensorflow_testing_rbe_linux: RBE options to use RBE with tensorflow-testing project on linux +# tensorflow_testing_rbe_win: RBE options to use RBE with tensorflow-testing project on windows +# +# Embedded Linux options (experimental and only tested with TFLite build yet) +# elinux: General Embedded Linux options shared by all flavors. +# elinux_aarch64: Embedded Linux options for aarch64 (ARM64) CPU support. +# elinux_armhf: Embedded Linux options for armhf (ARMv7) CPU support. +# +# Release build options (for all operating systems) +# release_base: Common options for all builds on all operating systems. +# release_gpu_base: Common options for GPU builds on Linux and Windows. +# release_cpu_linux: Toolchain and CUDA options for Linux CPU builds. +# release_cpu_macos: Toolchain and CUDA options for MacOS CPU builds. +# release_gpu_linux: Toolchain and CUDA options for Linux GPU builds. +# release_cpu_windows: Toolchain and CUDA options for Windows CPU builds. +# release_gpu_windows: Toolchain and CUDA options for Windows GPU builds. + +# Default build options. These are applied first and unconditionally. + +# For projects which use TensorFlow as part of a Bazel build process, putting +# nothing in a bazelrc will default to a monolithic build. The following line +# opts in to modular op registration support by default. +build --define framework_shared_object=true + +# For workaround https://github.com/bazelbuild/bazel/issues/8772 with Bazel >= 0.29.1 +build --java_toolchain=@tf_toolchains//toolchains/java:tf_java_toolchain +build --host_java_toolchain=@tf_toolchains//toolchains/java:tf_java_toolchain + +build --define=use_fast_cpp_protos=true +build --define=allow_oversize_protos=true + +build --spawn_strategy=standalone +build -c opt + +# Make Bazel print out all options from rc files. +build --announce_rc + +build --define=grpc_no_ares=true + +# See https://github.com/bazelbuild/bazel/issues/7362 for information on what +# --incompatible_remove_legacy_whole_archive flag does. +# This flag is set to true in Bazel 1.0 and newer versions. We tried to migrate +# Tensorflow to the default, however test coverage wasn't enough to catch the +# errors. +# There is ongoing work on Bazel team's side to provide support for transitive +# shared libraries. As part of migrating to transitive shared libraries, we +# hope to provide a better mechanism for control over symbol exporting, and +# then tackle this issue again. +# +# TODO: Remove this line once TF doesn't depend on Bazel wrapping all library +# archives in -whole_archive -no_whole_archive. +build --noincompatible_remove_legacy_whole_archive + +# These are bazel 2.0's incompatible flags. Tensorflow needs to use bazel 2.0.0 +# to use cc_shared_library, as part of the Tensorflow Build Improvements RFC: +# https://github.com/tensorflow/community/pull/179 +build --noincompatible_prohibit_aapt1 + +build --enable_platform_specific_config + +# Enable XLA support by default. +build --define=with_xla_support=true + +build --config=short_logs + +build --config=v2 + +# Default options should come above this line. + +# Allow builds using libc++ as a linker library +# This is mostly for OSSFuzz, so we also pass in the flags from environment to clean build file +build:libc++ --action_env=CC +build:libc++ --action_env=CXX +build:libc++ --action_env=CXXFLAGS=-stdlib=libc++ +build:libc++ --action_env=PATH +build:libc++ --define force_libcpp=enabled +build:libc++ --linkopt -fuse-ld=lld + +# Android configs. Bazel needs to have --cpu and --fat_apk_cpu both set to the +# target CPU to build transient dependencies correctly. See +# https://docs.bazel.build/versions/master/user-manual.html#flag--fat_apk_cpu +build:android --crosstool_top=//external:android/crosstool +build:android --host_crosstool_top=@bazel_tools//tools/cpp:toolchain +build:android_arm --config=android +build:android_arm --cpu=armeabi-v7a +build:android_arm --fat_apk_cpu=armeabi-v7a +build:android_arm64 --config=android +build:android_arm64 --cpu=arm64-v8a +build:android_arm64 --fat_apk_cpu=arm64-v8a +build:android_x86 --config=android +build:android_x86 --cpu=x86 +build:android_x86 --fat_apk_cpu=x86 +build:android_x86_64 --config=android +build:android_x86_64 --cpu=x86_64 +build:android_x86_64 --fat_apk_cpu=x86_64 + +# Sets the default Apple platform to macOS. +build:macos --apple_platform_type=macos + +# gRPC on MacOS requires this #define +build:macos --copt=-DGRPC_BAZEL_BUILD + +# Settings for MacOS on ARM CPUs. +build:macos_arm64 --cpu=darwin_arm64 + +# iOS configs for each architecture and the fat binary builds. +build:ios --apple_platform_type=ios +build:ios --apple_bitcode=embedded --copt=-fembed-bitcode +build:ios --copt=-Wno-c++11-narrowing +build:ios_armv7 --config=ios +build:ios_armv7 --cpu=ios_armv7 +build:ios_arm64 --config=ios +build:ios_arm64 --cpu=ios_arm64 +build:ios_i386 --config=ios +build:ios_i386 --cpu=ios_i386 +build:ios_x86_64 --config=ios +build:ios_x86_64 --cpu=ios_x86_64 +build:ios_fat --config=ios +build:ios_fat --ios_multi_cpus=armv7,arm64,i386,x86_64 + +# Config to use a mostly-static build and disable modular op registration +# support (this will revert to loading TensorFlow with RTLD_GLOBAL in Python). +# By default, TensorFlow will build with a dependence on +# //tensorflow:libtensorflow_framework.so. +build:monolithic --define framework_shared_object=false + +# Please note that MKL on MacOS or windows is still not supported. +# If you would like to use a local MKL instead of downloading, please set the +# environment variable "TF_MKL_ROOT" every time before build. +build:mkl --define=build_with_mkl=true --define=enable_mkl=true +build:mkl --define=tensorflow_mkldnn_contraction_kernel=0 +build:mkl --define=build_with_openmp=true +build:mkl -c opt + +# config to build OneDNN backend with a user specified threadpool. +build:mkl_threadpool --define=build_with_mkl=true --define=enable_mkl=true +build:mkl_threadpool --define=tensorflow_mkldnn_contraction_kernel=0 +build:mkl_threadpool --define=build_with_mkl_opensource=true +build:mkl_threadpool -c opt + +# Config setting to build oneDNN with Compute Library for the Arm Architecture (ACL). +# This build is for the inference regime only. +build:mkl_aarch64 --define=build_with_mkl_aarch64=true --define=enable_mkl=true +build:mkl_aarch64 --define=tensorflow_mkldnn_contraction_kernel=0 +build:mkl_aarch64 --define=build_with_mkl_opensource=true +build:mkl_aarch64 --define=build_with_openmp=true +build:mkl_aarch64 -c opt + +# This config refers to building CUDA op kernels with nvcc. +build:cuda --repo_env TF_NEED_CUDA=1 +build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain +build:cuda --@local_config_cuda//:enable_cuda + +# This config refers to building CUDA op kernels with clang. +build:cuda_clang --config=cuda +build:cuda_clang --repo_env TF_CUDA_CLANG=1 +build:cuda_clang --@local_config_cuda//:cuda_compiler=clang + +# Debug config +build:dbg -c dbg +# for now, disable arm_neon. see: https://github.com/tensorflow/tensorflow/issues/33360 +build:dbg --cxxopt -DTF_LITE_DISABLE_X86_NEON +# AWS SDK must be compiled in release mode. see: https://github.com/tensorflow/tensorflow/issues/37498 +build:dbg --copt -DDEBUG_BUILD + +# Config to build TPU backend +build:tpu --define=with_tpu_support=true + +build:tensorrt --repo_env TF_NEED_TENSORRT=1 + +build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain +build:rocm --define=using_rocm=true --define=using_rocm_hipcc=true +build:rocm --repo_env TF_NEED_ROCM=1 + +# Options extracted from configure script +build:numa --define=with_numa_support=true + +# Options to disable default on features +build:noaws --define=no_aws_support=true +build:nogcp --define=no_gcp_support=true +build:nohdfs --define=no_hdfs_support=true +build:nonccl --define=no_nccl_support=true + +build:stackdriver_support --define=stackdriver_support=true + +# Modular TF build options +build:dynamic_kernels --define=dynamic_loaded_kernels=true +build:dynamic_kernels --copt=-DAUTOLOAD_DYNAMIC_KERNELS + +# Build TF with C++ 17 features. +build:c++17 --cxxopt=-std=c++1z +build:c++17 --cxxopt=-stdlib=libc++ +build:c++1z --config=c++17 +build:c++17_gcc --cxxopt=-std=c++1z +build:c++1z_gcc --config=c++17_gcc + +# Don't trigger --config= when cross-compiling. +build:android --noenable_platform_specific_config +build:ios --noenable_platform_specific_config + +# Suppress C++ compiler warnings, otherwise build logs become 10s of MBs. +build:android --copt=-w +build:ios --copt=-w +build:linux --copt=-w +build:linux --host_copt=-w +build:macos --copt=-w +build:windows --copt=/W0 + +# Tensorflow uses M_* math constants that only get defined by MSVC headers if +# _USE_MATH_DEFINES is defined. +build:windows --copt=/D_USE_MATH_DEFINES +build:windows --host_copt=/D_USE_MATH_DEFINES + +# Default paths for TF_SYSTEM_LIBS +build:linux --define=PREFIX=/usr +build:linux --define=LIBDIR=$(PREFIX)/lib +build:linux --define=INCLUDEDIR=$(PREFIX)/include +build:linux --define=PROTOBUF_INCLUDE_PATH=$(PREFIX)/include +build:macos --define=PREFIX=/usr +build:macos --define=LIBDIR=$(PREFIX)/lib +build:macos --define=INCLUDEDIR=$(PREFIX)/include +build:macos --define=PROTOBUF_INCLUDE_PATH=$(PREFIX)/include +# TF_SYSTEM_LIBS do not work on windows. + +# By default, build TF in C++ 14 mode. +build:android --cxxopt=-std=c++14 +build:android --host_cxxopt=-std=c++14 +build:ios --cxxopt=-std=c++14 +build:ios --host_cxxopt=-std=c++14 +build:linux --cxxopt=-std=c++14 +build:linux --host_cxxopt=-std=c++14 +build:macos --cxxopt=-std=c++14 +build:macos --host_cxxopt=-std=c++14 +build:windows --cxxopt=/std:c++14 +build:windows --host_cxxopt=/std:c++14 + +# On windows, we still link everything into a single DLL. +build:windows --config=monolithic + +# On linux, we dynamically link small amount of kernels +build:linux --config=dynamic_kernels + +# Make sure to include as little of windows.h as possible +build:windows --copt=-DWIN32_LEAN_AND_MEAN +build:windows --host_copt=-DWIN32_LEAN_AND_MEAN +build:windows --copt=-DNOGDI +build:windows --host_copt=-DNOGDI + +# MSVC (Windows): Standards-conformant preprocessor mode +# See https://docs.microsoft.com/en-us/cpp/preprocessor/preprocessor-experimental-overview +build:windows --copt=/experimental:preprocessor +build:windows --host_copt=/experimental:preprocessor + +# Misc build options we need for windows. +build:windows --linkopt=/DEBUG +build:windows --host_linkopt=/DEBUG +build:windows --linkopt=/OPT:REF +build:windows --host_linkopt=/OPT:REF +build:windows --linkopt=/OPT:ICF +build:windows --host_linkopt=/OPT:ICF +build:windows --experimental_strict_action_env=true + +# Verbose failure logs when something goes wrong +build:windows --verbose_failures + +# On windows, we never cross compile +build:windows --distinct_host_configuration=false + +# Configure short or long logs +build:short_logs --output_filter=DONT_MATCH_ANYTHING +build:verbose_logs --output_filter= + +# Instruction set optimizations +# TODO(gunan): Create a feature in toolchains for avx/avx2 to +# avoid having to define linux/win separately. +build:avx_linux --copt=-mavx +build:avx_linux --host_copt=-mavx +build:avx2_linux --copt=-mavx2 +build:native_arch_linux --copt=-march=native +build:avx_win --copt=/arch=AVX +build:avx2_win --copt=/arch=AVX2 + +# Options to build TensorFlow 1.x or 2.x. +build:v1 --define=tf_api_version=1 --action_env=TF2_BEHAVIOR=0 +build:v2 --define=tf_api_version=2 --action_env=TF2_BEHAVIOR=1 + +# Disable XLA on mobile. +build:xla --define=with_xla_supprt=true # TODO: remove, it's on by default. +build:android --define=with_xla_support=false +build:ios --define=with_xla_support=false + +# BEGIN TF REMOTE BUILD EXECUTION OPTIONS +# Options when using remote execution +# WARNING: THESE OPTIONS WONT WORK IF YOU DO NOT HAVE PROPER AUTHENTICATION AND PERMISSIONS + +# Flag to enable remote config +common --experimental_repo_remote_exec + +build:rbe --repo_env=BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1 +build:rbe --google_default_credentials +build:rbe --bes_backend=buildeventservice.googleapis.com +build:rbe --bes_results_url="https://source.cloud.google.com/results/invocations" +build:rbe --bes_timeout=600s +build:rbe --define=EXECUTOR=remote +build:rbe --distinct_host_configuration=false +build:rbe --flaky_test_attempts=3 +build:rbe --jobs=200 +build:rbe --remote_executor=grpcs://remotebuildexecution.googleapis.com +build:rbe --remote_timeout=3600 +build:rbe --spawn_strategy=remote,worker,standalone,local +test:rbe --test_env=USER=anon +# Attempt to minimize the amount of data transfer between bazel and the remote +# workers: +build:rbe --remote_download_toplevel + +build:rbe_linux --config=rbe +build:rbe_linux --action_env=PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/go/bin" +build:rbe_linux --host_javabase=@bazel_toolchains//configs/ubuntu16_04_clang/1.1:jdk8 +build:rbe_linux --javabase=@bazel_toolchains//configs/ubuntu16_04_clang/1.1:jdk8 +build:rbe_linux --host_java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8 +build:rbe_linux --java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8 + +# Non-rbe settings we should include because we do not run configure +build:rbe_linux --config=avx_linux +# TODO(gunan): Check why we need this specified in rbe, but not in other builds. +build:rbe_linux --linkopt=-lrt +build:rbe_linux --host_linkopt=-lrt +build:rbe_linux --linkopt=-lm +build:rbe_linux --host_linkopt=-lm + +build:rbe_cpu_linux --config=rbe_linux +build:rbe_cpu_linux --host_crosstool_top="@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010:toolchain" +build:rbe_cpu_linux --crosstool_top="@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010:toolchain" +build:rbe_cpu_linux --extra_toolchains="@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010:cc-toolchain-k8" +build:rbe_cpu_linux --extra_execution_platforms="@ubuntu16.04-manylinux2010-py3_config_platform//:platform" +build:rbe_cpu_linux --extra_execution_platforms="@ubuntu16.04-manylinux2010-py3_config_platform//:platform" +build:rbe_cpu_linux --host_platform="@ubuntu16.04-manylinux2010-py3_config_platform//:platform" +build:rbe_cpu_linux --platforms="@ubuntu16.04-manylinux2010-py3_config_platform//:platform" + +build:rbe_linux_cuda_base --config=rbe_linux +build:rbe_linux_cuda_base --config=cuda +build:rbe_linux_cuda_base --config=tensorrt +build:rbe_linux_cuda_base --action_env=TF_CUDA_VERSION=11 +build:rbe_linux_cuda_base --action_env=TF_CUDNN_VERSION=8 +build:rbe_linux_cuda_base --repo_env=REMOTE_GPU_TESTING=1 +test:rbe_linux_cuda_base --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64" + +build:rbe_linux_cuda11.2_nvcc_base --config=rbe_linux_cuda_base +build:rbe_linux_cuda11.2_nvcc_base --host_crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain" +build:rbe_linux_cuda11.2_nvcc_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain" +build:rbe_linux_cuda11.2_nvcc_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64" +build:rbe_linux_cuda11.2_nvcc_base --extra_execution_platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform" +build:rbe_linux_cuda11.2_nvcc_base --host_platform="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform" +build:rbe_linux_cuda11.2_nvcc_base --platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform" +build:rbe_linux_cuda11.2_nvcc_base --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda" +build:rbe_linux_cuda11.2_nvcc_base --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_tensorrt" +build:rbe_linux_cuda11.2_nvcc_base --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_nccl" +build:rbe_linux_cuda11.2_nvcc_py3.6 --config=rbe_linux_cuda11.2_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.6" +build:rbe_linux_cuda11.2_nvcc_py3.7 --config=rbe_linux_cuda11.2_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.7" +build:rbe_linux_cuda11.2_nvcc_py3.8 --config=rbe_linux_cuda11.2_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.8" +build:rbe_linux_cuda11.2_nvcc_py3.9 --config=rbe_linux_cuda11.2_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.9" + +# Map default to CUDA 11.2. +build:rbe_linux_cuda_nvcc_py36 --config=rbe_linux_cuda11.2_nvcc_py3.6 +build:rbe_linux_cuda_nvcc_py37 --config=rbe_linux_cuda11.2_nvcc_py3.7 +build:rbe_linux_cuda_nvcc_py38 --config=rbe_linux_cuda11.2_nvcc_py3.8 +build:rbe_linux_cuda_nvcc_py39 --config=rbe_linux_cuda11.2_nvcc_py3.9 + +# Deprecated configs that people might still use. +build:rbe_linux_cuda_nvcc --config=rbe_linux_cuda_nvcc_py36 +build:rbe_gpu_linux --config=rbe_linux_cuda_nvcc + +build:rbe_linux_cuda_clang_base --config=rbe_linux_cuda_base +build:rbe_linux_cuda_clang_base --repo_env TF_CUDA_CLANG=1 +build:rbe_linux_cuda_clang_base --@local_config_cuda//:cuda_compiler=clang +build:rbe_linux_cuda_clang_base --crosstool_top="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain" +build:rbe_linux_cuda_clang_base --extra_toolchains="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64" +build:rbe_linux_cuda_clang_base --extra_execution_platforms="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform" +build:rbe_linux_cuda_clang_base --host_platform="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform" +build:rbe_linux_cuda_clang_base --platforms="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform" +build:rbe_linux_cuda_clang_base --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda" +build:rbe_linux_cuda_clang_base --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_tensorrt" +build:rbe_linux_cuda_clang_base --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_nccl" +build:rbe_linux_cuda_clang_py27 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python2.7" +build:rbe_linux_cuda_clang_py35 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.5" +build:rbe_linux_cuda_clang_py36 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.6" +build:rbe_linux_cuda_clang_py37 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.7" +build:rbe_linux_cuda_clang_py38 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.8" + +# ROCm +build:rbe_linux_rocm_base --config=rbe_linux +build:rbe_linux_rocm_base --repo_env=TF_NEED_ROCM=1 +build:rbe_linux_rocm_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-rocm_config_rocm//crosstool:toolchain" +build:rbe_linux_rocm_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-rocm_config_rocm//crosstool:toolchain-linux-x86_64" +build:rbe_linux_rocm_base --extra_execution_platforms="@ubuntu18.04-gcc7_manylinux2010-rocm_config_platform//:platform" +build:rbe_linux_rocm_base --host_platform="@ubuntu18.04-gcc7_manylinux2010-rocm_config_platform//:platform" +build:rbe_linux_rocm_base --platforms="@ubuntu18.04-gcc7_manylinux2010-rocm_config_platform//:platform" +build:rbe_linux_rocm_base --action_env=TF_ROCM_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_rocm" +build:rbe_linux_rocm_base --define=using_rocm_hipcc=true +build:rbe_linux_rocm_py2.7 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python2.7" +build:rbe_linux_rocm_py3.5 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python3.5" +build:rbe_linux_rocm_py3.6 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python3.6" +build:rbe_linux_rocm_py3.7 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python3.7" +build:rbe_linux_rocm_py3.8 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python3.8" + +# Linux CPU +build:rbe_linux_py2 --config=rbe_linux +build:rbe_linux_py2 --repo_env=PYTHON_BIN_PATH="/usr/bin/python2" +build:rbe_linux_py2 --python_path="/usr/bin/python2" +build:rbe_linux_py2 --repo_env=TF_PYTHON_CONFIG_REPO="@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/py" + +build:rbe_linux_py3 --config=rbe_linux +build:rbe_linux_py3 --python_path="/usr/bin/python3" +build:rbe_linux_py3 --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu16.04-manylinux2010-py3_config_python" + +build:rbe_win --config=rbe +build:rbe_win --crosstool_top="@org_tensorflow//third_party/toolchains/preconfig/win/tf_win_08062020:toolchain" +build:rbe_win --extra_toolchains="@org_tensorflow//third_party/toolchains/preconfig/win/tf_win_08062020:cc-toolchain-x64_windows" +build:rbe_win --host_javabase="@org_tensorflow//third_party/toolchains/preconfig/win:windows_jdk8" +build:rbe_win --javabase="@org_tensorflow//third_party/toolchains/preconfig/win:windows_jdk8" +build:rbe_win --extra_execution_platforms="@org_tensorflow//third_party/toolchains/preconfig/win:rbe_windows_ltsc2019" +build:rbe_win --host_platform="@org_tensorflow//third_party/toolchains/preconfig/win:rbe_windows_ltsc2019" +build:rbe_win --platforms="@org_tensorflow//third_party/toolchains/preconfig/win:rbe_windows_ltsc2019" +build:rbe_win --shell_executable=C:\\tools\\msys64\\usr\\bin\\bash.exe + +# TODO(gunan): Remove once we use MSVC 2019 with latest patches. +build:rbe_win --define=override_eigen_strong_inline=true +build:rbe_win --jobs=100 + +build:rbe_win_py37 --config=rbe +build:rbe_win_py37 --repo_env=TF_PYTHON_CONFIG_REPO="@windows_py37_config_python" +build:rbe_win_py37 --python_path=C:\\Python37\\python.exe + +build:rbe_win_py38 --config=rbe +build:rbe_win_py38 --repo_env=PYTHON_BIN_PATH=C:\\Python38\\python.exe +build:rbe_win_py38 --repo_env=PYTHON_LIB_PATH=C:\\Python38\\lib\\site-packages +build:rbe_win_py38 --repo_env=TF_PYTHON_CONFIG_REPO=@org_tensorflow//third_party/toolchains/preconfig/win_1803/py38 +build:rbe_win_py38 --python_path=C:\\Python38\\python.exe + +# These you may need to change for your own GCP project. +build:tensorflow_testing_rbe --project_id=tensorflow-testing +common:tensorflow_testing_rbe_linux --remote_instance_name=projects/tensorflow-testing/instances/default_instance +build:tensorflow_testing_rbe_linux --config=tensorflow_testing_rbe + +common:tensorflow_testing_rbe_win --remote_instance_name=projects/tensorflow-testing/instances/windows +build:tensorflow_testing_rbe_win --config=tensorflow_testing_rbe + +# TFLite build configs for generic embedded Linux +build:elinux --crosstool_top=@local_config_embedded_arm//:toolchain +build:elinux --host_crosstool_top=@bazel_tools//tools/cpp:toolchain +build:elinux_aarch64 --config=elinux +build:elinux_aarch64 --cpu=aarch64 +build:elinux_armhf --config=elinux +build:elinux_armhf --cpu=armhf +# END TF REMOTE BUILD EXECUTION OPTIONS + +# Config-specific options should come above this line. + +# Load rc file written by ./configure. +try-import %workspace%/.tf_configure.bazelrc + +# Load rc file with user-specific options. +try-import %workspace%/.bazelrc.user + +# Here are bazelrc configs for release builds +build:release_base --config=v2 +build:release_base --distinct_host_configuration=false +test:release_base --flaky_test_attempts=3 +test:release_base --test_size_filters=small,medium + +build:release_cpu_linux --config=release_base +build:release_cpu_linux --config=avx_linux +build:release_cpu_linux --crosstool_top=@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda11.2:toolchain +test:release_cpu_linux --test_env=LD_LIBRARY_PATH + +build:release_cpu_macos --config=release_base +build:release_cpu_macos --config=avx_linux + +build:release_gpu_base --config=cuda +build:release_gpu_base --action_env=TF_CUDA_VERSION="11" +build:release_gpu_base --action_env=TF_CUDNN_VERSION="8" +build:release_gpu_base --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_35,sm_50,sm_60,sm_70,sm_75,compute_80" + +build:release_gpu_linux --config=release_cpu_linux +build:release_gpu_linux --config=release_gpu_base +build:release_gpu_linux --config=tensorrt +build:release_gpu_linux --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-11.2" +build:release_gpu_linux --action_env=LD_LIBRARY_PATH="/usr/local/cuda:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/tensorrt/lib" +build:release_gpu_linux --action_env=GCC_HOST_COMPILER_PATH="/usr/bin/gcc-5" +build:release_gpu_linux --crosstool_top=@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda11.2:toolchain + +build:release_cpu_windows --config=release_base +build:release_cpu_windows --config=avx_win +build:release_cpu_windows --define=no_tensorflow_py_deps=true +# First available in VS 16.4. Speeds Windows compile times by a lot. See +# https://groups.google.com/a/tensorflow.org/d/topic/build/SsW98Eo7l3o/discussion +build:release_cpu_windows --copt=/d2ReducedOptimizeHugeFunctions --host_copt=/d2ReducedOptimizeHugeFunctions + +build:release_gpu_windows --config=release_cpu_windows +build:release_gpu_windows --config=release_gpu_base + +# Address sanitizer +# CC=clang bazel build --config asan +build:asan --strip=never +build:asan --copt -fsanitize=address +build:asan --copt -DADDRESS_SANITIZER +build:asan --copt -g +build:asan --copt -O3 +build:asan --copt -fno-omit-frame-pointer +build:asan --linkopt -fsanitize=address + +# Memory sanitizer +# CC=clang bazel build --config msan +build:msan --strip=never +build:msan --copt -fsanitize=memory +build:msan --copt -DADDRESS_SANITIZER +build:msan --copt -g +build:msan --copt -O3 +build:msan --copt -fno-omit-frame-pointer +build:msan --linkopt -fsanitize=memory + +# Undefined Behavior Sanitizer +# CC=clang bazel build --config ubsan +build:ubsan --strip=never +build:ubsan --copt -fsanitize=undefined +build:ubsan --copt -g +build:ubsan --copt -O3 +build:ubsan --copt -fno-omit-frame-pointer +build:ubsan --linkopt -fsanitize=undefined +build:ubsan --linkopt -lubsan diff --git a/density_functional_approximation_dm21/BUILD.bazel b/density_functional_approximation_dm21/BUILD.bazel new file mode 100644 index 0000000..838a602 --- /dev/null +++ b/density_functional_approximation_dm21/BUILD.bazel @@ -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"], +) diff --git a/density_functional_approximation_dm21/README.md b/density_functional_approximation_dm21/README.md new file mode 100644 index 0000000..658987b --- /dev/null +++ b/density_functional_approximation_dm21/README.md @@ -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 = `, the +functional is specified using `mf._numint = dm21.NeuralNumInt()`, where `` 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. diff --git a/density_functional_approximation_dm21/WORKSPACE.bazel b/density_functional_approximation_dm21/WORKSPACE.bazel new file mode 100644 index 0000000..e4e7ee3 --- /dev/null +++ b/density_functional_approximation_dm21/WORKSPACE.bazel @@ -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() diff --git a/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.cc b/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.cc new file mode 100644 index 0000000..86059be --- /dev/null +++ b/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.cc @@ -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 + +// 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."); + } +} diff --git a/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.h b/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.h new file mode 100644 index 0000000..529e0a4 --- /dev/null +++ b/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.h @@ -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_ diff --git a/density_functional_approximation_dm21/cc/run_dm21_aot_compiled_example.cc b/density_functional_approximation_dm21/cc/run_dm21_aot_compiled_example.cc new file mode 100644 index 0000000..17f96a0 --- /dev/null +++ b/density_functional_approximation_dm21/cc/run_dm21_aot_compiled_example.cc @@ -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 + +#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 +} + diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/__init__.py b/density_functional_approximation_dm21/density_functional_approximation_dm21/__init__.py new file mode 100644 index 0000000..eb3c862 --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/__init__.py @@ -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 diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/saved_model.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/saved_model.pb new file mode 100644 index 0000000..112a158 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/saved_model.pb differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/smart_module.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/smart_module.pb new file mode 100644 index 0000000..b33834c Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/smart_module.pb differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/tfhub_module.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/tfhub_module.pb new file mode 100644 index 0000000..d65dd8f --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/tfhub_module.pb @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.data-00000-of-00001 b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.data-00000-of-00001 new file mode 100644 index 0000000..8bee898 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.data-00000-of-00001 differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.index b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.index new file mode 100644 index 0000000..245f492 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.index differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/saved_model.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/saved_model.pb new file mode 100644 index 0000000..7b1ef6e Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/saved_model.pb differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/smart_module.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/smart_module.pb new file mode 100644 index 0000000..acd48e8 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/smart_module.pb differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/tfhub_module.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/tfhub_module.pb new file mode 100644 index 0000000..d65dd8f --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/tfhub_module.pb @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.data-00000-of-00001 b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.data-00000-of-00001 new file mode 100644 index 0000000..2a9952b Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.data-00000-of-00001 differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.index b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.index new file mode 100644 index 0000000..dfb7433 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.index differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/saved_model.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/saved_model.pb new file mode 100644 index 0000000..47f2664 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/saved_model.pb differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/smart_module.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/smart_module.pb new file mode 100644 index 0000000..adbaff6 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/smart_module.pb differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/tfhub_module.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/tfhub_module.pb new file mode 100644 index 0000000..d65dd8f --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/tfhub_module.pb @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.data-00000-of-00001 b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.data-00000-of-00001 new file mode 100644 index 0000000..684a385 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.data-00000-of-00001 differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.index b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.index new file mode 100644 index 0000000..ac424ff Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.index differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/saved_model.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/saved_model.pb new file mode 100644 index 0000000..5bd3cff Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/saved_model.pb differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/smart_module.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/smart_module.pb new file mode 100644 index 0000000..0749a4e Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/smart_module.pb differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/tfhub_module.pb b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/tfhub_module.pb new file mode 100644 index 0000000..d65dd8f --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/tfhub_module.pb @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.data-00000-of-00001 b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.data-00000-of-00001 new file mode 100644 index 0000000..a8ad26d Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.data-00000-of-00001 differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.index b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.index new file mode 100644 index 0000000..d80edc8 Binary files /dev/null and b/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.index differ diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density.py b/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density.py new file mode 100644 index 0000000..edea1fc --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density.py @@ -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) = , 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 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 diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density_test.py b/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density_test.py new file mode 100644 index 0000000..96277c9 --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density_test.py @@ -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() diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/export_saved_model.py b/density_functional_approximation_dm21/density_functional_approximation_dm21/export_saved_model.py new file mode 100644 index 0000000..0d872b4 --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/export_saved_model.py @@ -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) diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py new file mode 100644 index 0000000..bceea7b --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py @@ -0,0 +1,771 @@ +# 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.""" + +import enum +import os +from typing import Generator, Optional, Sequence, Tuple, Union + +import attr +import numpy as np +from pyscf import dft +from pyscf import gto +from pyscf.dft import numint +import tensorflow.compat.v1 as tf +import tensorflow_hub as hub + +from density_functional_approximation_dm21 import compute_hfx_density + +tf.disable_v2_behavior() + +# TODO(b/196260242): avoid depending upon private function +_dot_ao_ao = numint._dot_ao_ao # pylint: disable=protected-access + + +@enum.unique +class Functional(enum.Enum): + """Enum for exchange-correlation functionals in the DM21 family. + + Attributes: + 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. + """ + # Break pylint's preferred naming pattern to match the functional names used + # in the paper. + # pylint: disable=invalid-name + DM21 = enum.auto() + DM21m = enum.auto() + DM21mc = enum.auto() + DM21mu = enum.auto() + # pylint: enable=invalid-name + + +# We use attr.s instead of here instead of dataclasses.dataclass as +# dataclasses.asdict returns a deepcopy of the attributes. This is wasteful in +# memory if they are large and breaks (as in the case of tf.Tensors) if they are +# not serializable. attr.asdict does not perform this copy and so works with +# both np.ndarrays and tf.Tensors. +@attr.s(auto_attribs=True) +class FunctionalInputs: + r""""Inputs required for DM21 functionals. + + Depending upon the context, this is either a set of numpy arrays (feature + construction) or TF tensors (constructing placeholders/running functionals). + + Attributes: + rho_a: Density information for the alpha electrons. + PySCF for meta-GGAs supplies a single array for the total density + (restricted calculations) and a pair of arrays, one for each spin channel + (unrestricted calculations). + Each array/tensor is of shape (6, N) and contains the density and density + derivatives, where: + rho(0, :) - density at each grid point + rho(1, :) - norm of the derivative of the density at each grid point + along x + rho(2, :) - norm of the derivative of the density at each grid point + along y + rho(3, :) - norm of the derivative of the density at each grid point + along z + rho(4, :) - \nabla^2 \rho [not used] + rho(5, :) - tau (1/2 (\nabla \rho)^2) at each grid point. + See pyscf.dft.numint.eval_rho for more details. + We require separate inputs for both alpha- and beta-spin densities, even + in restricted calculations (where rho_a = rho_b = rho/2, where rho is the + total density). + rho_b: as for rho_a for the beta electrons. + hfx_a: local Hartree-Fock energy density at each grid point for the alpha- + spin density for each value of omega. Shape [N, len(omega_values)]. + See compute_hfx_density for more details. + hfx_b: as for hfx_a for the beta-spin density. + grid_coords: grid coordinates at which to evaluate the density. Shape + (N, 3), where N is the number of grid points. Note that this is currently + unused by the functional, but is still a required input. + grid_weights: weight of each grid point. Shape (N). + """ + rho_a: Union[tf.Tensor, np.ndarray] + rho_b: Union[tf.Tensor, np.ndarray] + hfx_a: Union[tf.Tensor, np.ndarray] + hfx_b: Union[tf.Tensor, np.ndarray] + grid_coords: Union[tf.Tensor, np.ndarray] + grid_weights: Union[tf.Tensor, np.ndarray] + + +@attr.s(auto_attribs=True) +class _GridState: + """Internal state required for the numerical grid. + + Attributes: + coords: coordinates of the grid. Shape (N, 3), where N is the number of grid + points. + weight: weight associated with each grid point. Shape (N). + mask: mask indicating whether a shell is zero at a grid point. Shape + (N, nbas) where nbas is the number of shells in the basis set. See + pyscf.dft.gen_grids.make_mask. + ao: atomic orbitals evaluated on the grid. Shape (N, nao), where nao is the + number of atomic orbitals, or shape (:, N, nao), where the 0-th element + contains the ao values, the next three elements contain the first + derivatives, and so on. + """ + coords: np.ndarray + weight: np.ndarray + mask: np.ndarray + ao: np.ndarray + + +@attr.s(auto_attribs=True) +class _SystemState: + """Internal state required for system of interest. + + Attributes: + mol: PySCF molecule + dms: density matrix or matrices (unrestricted calculations only). + Restricted calculations: shape (nao, nao), where nao is the number of + atomic orbitals. + Unrestricted calculations: shape (2, nao, nao) or a sequence (length 2) of + arrays of shape (nao, nao), and dms[0] and dms[1] are the density matrices + of the alpha and beta electrons respectively. + """ + mol: gto.Mole + dms: Union[np.ndarray, Sequence[np.ndarray]] + + +def _get_number_of_density_matrices(dms): + """Returns the number of density matrices in dms.""" + # See pyscf.numint.NumInt._gen_rho_evaluator + if isinstance(dms, np.ndarray) and dms.ndim == 2: + return 1 + return len(dms) + + +class NeuralNumInt(numint.NumInt): + """A wrapper around pyscf.dft.numint.NumInt for the DM21 functionals. + + In order to supply the local Hartree-Fock features required for the DM21 + functionals, we lightly wrap the NumInt class. The actual evaluation of the + exchange-correlation functional is performed in NeuralNumInt.eval_xc. + + Usage: + mf = dft.RKS(...) # dft.ROKS and dft.UKS are also supported. + # Specify the functional by monkey-patching mf._numint rather than using + # mf._xc or mf._define_xc_. + mf._numint = NeuralNumInt(Functional.DM21) + mf.kernel() + """ + + def __init__(self, functional: Functional): + """Constructs a NeuralNumInt object. + + Args: + functional: member of Functional enum giving the name of the + functional. + """ + + self._functional_name = functional.name + self._model_path = os.path.join( + os.path.dirname(__file__), 'checkpoints', self._functional_name) + + # All DM21 functionals use local Hartree-Fock features with a non-range + # separated 1/r kernel and a range-seperated kernel with \omega = 0.4. + # Note an omega of 0.0 is interpreted by PySCF and libcint to indicate no + # range-separation. + self._omega_values = [0.0, 0.4] + self._graph = tf.Graph() + with self._graph.as_default(): + self._build_graph() + self._session = tf.Session() + self._session.run(tf.global_variables_initializer()) + + self._grid_state = None + self._system_state = None + self._vmat_hf = None + super().__init__() + + def _build_graph(self, batch_dim: Optional[int] = None): + """Builds the TensorFlow graph for evaluating the functional. + + Args: + batch_dim: the batch dimension of the grid to use in the model. Default: + None (determine at runtime). This should only be set if building a model + in order to export and ahead-of-time compile it into a standalone + library. + """ + + self._functional = hub.Module(spec=self._model_path) + + grid_coords = tf.placeholder( + tf.float32, shape=[batch_dim, 3], name='grid_coords') + grid_weights = tf.placeholder( + tf.float32, shape=[batch_dim], name='grid_weights') + + # Density information. + rho_a = tf.placeholder(tf.float32, shape=[6, batch_dim], name='rho_a') + rho_b = tf.placeholder(tf.float32, shape=[6, batch_dim], name='rho_b') + + # Split into corresponding terms. + rho_only_a, grad_a_x, grad_a_y, grad_a_z, _, tau_a = tf.unstack( + rho_a, axis=0) + rho_only_b, grad_b_x, grad_b_y, grad_b_z, _, tau_b = tf.unstack( + rho_b, axis=0) + + # Evaluate |\del \rho|^2 for each spin density and for the total density. + norm_grad_a = (grad_a_x**2 + grad_a_y**2 + grad_a_z**2) + norm_grad_b = (grad_b_x**2 + grad_b_y**2 + grad_b_z**2) + grad_x = grad_a_x + grad_b_x + grad_y = grad_a_y + grad_b_y + grad_z = grad_a_z + grad_b_z + norm_grad = (grad_x**2 + grad_y**2 + grad_z**2) + + # The local Hartree-Fock energy densities at each grid point for the alpha- + # and beta-spin densities for each value of omega. + # Note an omega of 0 indicates no screening of the Coulomb potential. + hfxa = tf.placeholder( + tf.float32, shape=[batch_dim, len(self._omega_values)], name='hfxa') + hfxb = tf.placeholder( + tf.float32, shape=[batch_dim, len(self._omega_values)], name='hfxb') + + # Make all features 2D arrays on input for ease of handling inside the + # functional. + features = { + 'grid_coords': grid_coords, + 'grid_weights': tf.expand_dims(grid_weights, 1), + 'rho_a': tf.expand_dims(rho_only_a, 1), + 'rho_b': tf.expand_dims(rho_only_b, 1), + 'tau_a': tf.expand_dims(tau_a, 1), + 'tau_b': tf.expand_dims(tau_b, 1), + 'norm_grad_rho_a': tf.expand_dims(norm_grad_a, 1), + 'norm_grad_rho_b': tf.expand_dims(norm_grad_b, 1), + 'norm_grad_rho': tf.expand_dims(norm_grad, 1), + 'hfxa': hfxa, + 'hfxb': hfxb, + } + tensor_dict = {f'tensor_dict${k}': v for k, v in features.items()} + + predictions = self._functional(tensor_dict, as_dict=True) + local_xc = predictions['grid_contribution'] + weighted_local_xc = local_xc * grid_weights + unweighted_xc = tf.reduce_sum(local_xc, axis=0) + xc = tf.reduce_sum(weighted_local_xc, axis=0) + + # The potential is the local exchange correlation divided by the + # total density. Add a small constant to deal with zero density. + self._vxc = local_xc / (rho_only_a + rho_only_b + 1E-12) + + # The derivatives of the exchange-correlation (XC) energy with respect to + # input features. PySCF weights the (standard) derivatives by the grid + # weights, so we need to compute this with respect to the unweighted sum + # over grid points. + self._vrho = tf.gradients( + unweighted_xc, [features['rho_a'], features['rho_b']], + name='GRAD_RHO', + unconnected_gradients=tf.UnconnectedGradients.ZERO) + self._vsigma = tf.gradients( + unweighted_xc, [ + features['norm_grad_rho_a'], features['norm_grad_rho_b'], + features['norm_grad_rho'] + ], + name='GRAD_SIGMA', + unconnected_gradients=tf.UnconnectedGradients.ZERO) + self._vtau = tf.gradients( + unweighted_xc, [features['tau_a'], features['tau_b']], + name='GRAD_TAU', + unconnected_gradients=tf.UnconnectedGradients.ZERO) + # Standard meta-GGAs do not have a dependency on local HF, so we need to + # compute the contribution to the Fock matrix ourselves. Just use the + # weighted XC energy to avoid having to weight this later. + self._vhf = tf.gradients( + xc, [features['hfxa'], features['hfxb']], + name='GRAD_HFX', + unconnected_gradients=tf.UnconnectedGradients.ZERO) + + self._placeholders = FunctionalInputs( + rho_a=rho_a, + rho_b=rho_b, + hfx_a=hfxa, + hfx_b=hfxb, + grid_coords=grid_coords, + grid_weights=grid_weights) + + outputs = { + 'vxc': self._vxc, + 'vrho': tf.stack(self._vrho), + 'vsigma': tf.stack(self._vsigma), + 'vtau': tf.stack(self._vtau), + 'vhf': tf.stack(self._vhf), + } + # Create the signature for TF-Hub, including both the energy and functional + # derivatives. + # This is a no-op if _build_graph is called outside of + # hub.create_module_spec. + hub.add_signature( + inputs=attr.asdict(self._placeholders), outputs=outputs) + + def export_functional_and_derivatives( + self, + export_path: str, + batch_dim: Optional[int] = None, + ): + """Exports the TensorFlow graph containing the functional and derivatives. + + The hub modules supplied contain the TensorFlow operations for the + evaluation of the exchange-correlation energy. Evaluation of the functional + derivatives, required for a self-consistent calculation, are added in + _build_graph. The module created by export_functional_and_derivatives + contains the evaluation of the functional and the functional derivatives. + This is much simpler to use from languages other than Python, e.g. using the + C or C++ TensorFlow API, or using tfcompile to create a standalone C++ + library. + + Args: + export_path: path to write the Hub model to. The exported model can be + loaded using either TF-Hub or SavedModel APIs. + batch_dim: the batch dimension of the grid to use in the model. Default: + None (determine at runtime). This should only be set if the exported + model is to be ahead-of-time compiled into a standalone library. + """ + with tf.Graph().as_default(): + spec = hub.create_module_spec( + self._build_graph, tags_and_args=[(set(), {'batch_dim': batch_dim})]) + functional_and_derivatives = hub.Module(spec=spec) + with tf.Session() as session: + session.run(tf.global_variables_initializer()) + functional_and_derivatives.export(export_path, session) + + # DM21* functionals include the hybrid term directly, so set the + # range-separated and hybrid parameters expected by PySCF to 0 so PySCF + # doesn't also add these contributions in separately. + def rsh_coeff(self, *args): + """Returns the range separated parameters, omega, alpha, beta.""" + return [0.0, 0.0, 0.0] + + def hybrid_coeff(self, *args, **kwargs): + """Returns the fraction of Hartree-Fock exchange to include.""" + return 0.0 + + def _xc_type(self, *args, **kwargs): + return 'MGGA' + + def nr_rks(self, + mol: gto.Mole, + grids: dft.Grids, + xc_code: str, + dms: Union[np.ndarray, Sequence[np.ndarray]], + relativity: int = 0, + hermi: int = 0, + max_memory: float = 20000, + verbose=None) -> Tuple[float, float, np.ndarray]: + """Calculates RKS XC functional and potential matrix on a given grid. + + Args: + mol: PySCF molecule. + grids: grid on which to evaluate the functional. + xc_code: XC code. Unused. NeuralNumInt hard codes the XC functional + based upon the functional argument given to the constructor. + dms: the density matrix or sequence of density matrices. Multiple density + matrices are not currently supported. Shape (nao, nao), where nao is the + number of atomic orbitals. + relativity: Unused. (pyscf.numint.NumInt.nr_rks does not currently use + this argument.) + hermi: 0 if the density matrix is Hermitian, 1 if the density matrix is + non-Hermitian. + max_memory: the maximum cache to use, in MB. + verbose: verbosity level. Unused. (PySCF currently does not handle the + verbosity level passed in here.) + + Returns: + nelec, excsum, vmat, where + nelec is the number of electrons obtained by numerical integration of + the density matrix. + excsum is the functional's XC energy. + vmat is the functional's XC potential matrix, shape (nao, nao). + + Raises: + NotImplementedError: if multiple density matrices are supplied. + """ + # Wrap nr_rks so we can store internal variables required to evaluate the + # contribution to the XC potential from local Hartree-Fock features. + # See pyscf.dft.numint.nr_rks for more details. + ndms = _get_number_of_density_matrices(dms) + if ndms > 1: + raise NotImplementedError( + 'NeuralNumInt does not support multiple density matrices. ' + 'Only ground state DFT calculations are currently implemented.') + nao = mol.nao_nr() + self._vmat_hf = np.zeros((nao, nao)) + self._system_state = _SystemState(mol=mol, dms=dms) + nelec, excsum, vmat = super().nr_rks( + mol=mol, + grids=grids, + xc_code=xc_code, + dms=dms, + relativity=relativity, + hermi=hermi, + max_memory=max_memory, + verbose=verbose) + vmat += self._vmat_hf + self._vmat_hf.T + + # Clear internal state to prevent accidental re-use. + self._system_state = None + self._grid_state = None + return nelec, excsum, vmat + + def nr_uks(self, + mol: gto.Mole, + grids: dft.Grids, + xc_code: str, + dms: Union[Sequence[np.ndarray], Sequence[Sequence[np.ndarray]]], + relativity: int = 0, + hermi: int = 0, + max_memory: float = 20000, + verbose=None) -> Tuple[np.ndarray, float, np.ndarray]: + """Calculates UKS XC functional and potential matrix on a given grid. + + Args: + mol: PySCF molecule. + grids: grid on which to evaluate the functional. + xc_code: XC code. Unused. NeuralNumInt hard codes the XC functional + based upon the functional argument given to the constructor. + dms: the density matrix or sequence of density matrices for each spin + channel. Multiple density matrices for each spin channel are not + currently supported. Each density matrix is shape (nao, nao), where nao + is the number of atomic orbitals. + relativity: Unused. (pyscf.dft.numint.NumInt.nr_rks does not currently use + this argument.) + hermi: 0 if the density matrix is Hermitian, 1 if the density matrix is + non-Hermitian. + max_memory: the maximum cache to use, in MB. + verbose: verbosity level. Unused. (PySCF currently does not handle the + verbosity level passed in here.) + + Returns: + nelec, excsum, vmat, where + nelec is the number of alpha, beta electrons obtained by numerical + integration of the density matrix as an array of size 2. + excsum is the functional's XC energy. + vmat is the functional's XC potential matrix, shape (2, nao, nao), where + vmat[0] and vmat[1] are the potential matrices for the alpha and beta + spin channels respectively. + + Raises: + NotImplementedError: if multiple density matrices for each spin channel + are supplied. + """ + # Wrap nr_uks so we can store internal variables required to evaluate the + # contribution to the XC potential from local Hartree-Fock features. + # See pyscf.dft.numint.nr_uks for more details. + if isinstance(dms, np.ndarray) and dms.ndim == 2: # RHF DM + ndms = _get_number_of_density_matrices(dms) + else: + ndms = _get_number_of_density_matrices(dms[0]) + if ndms > 1: + raise NotImplementedError( + 'NeuralNumInt does not support multiple density matrices. ' + 'Only ground state DFT calculations are currently implemented.') + + nao = mol.nao_nr() + self._vmat_hf = np.zeros((2, nao, nao)) + self._system_state = _SystemState(mol=mol, dms=dms) + nelec, excsum, vmat = super().nr_uks( + mol=mol, + grids=grids, + xc_code=xc_code, + dms=dms, + relativity=relativity, + hermi=hermi, + max_memory=max_memory, + verbose=verbose) + vmat[0] += self._vmat_hf[0] + self._vmat_hf[0].T + vmat[1] += self._vmat_hf[1] + self._vmat_hf[1].T + + # Clear internal state to prevent accidental re-use. + self._system_state = None + self._grid_state = None + self._vmat_hf = None + return nelec, excsum, vmat + + def block_loop( + self, + mol: gto.Mole, + grids: dft.Grids, + nao: Optional[int] = None, + deriv: int = 0, + max_memory: float = 2000, + non0tab: Optional[np.ndarray] = None, + blksize: Optional[int] = None, + buf: Optional[np.ndarray] = None + ) -> Generator[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], None, + None]: + """Loops over the grid by blocks. See pyscf.dft.numint.NumInt.block_loop. + + Args: + mol: PySCF molecule. + grids: grid on which to evaluate the functional. + nao: number of basis functions. If None, obtained from mol. + deriv: unused. The first functional derivatives are always computed. + max_memory: the maximum cache to use for the information on the grid, in + MB. Determines the size of each block if blksize is None. + non0tab: mask determining if a shell in the basis set is zero at a grid + point. Shape (N, nbas), where N is the number of grid points and nbas + the number of shells in the basis set. Obtained from grids if not + supplied. + blksize: size of each block. Calculated from max_memory if None. + buf: buffer to use for storing ao. If None, a new array for ao is created + for each block. + + Yields: + ao, mask, weight, coords: information on a block of the grid containing N' + points, where + ao: atomic orbitals evaluated on the grid. Shape (N', nao), where nao is + the number of atomic orbitals. + mask: mask indicating whether a shell in the basis set is zero at a grid + point. Shape (N', nbas). + weight: weight associated with each grid point. Shape (N'). + coords: coordinates of the grid. Shape (N', 3). + """ + # Wrap block_loop so we can store internal variables required to evaluate + # the contribution to the XC potential from local Hartree-Fock features. + for ao, mask, weight, coords in super().block_loop( + mol=mol, + grids=grids, + nao=nao, + deriv=deriv, + max_memory=max_memory, + non0tab=non0tab, + blksize=blksize, + buf=buf): + # Cache the curent block so we can access it in eval_xc. + self._grid_state = _GridState( + ao=ao, mask=mask, weight=weight, coords=coords) + yield ao, mask, weight, coords + + def construct_functional_inputs( + self, + mol: gto.Mole, + dms: Union[np.ndarray, Sequence[np.ndarray]], + spin: int, + coords: np.ndarray, + weights: np.ndarray, + rho: Union[np.ndarray, Tuple[np.ndarray, np.ndarray]], + ao: Optional[np.ndarray] = None, + ) -> Tuple[FunctionalInputs, Tuple[np.ndarray, np.ndarray]]: + """Constructs the input features required for the functional. + + Args: + mol: PySCF molecule. + dms: density matrix of shape (nao, nao) (restricted calculations) or of + shape (2, nao, nao) (unrestricted calculations) or tuple of density + matrices for each spin channel, each of shape (nao, nao) (unrestricted + calculations). + spin: 0 for a spin-unpolarized (restricted Kohn-Sham) calculation, and + spin-polarized (unrestricted) otherwise. + coords: coordinates of the grid. Shape (N, 3), where N is the number of + grid points. + weights: weight associated with each grid point. Shape (N). + rho: density and density derivatives at each grid point. Single array + containing the total density for restricted calculations, tuple of + arrays for each spin channel for unrestricted calculations. Each array + has shape (6, N). See pyscf.dft.numint.eval_rho and comments in + FunctionalInputs for more details. + ao: The atomic orbitals evaluated on the grid, shape (N, nao). Computed if + not supplied. + + Returns: + inputs, fxx, where + inputs: FunctionalInputs object containing the inputs (as np.ndarrays) + for the functional. + fxx: intermediates, shape (N, nao) for the alpha- and beta-spin + channels, required for computing the first derivative of the local + Hartree-Fock density with respect to the density matrices. See + compute_hfx_density for more details. + """ + if spin == 0: + # RKS + rhoa = rho / 2 + rhob = rho / 2 + else: + # UKS + rhoa, rhob = rho + + # Local HF features. + exxa, exxb = [], [] + fxxa, fxxb = [], [] + for omega in sorted(self._omega_values): + hfx_results = compute_hfx_density.get_hf_density( + mol, + dms, + coords=coords, + omega=omega, + deriv=1, + ao=ao) + exxa.append(hfx_results.exx[0]) + exxb.append(hfx_results.exx[1]) + fxxa.append(hfx_results.fxx[0]) + fxxb.append(hfx_results.fxx[1]) + exxa = np.stack(exxa, axis=-1) + fxxa = np.stack(fxxa, axis=-1) + if spin == 0: + exx = (exxa, exxa) + fxx = (fxxa, fxxa) + else: + exxb = np.stack(exxb, axis=-1) + fxxb = np.stack(fxxb, axis=-1) + exx = (exxa, exxb) + fxx = (fxxa, fxxb) + + return FunctionalInputs( + rho_a=rhoa, + rho_b=rhob, + hfx_a=exx[0], + hfx_b=exx[1], + grid_coords=coords, + grid_weights=weights), fxx + + def eval_xc( + self, + xc_code: str, + rho: Union[np.ndarray, Tuple[np.ndarray, np.ndarray]], + spin: int = 0, + relativity: int = 0, + deriv: int = 1, + verbose=None + ) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], + None, None]: + """Evaluates the XC energy and functional derivatives. + + See pyscf.dft.libxc.eval_xc for more details on the interface. + + Note: this also sets self._vmat_extra, which contains the contribution the + the potential matrix from the local Hartree-Fock terms in the functional. + + Args: + xc_code: unused. + rho: density and density derivatives at each grid point. Single array + containing the total density for restricted calculations, tuple of + arrays for each spin channel for unrestricted calculations. Each array + has shape (6, N), where N is the number of grid points. See + pyscf.dft.numint.eval_rho and comments in FunctionalInputs for more + details. + spin: 0 for a spin-unpolarized (restricted Kohn-Sham) calculation, and + spin-polarized (unrestricted) otherwise. + relativity: unused. + deriv: unused. The first functional derivatives are always computed. + verbose: unused. + + Returns: + exc, vxc, fxc, kxc, where: + exc is the exchange-correlation potential matrix evaluated at each grid + point, shape (N). + vxc is (vrho, vgamma, vlapl, vtau), the first-order functional + derivatives evaluated at each grid point, each shape (N). + fxc is set to None. (The second-order functional derivatives are not + computed.) + kxc is set to None. (The third-order functional derivatives are not + computed.) + """ + del xc_code, verbose, relativity, deriv # unused + + # Retrieve cached state. + ao = self._grid_state.ao + if ao.ndim == 3: + # Just need the AO values, not the gradients. + ao = ao[0] + if self._grid_state.weight is None: + weights = np.array([1.]) + else: + weights = self._grid_state.weight + mask = self._grid_state.mask + + inputs, (fxxa, fxxb) = self.construct_functional_inputs( + mol=self._system_state.mol, + dms=self._system_state.dms, + spin=spin, + rho=rho, + weights=weights, + coords=self._grid_state.coords, + ao=ao) + + with self._graph.as_default(): + feed_dict = dict( + zip( + attr.asdict(self._placeholders).values(), + attr.asdict(inputs).values(), + )) + tensor_list = [ + self._vxc, + self._vrho, + self._vsigma, + self._vtau, + self._vhf, + ] + exc, vrho, vsigma, vtau, vhf = ( + self._session.run(tensor_list, feed_dict=feed_dict)) + + mol = self._system_state.mol + shls_slice = (0, mol.nbas) + ao_loc_nr = mol.ao_loc_nr() + # Note: tf.gradients returns a list of gradients. + # vrho, vsigma, vtau are derivatives of objects that had + # tf.expand_dims(..., 1) applied. The [:, 0] indexing undoes this by + # selecting the 0-th (and only) element from the second dimension. + if spin == 0: + vxc_0 = (vrho[0][:, 0] + vrho[1][:, 0]) / 2. + # pyscf expects derivatives with respect to: + # grad_rho . grad_rho. + # The functional uses the first and last as inputs, but then has + # grad_(rho_a + rho_b) . grad_(rho_a + rho_b) + # as input. The following computes the correct total derivatives. + vxc_1 = (vsigma[0][:, 0] / 4. + vsigma[1][:, 0] / 4. + vsigma[2][:, 0]) + vxc_3 = (vtau[0][:, 0] + vtau[1][:, 0]) / 2. + vxc_2 = np.zeros_like(vxc_3) + vhfs = (vhf[0] + vhf[1]) / 2. + # Local Hartree-Fock terms + for i in range(len(self._omega_values)): + # Factor of 1/2 is to account for adding vmat_hf + vmat_hf.T to vmat, + # which we do to match existing PySCF style. Unlike other terms, vmat_hf + # is already symmetric though. + aow = np.einsum('pi,p->pi', fxxa[:, :, i], -0.5 * vhfs[:, i]) + self._vmat_hf += _dot_ao_ao(mol, ao, aow, mask, shls_slice, + ao_loc_nr) + else: + vxc_0 = np.stack([vrho[0][:, 0], vrho[1][:, 0]], axis=1) + # pyscf expects derivatives with respect to: + # grad_rho_a . grad_rho_a + # grad_rho_a . grad_rho_b + # grad_rho_b . grad_rho_b + # The functional uses the first and last as inputs, but then has + # grad_(rho_a + rho_b) . grad_(rho_a + rho_b) + # as input. The following computes the correct total derivatives. + vxc_1 = np.stack([ + vsigma[0][:, 0] + vsigma[2][:, 0], 2. * vsigma[2][:, 0], + vsigma[1][:, 0] + vsigma[2][:, 0] + ], + axis=1) + vxc_3 = np.stack([vtau[0][:, 0], vtau[1][:, 0]], axis=1) + vxc_2 = np.zeros_like(vxc_3) + vhfs = np.stack([vhf[0], vhf[1]], axis=2) + for i in range(len(self._omega_values)): + # Factors of 1/2 are due to the same reason as in the spin=0 case. + aow = np.einsum('pi,p->pi', fxxa[:, :, i], -0.5 * vhfs[:, i, 0]) + self._vmat_hf[0] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, + ao_loc_nr) + aow = np.einsum('pi,p->pi', fxxb[:, :, i], -0.5 * vhfs[:, i, 1]) + self._vmat_hf[1] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, + ao_loc_nr) + + fxc = None # Second derivative not implemented + kxc = None # Second derivative not implemented + return exc, (vxc_0, vxc_1, vxc_2, vxc_3), fxc, kxc diff --git a/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint_test.py b/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint_test.py new file mode 100644 index 0000000..0f62e09 --- /dev/null +++ b/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint_test.py @@ -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() diff --git a/density_functional_approximation_dm21/external/tf_bazel.patch b/density_functional_approximation_dm21/external/tf_bazel.patch new file mode 100644 index 0000000..380a589 --- /dev/null +++ b/density_functional_approximation_dm21/external/tf_bazel.patch @@ -0,0 +1,78 @@ +From 76edcc2da8a754cb2b283d7de9a7b37f51f7f380 Mon Sep 17 00:00:00 2001 +From: James Spencer +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 + diff --git a/density_functional_approximation_dm21/requirements.txt b/density_functional_approximation_dm21/requirements.txt new file mode 100644 index 0000000..ff94956 --- /dev/null +++ b/density_functional_approximation_dm21/requirements.txt @@ -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 diff --git a/density_functional_approximation_dm21/requirements_aot_compilation.txt b/density_functional_approximation_dm21/requirements_aot_compilation.txt new file mode 100644 index 0000000..eb9409b --- /dev/null +++ b/density_functional_approximation_dm21/requirements_aot_compilation.txt @@ -0,0 +1,3 @@ +Keras-Preprocessing==1.1.2 +numpy==1.19.5 +tensorflow-estimator==2.7.0 diff --git a/density_functional_approximation_dm21/run.sh b/density_functional_approximation_dm21/run.sh new file mode 100755 index 0000000..319f2a3 --- /dev/null +++ b/density_functional_approximation_dm21/run.sh @@ -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/ diff --git a/density_functional_approximation_dm21/setup.py b/density_functional_approximation_dm21/setup.py new file mode 100644 index 0000000..028ad84 --- /dev/null +++ b/density_functional_approximation_dm21/setup.py @@ -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']}, +)